import sys
from unittest.mock import MagicMock

# MANDATORY: Mock lzma BEFORE any other imports
try:
    import lzma
except ImportError:
    mock_lzma = MagicMock()
    mock_lzma.FORMAT_XZ, mock_lzma.FORMAT_ALONE, mock_lzma.FORMAT_RAW = 1, 2, 3
    mock_lzma.CHECK_NONE, mock_lzma.CHECK_CRC32, mock_lzma.CHECK_CRC64, mock_lzma.CHECK_SHA256 = 0, 1, 4, 10
    sys.modules["_lzma"] = MagicMock()
    sys.modules["lzma"] = mock_lzma

import time
import argparse
import multiprocessing
import os
import re
import json
import requests
from collections import deque
from datetime import datetime, timezone


def build_paragraph_messages(previous_refined_context, previous_source_text, new_source_text):
    system_prompt = """You are a careful live transcript editor working from noisy translated source text.

Goal:
Produce the most faithful readable English update.

Decision rules:
1. NEW SOURCE TEXT is the primary evidence and should dominate the output.
2. Use PREVIOUS SOURCE TEXT and PREVIOUS REFINED CONTEXT only when they clearly help resolve or continue the new material.
3. The final answer must be English only.
4. Never copy non-English words into the output unless they are proper names.
5. If a non-English fragment appears as an isolated clause, trailing fragment, or side comment without a clear English anchor, drop it.
6. Only translate a non-English fragment when it is clearly central to the same idea and its meaning is reasonably obvious from context.
7. If the new material starts a fresh thought, output only the new material.
8. Remove exact or near-exact repetition unless the repetition is clearly intentional rhetoric.
9. Do not add explanations, disclaimers, or meta commentary.
10. Prefer conservative wording over guessed meaning.
11. Return only the revised transcript text."""

    prompt = f"""Edit this transcript update.

[PREVIOUS REFINED CONTEXT]
{previous_refined_context}

[PREVIOUS SOURCE TEXT]
{previous_source_text}

[NEW SOURCE TEXT]
{new_source_text}
"""
    return [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": prompt},
    ], prompt


def build_line_messages(prev1, prev2, corrected):
    system_prompt = """You are a real-time English caption corrector for live speech.

Task:
Clean only the current caption line.

Hard rules:
1. Preserve meaning exactly. Never replace current content with prior context.
2. Remove disfluencies and false starts.
3. Fix punctuation, casing, and obvious STT typos.
4. If uncertain, return the original line unchanged.
5. Output only one corrected English line."""

    prompt = (
        "Context (reference only):\n"
        f"Previous line 1: {prev1}\n"
        f"Previous line 2: {prev2}\n"
        "Current line to correct:\n"
        f"{corrected}\n"
        "Corrected:"
    )
    return [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": prompt},
    ], prompt


def append_llm_request_log(log_path, entry):
    if not log_path:
        return
    log_dir = os.path.dirname(log_path)
    if log_dir:
        os.makedirs(log_dir, exist_ok=True)
    with open(log_path, "a", encoding="utf-8") as handle:
        json.dump(entry, handle, ensure_ascii=False)
        handle.write("\n")


def load_llm_request_log_entry(log_path, entry_index):
    if not log_path or not os.path.exists(log_path):
        raise FileNotFoundError(f"LLM request log not found: {log_path}")

    with open(log_path, "r", encoding="utf-8") as handle:
        entries = [json.loads(line) for line in handle if line.strip()]

    if not entries:
        raise ValueError(f"LLM request log is empty: {log_path}")

    if entry_index is None or entry_index == -1:
        return entries[-1]

    if entry_index < 0:
        entry_index = len(entries) + entry_index

    if entry_index < 0 or entry_index >= len(entries):
        raise IndexError(f"LLM request log index {entry_index} is out of range for {len(entries)} entries")

    return entries[entry_index]


