from llama_cpp import Llama
from characters import CHARACTERS
import config
import os

try:
    import spaces
    def gpu(fn):
        return spaces.GPU(fn)
except ImportError:
    def gpu(fn):
        return fn


class TheaterEngine:
    def __init__(self):
        print("[TheaterEngine] Loading GGUF model via llama.cpp (CUDA enabled)...")
        # Check if the model exists
        model_path = "models/Qwen3.5-4B-q4_k_m.gguf"
        if not os.path.exists(model_path):
            print(f"[TheaterEngine] WARNING: Model not found at {model_path}!")
            
        self.model = Llama(
            model_path=model_path,
            n_gpu_layers=-1, # Offload all layers to RTX 4060
            n_ctx=4096,      # Context window
            verbose=False,
            chat_format="chatml" # Qwen format
        )
        self.char_map = {c["name"]: c for c in CHARACTERS}
        print("[TheaterEngine] Ready.")

    @gpu
    def generate_line(self, char_name: str, history: list[dict]) -> str:
        """
        Generate the next line for the given character.
        """
        char = self.char_map[char_name]
        messages = self._format_messages(char, history)

        response = self.model.create_chat_completion(
            messages=messages,
            max_tokens=config.MAX_NEW_TOKENS,
            temperature=config.TEMPERATURE,
            top_p=config.TOP_P,
        )

        text = response['choices'][0]['message']['content'].strip()

        # Akıl yürütme bloklarını (<think>...</think>) ve içindeki her şeyi tamamen sil (Kapatılmamış etiketleri de yakalar)
        import re
        text = re.sub(r'<think>.*?(?:</think>|$)\s*', '', text, flags=re.DOTALL)
        
        # Eğer sadece düşünmüş ve cevap vermemişse
        if not text.strip():
            text = "..."

        # Strip any accidental "CharacterName: " prefix
        if text.startswith(f"{char_name}:"):
            text = text[len(char_name) + 1:].strip()

        return text or "..."

    @gpu
    def generate_scene_description(self, history: list[dict]) -> str:
        """
        Generate a short atmospheric scene description based on recent dialog.
        """
        recent = history[-8:] if len(history) > 8 else history
        dialog_snippet = "\n".join(
            f"{e['speaker']}: {e['content']}" for e in recent
            if e["speaker"] not in ["[STAGE DIRECTION]", "[STRANGER]"]
        )

        messages = [
            {
                "role": "system",
                "content": (
                    "You are a stage direction writer. Based on a theater scene excerpt, "
                    "write exactly two sentences describing the physical scene as it appears RIGHT NOW. "
                    "Focus on: lighting quality, shadows, objects, textures, emotional atmosphere made tangible. "
                    "Be poetic and concrete. Never mention character names. Never include dialog."
                ),
            },
            {"role": "user", "content": dialog_snippet},
        ]

        response = self.model.create_chat_completion(
            messages=messages,
            max_tokens=80,
            temperature=0.7,
        )
        
        return response['choices'][0]['message']['content'].strip()

    def _format_messages(self, char: dict, history: list[dict]) -> list[dict]:
        """
        Format conversation history for a specific character.
        That character's lines → "assistant"
        Everyone else's lines → "user" (with speaker prefix)
        """
        messages = [{"role": "system", "content": char["system"]}]
        messages.append({"role": "user", "content": config.OPENING_NARRATION})

        for entry in history:
            speaker = entry["speaker"]
            content = entry["content"]

            if speaker == char["name"]:
                role = "assistant"
                text = content
            else:
                role = "user"
                if speaker.startswith("["):
                    text = content
                else:
                    text = f"{speaker}: {content}"

            if messages and messages[-1]["role"] == role:
                messages[-1]["content"] += f"\n{text}"
            else:
                messages.append({"role": role, "content": text})

        if len(messages) > config.MAX_HISTORY_MESSAGES:
            messages = messages[:2] + messages[-(config.MAX_HISTORY_MESSAGES - 2):]

        return messages