from typing import Dict, Any, Optional
from src.intent import detect_intent
from src.templates import TEMPLATES

DEFAULT_GEN_ARGS = {
    "max_tokens": 300,
    "temperature": 0.7,
    "top_p": 0.95
}

MSG_SEPARATOR = "\n"

class LocalChatbot:
    def __init__(self, llm, memory, tokenizer=None, default_template: Optional[str] = "general"):
        self.llm = llm
        self.memory = memory
        self.tokenizer = tokenizer
        self.default_template = default_template

    def _build_system_prompt(self, intent: str) -> str:
        return TEMPLATES.get(intent, TEMPLATES.get(self.default_template, TEMPLATES["general"]))

    def _build_prompt(self, user_message: str, intent: str, max_pairs: int = 12) -> str:
        try:
            self.memory.trim_to_recent_pairs(max_pairs)
        except Exception:
            pass

        system_prompt = self._build_system_prompt(intent)
        history_text = self.memory.get_formatted(separator=MSG_SEPARATOR)

        parts = [
            f"System: {system_prompt}",
            history_text,
            f"User: {user_message}",
            "Assistant:"
        ]
        return MSG_SEPARATOR.join([p for p in parts if p])

    def ask(self, user_message: str, gen_args: Optional[Dict[str, Any]] = None) -> str:
        if not user_message.strip():
            return "Please enter a message."

        intent = detect_intent(user_message)
        prompt = self._build_prompt(user_message, intent)

        gen = DEFAULT_GEN_ARGS.copy()
        if gen_args:
            gen.update(gen_args)

        try:
            if self.tokenizer:
                # Transformers-style generation
                inputs = self.tokenizer(prompt, return_tensors="pt")
                outputs = self.llm.generate(**inputs, max_new_tokens=gen.get("max_tokens", 300))
                bot_reply = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
            else:
                # Fallback: callable LLM
                bot_reply = self.llm(prompt, **gen)
                if isinstance(bot_reply, dict) and "choices" in bot_reply:
                    bot_reply = bot_reply["choices"][0].get("text", "").strip()
        except Exception:
            bot_reply = "Sorry — I couldn't generate a response. Please try again."

        if not bot_reply:
            bot_reply = "Sorry — I couldn't generate a response. Please try again."

        try:
            self.memory.add(user_message, bot_reply)
        except Exception:
            try:
                self.memory.add_message("user", user_message)
                self.memory.add_message("assistant", bot_reply)
            except Exception:
                pass

        return bot_reply