def run_llm_prompt_test(args):
    from dotenv import load_dotenv
    load_dotenv()

    api_key = os.environ.get("OPENAI_API_KEY")
    log_path = getattr(args, "llm_request_log_path", "logs/llm_requests.jsonl")

    def call_openai(model, messages):
        response = requests.post(
            "https://api.openai.com/v1/chat/completions",
            json={"model": model, "messages": messages},
            headers={"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"},
            timeout=getattr(args, "post_correct_llm_timeout", 8.0),
        )
        if response.status_code >= 400:
            try:
                payload = response.json()
                detail = json.dumps(payload, ensure_ascii=True)
            except Exception:
                detail = response.text.strip()
            raise RuntimeError(f"OpenAI API error {response.status_code}: {detail}")
        return response.json()["choices"][0]["message"]["content"].strip()

    def call_ollama(model, prompt, temperature, system=None, extra_options=None):
        options = {"temperature": temperature}
        if extra_options:
            options.update(extra_options)
        response = requests.post(
            getattr(args, "post_correct_ollama_url", "http://127.0.0.1:11434/api/generate"),
            json={
                "model": model,
                "prompt": prompt,
                "system": system or "",
                "stream": False,
                "options": options,
                "keep_alive": getattr(args, "post_correct_keep_alive", "30m"),
            },
            timeout=getattr(args, "post_correct_llm_timeout", 8.0),
        )
        response.raise_for_status()
        return response.json().get("response", "").strip()

    def send_test_request(mode, messages, prompt, temperature, metadata, model=None, extra_options=None):
        request_model = model or args.post_correct_model
        is_openai = request_model.startswith("gpt-")
        entry = {
            "timestamp": datetime.now(timezone.utc).isoformat(),
            "mode": mode,
            "provider": "openai" if is_openai else "ollama",
            "model": request_model,
            "temperature": temperature,
            "options": extra_options or {},
            "messages": messages,
            "metadata": metadata,
        }
        started_at = time.time()
        if is_openai:
            if not api_key:
                raise RuntimeError("OPENAI_API_KEY is not set")
            try:
                response_text = call_openai(request_model, messages)
                entry["response"] = {"status": "ok", "content": response_text}
                return response_text
            except Exception as exc:
                entry["response"] = {"status": "error", "error": str(exc)}
                raise
            finally:
                entry["duration_ms"] = round((time.time() - started_at) * 1000, 2)
                append_llm_request_log(log_path, entry)
        system = ""
        for message in messages:
            if message.get("role") == "system":
                system = message.get("content", "")
                break
        try:
            response_text = call_ollama(request_model, prompt, temperature, system=system, extra_options=extra_options)
            entry["response"] = {"status": "ok", "content": response_text}
            return response_text
        except Exception as exc:
            entry["response"] = {"status": "error", "error": str(exc)}
            raise
        finally:
            entry["duration_ms"] = round((time.time() - started_at) * 1000, 2)
            append_llm_request_log(log_path, entry)

    if getattr(args, "llm_test_from_log", None) is not None:
        logged_entry = load_llm_request_log_entry(log_path, args.llm_test_from_log)
        logged_messages = logged_entry.get("messages")
        if not logged_messages:
            raise ValueError("Selected log entry does not contain messages")
        logged_prompt = next((m.get("content", "") for m in logged_messages if m.get("role") == "user"), "")
        logged_model = logged_entry.get("model", args.post_correct_model)
        replay_model = logged_model if getattr(args, "llm_test_use_logged_model", False) else args.post_correct_model
        response = send_test_request(
            logged_entry.get("mode", "replay"),
            logged_messages,
            logged_prompt,
            logged_entry.get("temperature", 0.1),
            {
                "replay_source_log_path": log_path,
                "replay_source_index": args.llm_test_from_log,
                "replay_source_timestamp": logged_entry.get("timestamp"),
                "replay_source_model": logged_model,
            },
            model=replay_model,
        )
        print("[LLM TEST] Replayed request from log:")
        print(f"source_index={args.llm_test_from_log} source_model={logged_model} replay_model={replay_model}")
        print(response)

    if getattr(args, "llm_test_line", None):
        messages, prompt = build_line_messages(
            getattr(args, "llm_test_prev1", ""),
            getattr(args, "llm_test_prev2", ""),
            args.llm_test_line,
        )
        response = send_test_request(
            "line",
            messages,
            prompt,
            0.1,
            {
                "prev1": getattr(args, "llm_test_prev1", ""),
                "prev2": getattr(args, "llm_test_prev2", ""),
                "input": args.llm_test_line,
            },
        )
        print("[LLM TEST] Line response:")
        print(response)

    if getattr(args, "llm_test_segments", None):
        messages, prompt = build_paragraph_messages(
            getattr(args, "llm_test_context", ""),
            args.llm_test_segments,
        )
        response = send_test_request(
            "paragraph",
            messages,
            prompt,
            0.1,
            {
                "context": getattr(args, "llm_test_context", ""),
                "segments": args.llm_test_segments,
            },
        )
        print("[LLM TEST] Paragraph response:")
        print(response)

