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 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-")
    target_lang = "English"
    if args.lang and args.lang != 'en': target_lang = args.lang

    # State: This is the ONLY text we send to the LLM as context
    active_context = ""
    
    def call_llm_rolling_refine(new_segments_str, context):
        if not is_openai: return ""
        
        url = "https://api.openai.com/v1/chat/completions"
        headers = {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"}
        
        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 {target_lang}.
5. OUTPUT: Provide ONLY the refined, consolidated text. Do not explain anything.
"""
        payload = {
            "model": args.post_correct_model,
            "messages": [{"role": "system", "content": "You are a live transcript editor."},
                         {"role": "user", "content": prompt}],
            "temperature": 0.1
        }
        try:
            resp = requests.post(url, json=payload, headers=headers, timeout=20)
            if resp.status_code == 200:
                return resp.json()["choices"][0]["message"]["content"].strip()
            return ""
        except: return ""

    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

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

            pending_buffer.append(bridge_text)

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

            structured_paragraph = None
            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
                    
                    pending_buffer = []
                    last_llm_call_time = time.time()

            out_queue.put({
                "raw": raw_text,
                "corrected": bridge_text,
                "paragraph": structured_paragraph,
                "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())