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

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}

    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",
                    "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):
        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,
                "stream": False,
                "options": {"temperature": temperature},
                "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):
        try:
            if is_openai:
                if not api_key:
                    raise RuntimeError("OPENAI_API_KEY is not set")
                return call_openai(messages)
            return call_ollama(prompt, temperature)
        except Exception as 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 ""

    # State: This is the ONLY text we send to the LLM as context
    active_context = ""
    recent_lines = deque(maxlen=2)
    
    def call_llm_rolling_refine(new_segments_str, context):
        prompt = f"""Task: Refine the following live transcription stream into clean, professional paragraphs.

[PREVIOUS WORKING CONTEXT]
"{context}"

[NEW RAW ASR SEGMENTS]
"{new_segments_str}"

[INSTRUCTIONS]
1. INTEGRATE: Polished and merge the new segments into the flow of the 'PREVIOUS WORKING CONTEXT'.
2. CONSOLIDATE: Remove redundant repetitions and translator echoes.
3. ORGANIZE: Use a double newline (\\n\\n) to start a new paragraph when a topic changes or the current one is complete.
4. TARGET LANGUAGE: Output ONLY in English.
5. OUTPUT: Provide ONLY the refined, consolidated text. Do not explain anything.
"""
        messages = [
            {"role": "system", "content": "You are a live transcript editor."},
            {"role": "user", "content": prompt},
        ]
        return call_llm(messages, prompt, 0.1, "paragraph_warned")

    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 ""
        prompt = (
            "You are a real-time English caption corrector for live speech.\n"
            "Task:\n"
            "Clean ONLY the current caption line.\n"
            "Hard rules:\n"
            "1. Preserve meaning exactly. Never replace current content with prior context.\n"
            "2. Remove disfluencies and false starts.\n"
            "3. Fix punctuation, casing, and obvious STT typos.\n"
            "4. If uncertain, return the original line unchanged.\n"
            "5. Output ONLY one corrected English line.\n"
            "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:"
        )
        messages = [
            {"role": "system", "content": "You are a real-time English caption corrector."},
            {"role": "user", "content": prompt},
        ]
        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)
                
                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

                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())