def run_llm_processor(in_queue, out_queue, args):
    from dotenv import load_dotenv
    load_dotenv()
    
    print("[LLM] Starting Rolling Paragraph Engine...")
    
    api_key = os.environ.get("OPENAI_API_KEY")
    is_openai = args.post_correct_model.startswith("gpt-")
    llm_state = {"line_warned": False, "paragraph_warned": False}
    log_path = getattr(args, "llm_request_log_path", "logs/llm_requests.jsonl")

    def check_ollama_health():
        try:
            response = requests.get("http://127.0.0.1:11434/api/tags", timeout=2.5)
            response.raise_for_status()
            data = response.json()
            models = data.get("models", [])
            available = {model.get("name") for model in models if model.get("name")}
            if args.post_correct_model not in available and f"{args.post_correct_model}:latest" not in available:
                print(f"[LLM] Configured Ollama model '{args.post_correct_model}' is not installed.")
            return True
        except Exception as exc:
            print(f"[LLM] Ollama health check failed ({exc}). Local LLM features may fall back.")
            return False

    def warmup_ollama_model():
        try:
            response = requests.post(
                getattr(args, "post_correct_ollama_url", "http://127.0.0.1:11434/api/generate"),
                json={
                    "model": args.post_correct_model,
                    "prompt": "Return exactly: ok",
                    "system": "You are a concise assistant.",
                    "stream": False,
                    "options": {"temperature": 0.0},
                    "keep_alive": getattr(args, "post_correct_keep_alive", "30m"),
                },
                timeout=getattr(args, "post_correct_warmup_timeout", 20.0),
            )
            response.raise_for_status()
            body = response.json()
            if body.get("response", "").strip():
                total_s = body.get("total_duration", 0) / 1_000_000_000
                load_s = body.get("load_duration", 0) / 1_000_000_000
                print(f"[LLM] Ollama warmup OK for {args.post_correct_model} (load={load_s:.2f}s total={total_s:.2f}s).")
                return True
        except Exception as exc:
            print(f"[LLM] Ollama warmup failed for {args.post_correct_model} ({exc}).")
        return False

    if not is_openai:
        check_ollama_health()
        if getattr(args, "post_correct_llm", False) or getattr(args, "llm_paragraph", False):
            warmup_ollama_model()

    def normalize_english_caption(text):
        normalized = " ".join((text or "").split()).strip()
        if normalized and normalized[0].isalpha():
            normalized = normalized[0].upper() + normalized[1:]
        return normalized

    def apply_rule_based_post_correction(text):
        corrected = normalize_english_caption(text)
        corrected = re.sub(r"\s+", " ", corrected).strip()
        corrected = re.sub(r"\s+([,.;!?])", r"\1", corrected)
        corrected = re.sub(r"([,.;!?]){2,}", r"\1", corrected)
        corrected = re.sub(r"\b(uh+|um+|erm+|ah+|hmm+)\b", "", corrected, flags=re.IGNORECASE)
        corrected = re.sub(r"\s+", " ", corrected).strip()
        corrected = re.sub(r"\b(\w+)\s+\1\b", r"\1", corrected, flags=re.IGNORECASE)
        return normalize_english_caption(corrected)

    def word_overlap_ratio(source_text, candidate_text):
        src_tokens = set(re.findall(r"[a-z0-9']+", source_text.lower()))
        cand_tokens = set(re.findall(r"[a-z0-9']+", candidate_text.lower()))
        if not src_tokens:
            return 1.0
        return len(src_tokens & cand_tokens) / len(src_tokens)

    def call_openai(messages):
        response = requests.post(
            "https://api.openai.com/v1/chat/completions",
            json={"model": args.post_correct_model, "messages": messages},
            headers={"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"},
            timeout=getattr(args, "post_correct_llm_timeout", 8.0),
        )
        if response.status_code >= 400:
            try:
                payload = response.json()
                detail = json.dumps(payload, ensure_ascii=True)
            except Exception:
                detail = response.text.strip()
            raise RuntimeError(f"OpenAI API error {response.status_code}: {detail}")
        return response.json()["choices"][0]["message"]["content"].strip()

    def call_ollama(prompt, temperature, system=None, extra_options=None):
        options = {"temperature": temperature}
        if extra_options:
            options.update(extra_options)
        response = requests.post(
            getattr(args, "post_correct_ollama_url", "http://127.0.0.1:11434/api/generate"),
            json={
                "model": args.post_correct_model,
                "prompt": prompt,
                "system": system or "",
                "stream": False,
                "options": options,
                "keep_alive": getattr(args, "post_correct_keep_alive", "30m"),
            },
            timeout=getattr(args, "post_correct_llm_timeout", 8.0),
        )
        response.raise_for_status()
        return response.json().get("response", "").strip()

    def call_llm(messages, prompt, temperature, warn_key, extra_options=None):
        entry = {
            "timestamp": datetime.now(timezone.utc).isoformat(),
            "mode": "paragraph" if warn_key == "paragraph_warned" else "line",
            "provider": "openai" if is_openai else "ollama",
            "model": args.post_correct_model,
            "temperature": temperature,
            "options": extra_options or {},
            "messages": messages,
        }
        started_at = time.time()
        try:
            if is_openai:
                if not api_key:
                    raise RuntimeError("OPENAI_API_KEY is not set")
                response_text = call_openai(messages)
                entry["response"] = {"status": "ok", "content": response_text}
                return response_text
            system = ""
            if messages:
                for message in messages:
                    if message.get("role") == "system":
                        system = message.get("content", "")
                        break
            response_text = call_ollama(prompt, temperature, system=system, extra_options=extra_options)
            entry["response"] = {"status": "ok", "content": response_text}
            return response_text
        except Exception as exc:
            entry["response"] = {"status": "error", "error": str(exc)}
            if not llm_state[warn_key]:
                label = "Paragraph LLM" if warn_key == "paragraph_warned" else "Line LLM"
                print(f"[LLM] {label} unavailable ({exc}). Falling back to deterministic text.")
                llm_state[warn_key] = True
            return ""
        finally:
            entry["duration_ms"] = round((time.time() - started_at) * 1000, 2)
            append_llm_request_log(log_path, entry)

    # State: This is the ONLY text we send to the LLM as context
    active_context = ""
    previous_source_window = ""
    recent_lines = deque(maxlen=2)
    
    def call_llm_rolling_refine(new_segments_str, context, previous_source_text):
        messages, prompt = build_paragraph_messages(context, previous_source_text, new_segments_str)
        paragraph_options = {
            "top_p": getattr(args, "llm_paragraph_top_p", 0.8),
            "repeat_penalty": getattr(args, "llm_paragraph_repeat_penalty", 1.0),
        }
        return call_llm(
            messages,
            prompt,
            getattr(args, "llm_paragraph_temperature", 0.2),
            "paragraph_warned",
            extra_options=paragraph_options,
        )

    def post_correct_line(text):
        if not getattr(args, "post_correct", False):
            return normalize_english_caption(text)

        corrected = apply_rule_based_post_correction(text)
        if not getattr(args, "post_correct_llm", False):
            return corrected

        prev2 = recent_lines[-2] if len(recent_lines) >= 2 else ""
        prev1 = recent_lines[-1] if len(recent_lines) >= 1 else ""
        messages, prompt = build_line_messages(prev1, prev2, corrected)
        candidate = normalize_english_caption(call_llm(messages, prompt, 0.1, "line_warned"))
        if candidate and word_overlap_ratio(corrected, candidate) >= getattr(args, "post_correct_min_overlap", 0.45):
            return candidate
        return corrected

    pending_buffer = []
    last_llm_call_time = time.time()
    
    while True:
        try:
            item = in_queue.get()
            if item is None: break
            if "draft" in item:
                out_queue.put(item)
                continue

            raw_text = item.get("original", "").strip()
            bridge_text = item.get("en_bridge", raw_text).strip()
            if not raw_text: continue

            corrected_text = post_correct_line(bridge_text)

            # Hallucination check
            words = corrected_text.lower().split()
            if len(words) > 10 and len(set(words)) < 3: continue

            pending_buffer.append(corrected_text)
            recent_lines.append(corrected_text)

            should_refine = False
            if len(pending_buffer) >= 5: should_refine = True
            elif len(pending_buffer) >= 2 and corrected_text.endswith((".", "?", "!")): should_refine = True
            elif len(pending_buffer) >= 1 and (time.time() - last_llm_call_time) > 15: should_refine = True

            structured_paragraph = None
            paragraph_from_llm = False
            if args.llm_paragraph and should_refine:
                new_batch = " ".join(pending_buffer)
                result = call_llm_rolling_refine(new_batch, active_context, previous_source_window)
                
                if result:
                    # Logic to handle paragraph breaks
                    if "\n\n" in result:
                        # Split by double newline
                        parts = result.split("\n\n")
                        # The last part is the new "Active Context"
                        active_context = parts[-1].strip()
                        # The whole result is sent to the user (contains the break)
                        structured_paragraph = result
                    else:
                        # No break, just update the context
                        active_context = result
                        structured_paragraph = result
                    paragraph_from_llm = True
                else:
                    structured_paragraph = new_batch
                    active_context = new_batch
                previous_source_window = new_batch

                pending_buffer = []
                last_llm_call_time = time.time()

            out_queue.put({
                "raw": raw_text,
                "corrected": corrected_text,
                "paragraph": structured_paragraph,
                "en_bridge": bridge_text,
                "used_llm_line": bool(getattr(args, "post_correct_llm", False)),
                "used_llm_paragraph": paragraph_from_llm,
                "detected_lang": item.get("detected_lang"),
                "speaker": item.get("speaker"),
                "ts": item.get("ts", time.time())
            })
            
        except Exception as e:
            print(f"[LLM] Error: {e}")

if __name__ == "__main__":
    run_llm_processor(multiprocessing.Queue(), multiprocessing.Queue(), argparse.Namespace())