
import json, torch, numpy as np
from sentence_transformers import SentenceTransformer, CrossEncoder
import faiss
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig

class Chronos:
    def __init__(self, model_dir="."):
        with open(f"{model_dir}/rag_config.json") as f:
            config = json.load(f)
        self.embedder = SentenceTransformer(config["embedder_model"])
        self.index = faiss.read_index(f"{model_dir}/jjk_index.faiss")
        with open(f"{model_dir}/chunks.txt", "r") as f:
            raw = f.read().split("<|CHUNK_END|>")
        self.chunks = [c.strip() for c in raw if c.strip()]
        self.reranker = CrossEncoder(f"{model_dir}/cross_encoder_model")
        bnb = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_use_double_quant=True,
            bnb_4bit_quant_type="nf4",
            bnb_4bit_compute_dtype=torch.float16
        )
        self.tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True)
        self.model = AutoModelForCausalLM.from_pretrained(
            model_dir,
            quantization_config=bnb,
            device_map="auto",
            trust_remote_code=True
        )

    def ask(self, question, max_tokens_casual=200, max_tokens_hist=400):
        q = question.strip().lower().rstrip("?!.")
        greetings = [
            "hi", "hello", "hey", "yo", "sup", "good morning",
            "good evening", "how are you", "hi bro", "hi there",
            "what's up", "howdy"
        ]
        if any(q == g or q.startswith(g) for g in greetings):
            return "Hello! I'm Chronos, your 20th‑century historian. Ask me anything about the wars, inventions, and events that shaped the modern world. Or just chat – I'm pretty friendly! 😄"

        identity_qs = [
            "who are you", "what is your name", "what are you",
            "are you chronos", "your name", "who is chronos",
            "what is chronos", "tell me about yourself",
            "what do you do", "who created you", "who made you"
        ]
        if any(idq in q for idq in identity_qs):
            return "I'm Chronos 🕰️, a friendly historian AI. I can talk about whatever you like, but my real passion is 20th‑century history – WW1, WW2, the Cold War, and all the inventions that came out of that era. So, what's on your mind?"

        safety_net = {
            "leader of germany during world war i": "Kaiser Wilhelm II was the German Emperor during World War I.",
            "dictator of germany during world war ii": "Adolf Hitler was the dictator of Germany during World War II.",
            "leader of germany during ww2": "Adolf Hitler was the dictator of Germany during World War II.",
            "leader of the soviet union during world war ii": "Joseph Stalin led the Soviet Union during World War II.",
            "prime minister of the uk during world war ii": "Winston Churchill was the UK Prime Minister during World War II.",
            "president of the usa during world war ii": "Franklin D. Roosevelt was the US President for most of WWII.",
            "when did world war i start": "World War I started on July 28, 1914.",
            "when did world war ii start": "World War II started on September 1, 1939.",
            "what was the holocaust": "The Holocaust was the systematic murder of six million Jews by the Nazi regime.",
        }
        for key, ans in safety_net.items():
            if key in q:
                return ans

        history_kw = [
            "ww1", "ww2", "world war", "cold war", "great war", "nazi", "hitler", "stalin", "mussolini", "tojo",
            "holocaust", "atomic bomb", "nuclear weapon", "manhattan project", "d-day", "pearl harbor",
            "battle of", "treaty of", "versailles", "yalta", "potsdam", "munich agreement",
            "kaiser", "tsar", "emperor", "dictator", "president", "prime minister", "chancellor",
            "soviet union", "ussr", "nazi germany", "third reich", "allies", "axis",
            "invent", "radar", "penicillin", "jet engine", "computer", "transistor", "satellite", "rocket",
            "space race", "apollo", "sputnik",
            "depression", "new deal", "marshall plan", "nato", "warsaw pact", "berlin wall", "iron curtain",
            "decolonization", "civil rights", "women's suffrage", "baby boom",
            "what caused", "how did", "why did", "explain", "describe", "timeline", "significance", "result of",
            "who was", "who were", "when did", "where did", "what was the", "what is the",
            "churchill", "roosevelt", "lenin", "mao", "kennedy", "eisenhower", "montgomery", "rommel",
            "blitzkrieg", "trench warfare", "enigma", "code talkers", "kamikaze", "fat man", "little boy"
        ]
        is_history = any(kw in q for kw in history_kw)

        if not is_history:
            try:
                prompt = (
                    "<|im_start|>system\n"
                    "You are Chronos, a friendly and knowledgeable historian. "
                    "You love chatting with people about anything, not just history. "
                    "Be warm, concise, and lively.\n"
                    "<|im_end|>\n"
                    f"<|im_start|>user\n{question}<|im_end|>\n"
                    "<|im_start|>assistant\n"
                )
                inputs = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)
                outputs = self.model.generate(
                    **inputs,
                    max_new_tokens=max_tokens_casual,
                    temperature=0.9,
                    do_sample=True,
                    pad_token_id=self.tokenizer.eos_token_id
                )
                return self.tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True).strip()
            except:
                return "I'm not sure what to say, but I'm happy to talk about 20th‑century history if you're interested! 😊"

        try:
            q_emb = self.embedder.encode([question]).astype("float32")
            _, indices = self.index.search(q_emb, 30)
            candidates = [self.chunks[i] for i in indices[0]]
            pairs = [(question, c) for c in candidates]
            scores = self.reranker.predict(pairs)
            if max(scores) < -4.5:
                return "I'm sorry, I don't have enough historical information to answer that question accurately."
            best = sorted(zip(scores, candidates), reverse=True)[:4]
            context = "\n\n".join([c for _, c in best])
            messages = [
                {"role": "system", "content": "You are Chronos, a historian. Use the provided context to answer accurately. If the information is not present, say 'I don't know'."},
                {"role": "user", "content": f"Context:\n{context}\n\nQuestion: {question}"}
            ]
            prompt = self.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
            inputs = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)
            outputs = self.model.generate(
                **inputs,
                max_new_tokens=max_tokens_hist,
                temperature=0.7,
                do_sample=True,
                pad_token_id=self.tokenizer.eos_token_id
            )
            return self.tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True).strip()
        except:
            return "I'm sorry, something went wrong with the historical search. Please try again."
