"""
ViuAI-500M Agentic Chat (Advanced UI & Streaming)
=================================================
This script implements a pause-and-resume agentic loop with word-by-word streaming.
- Manual Autoregressive Loop for true streaming.
- Dynamic colorization for [THINK] (grey) and <search> (yellow).
- Timeout protection (10s) on DuckDuckGo search.
- String interception for robust </search> handling (Repetition Bug Fix).
"""
import argparse
import os
import re
import sys
import threading
from concurrent.futures import ThreadPoolExecutor, TimeoutError
from datetime import datetime

import torch
import torch.nn.functional as F
from huggingface_hub import hf_hub_download
from transformers import PreTrainedTokenizerFast

from config import ViuAIConfig
from model import ViuAI

import warnings
warnings.filterwarnings("ignore")
def no_warning(*args, **kwargs): pass
warnings.showwarning = no_warning
import logging
logging.getLogger("transformers").setLevel(logging.ERROR)

# ============================================================================
# 1. Internet Search (DuckDuckGo with Timeout Protection)
# ============================================================================
def _ddgs_search_sync(query: str) -> str:
    results = []
    try:
        from duckduckgo_search import DDGS
        with DDGS() as ddgs:
            results = list(ddgs.text(query, max_results=3))
    except ImportError:
        return "Search failed: duckduckgo-search package not installed."
    except Exception as e:
        return f"Search failed: {e}"

    if not results:
        today = datetime.now().strftime("%d %B %Y")
        title1 = f"Latest on {query[:20]}"
        title2 = f"More info for {query[:20]}"
        title3 = f"Updates: {query[:20]}"
        return (
            f"[Today: {today}] Results:\n"
            f"1. {title1}: No relevant information found for this topic.\n"
            f"2. {title2}: No relevant information found for this topic.\n"
            f"3. {title3}: No relevant information found for this topic."
        )

    MAX_CHARS = 1200
    today = datetime.now().strftime("%d %B %Y")
    context = f"[Today: {today}] Results:\n"

    for i, r in enumerate(results[:3], 1):
        snippet = f"{i}. {r.get('title', '')}: {r.get('body', '')}\n"
        if len(context) + len(snippet) > MAX_CHARS:
            context += snippet[:MAX_CHARS - len(context)] + "...\n"
            break
        context += snippet
    return context.strip()

def do_search(query: str, timeout: int = 10) -> str:
    """Search with strict timeout to prevent hangs."""
    with ThreadPoolExecutor(max_workers=1) as executor:
        future = executor.submit(_ddgs_search_sync, query)
        try:
            return future.result(timeout=timeout)
        except TimeoutError:
            return "Search timed out after 10 seconds. Please rely on internal knowledge."
        except Exception as e:
            return f"Search error: {e}"

# ============================================================================
# 2. Model Loading
# ============================================================================
def load_tokenizer(repo_id, token):
    tokenizer_path = hf_hub_download(repo_id=repo_id, filename="tokenizer/tokenizer.json", token=token)
    tokenizer = PreTrainedTokenizerFast(
        tokenizer_file=tokenizer_path,
        pad_token="<pad>", bos_token="<bos>", eos_token="<eos>", unk_token="<unk>",
        additional_special_tokens=["<|user|>", "<|assistant|>", "<|endofturn|>", "[THINK]", "[/THINK]"]
    )
    tokenizer.add_tokens(["<search>", "</search>", "<search_result>", "</search_result>"], special_tokens=True)
    return tokenizer

def strip_training_prefixes(state_dict):
    cleaned = {}
    for key, value in state_dict.items():
        for prefix in ("module._orig_mod.", "_orig_mod.", "module."):
            if key.startswith(prefix):
                key = key[len(prefix):]
                break
        cleaned[key] = value
    return cleaned

# ============================================================================
# 3. Main Agentic Loop with Streaming
# ============================================================================
def main():
    parser = argparse.ArgumentParser(description="Agentic Chat with ViuAI-500M")
    parser.add_argument("--repo-id", default=os.environ.get("VIUAI_MODEL_REPO", "ViuAI/ViuAI-500M"))
    args = parser.parse_args()

    token = os.environ.get("HF_TOKEN")
    device = "cuda" if torch.cuda.is_available() else "cpu"

    print(f"🚀 Loading Agentic ViuAI-500M on {device.upper()}...")
    tokenizer = load_tokenizer(args.repo_id, token)
    
    try:
        ckpt_path = hf_hub_download(repo_id=args.repo_id, filename="sft_checkpoints/sft_v6/sft_ckpt_latest.pt", token=token)
    except:
        ckpt_path = "checkpoints/ckpt_latest_updated.pt"
        if not os.path.exists(ckpt_path):
            print("⚠️ Could not find SFT checkpoint on HF, downloading base updated local...")
            ckpt_path = hf_hub_download(repo_id=args.repo_id, filename="checkpoints/ckpt_latest.pt", token=token)

    checkpoint = torch.load(ckpt_path, map_location="cpu", weights_only=False)
    state_dict = strip_training_prefixes(checkpoint["model"])
    
    ckpt_vocab_size = state_dict['tok_emb.weight'].shape[0]
    config = ViuAIConfig(vocab_size=ckpt_vocab_size, use_checkpoint=False)
    model = ViuAI(config).to(device)
    model.load_state_dict(state_dict, strict=True)
    model.head.weight = model.tok_emb.weight
    model.eval()

    print(f"✅ Ready! (Vocab size: {ckpt_vocab_size})")
    print(f"💡 Tips: The model will autonomously decide when to search the internet.")
    print(f"         Type 'quit' or 'exit' to stop.\n")

    eot_id = tokenizer.convert_tokens_to_ids("<|endofturn|>")

    # Colors for Terminal UI
    RESET = "\033[0m"
    GREY = "\033[90m"
    YELLOW = "\033[33m"
    CYAN = "\033[36m"
    RED = "\033[31m"

    while True:
        try:
            message = input(f"👤 You: ").strip()
        except (EOFError, KeyboardInterrupt):
            print()
            break
        if message.lower() in {"quit", "exit"}:
            break
        if not message:
            continue

        # ── RAG Pre-Search ──────────────────────────────────────────────────────
        # Detect if query needs current/live info. Only trigger for explicit news/live queries.
        SEARCH_TRIGGERS = [
            "latest news", "current news", "today news", "news today", "latest phone",
            "price of", "stock price", "weather today", "score today", "match result",
            "who is the current", "who is the present", "aaj ki khabar", "aaj ka mausam",
            "india news", "news"
        ]
        
        # DO NOT trigger search for identity/basic questions
        IDENTITY_KEYWORDS = [
            "who are you", "what is your name", "what's your name", "your name", 
            "tell me about yourself", "who created you", "who made you"
        ]
        
        msg_lower = message.lower()
        is_identity = any(kw in msg_lower for kw in IDENTITY_KEYWORDS)
        needs_search = any(kw in msg_lower for kw in SEARCH_TRIGGERS) and not is_identity

        if needs_search:
            print(f"\n{CYAN}⚡ [Auto-Search Triggered for live query...]{RESET}")
            pre_result = do_search(message, timeout=12)
            print(f"{CYAN}   Result length: {len(pre_result)} chars{RESET}\n")
            if pre_result and "No results" not in pre_result and "timed out" not in pre_result:
                prompt = (
                    f"<|user|>\n"
                    f"{message}<|endofturn|>\n"
                    f"<|assistant|>\n"
                    f"<search>{message}</search>\n"
                    f"<search_result>\n{pre_result}\n</search_result>\n"
                )
                full_reply_text = ""
            else:
                prompt = f"<|user|>\n{message}<|endofturn|>\n<|assistant|>\n"
                full_reply_text = ""
        else:
            prompt = f"<|user|>\n{message}<|endofturn|>\n<|assistant|>\n"
            full_reply_text = ""
        # ────────────────────────────────────────────────────────────────────────

        input_ids = tokenizer(prompt, add_special_tokens=False, return_tensors="pt").input_ids.to(device)
        generated_ids = []
        last_decoded_len = 0
        
        print(f"\n🤖 ViuAI:")
        current_color = RESET
        search_count = 0
        # Track open/close counts for proper colorization
        think_open = 0
        think_close = 0
        search_open = 0
        search_close_count = 0
        
        # ── Streaming Generation Loop ───────────────────────────────────────────
        while True:
            # Safely truncate input_ids to prevent CUDA OOM
            if input_ids.size(1) > config.context_length:
                input_ids = input_ids[:, -config.context_length:]
                
            max_new = min(1024, max(1, config.context_length - input_ids.size(1)))
            
            triggered_search = False

            # --- Custom Autoregressive Streaming Loop ---
            for _ in range(max_new):
                # cond_idx is now just input_ids since we truncate above
                cond_idx = input_ids
                
                with torch.no_grad():
                    logits = model(cond_idx)[0]
                logits = logits[:, -1, :] # last token only
                
                temperature = 0.6
                top_k = 40
                top_p = 0.9
                repetition_penalty = 1.15
                
                # Apply Repetition Penalty (Vectorized & Only on generated tokens!)
                if repetition_penalty != 1.0 and len(generated_ids) > 0:
                    unique_tokens = torch.unique(torch.tensor(generated_ids, device=device))
                    score = logits[0, unique_tokens]
                    logits[0, unique_tokens] = torch.where(
                        score < 0, score * repetition_penalty, score / repetition_penalty
                    )
                
                # Temperature Scaling
                logits = logits / max(temperature, 1e-8)
                
                # Top K
                kth_vals, _ = torch.topk(logits, top_k)
                logits[logits < kth_vals[:, [-1]]] = float('-inf')
                
                # Top P
                sorted_logits, sorted_indices = torch.sort(logits, descending=True)
                cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
                sorted_indices_to_remove = cumulative_probs > top_p
                sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
                sorted_indices_to_remove[..., 0] = False
                indices_to_remove = torch.zeros_like(logits, dtype=torch.bool).scatter_(1, sorted_indices, sorted_indices_to_remove)
                logits[indices_to_remove] = float('-inf')
                
                # Sample
                probs = F.softmax(logits, dim=-1)
                idx_next = torch.multinomial(probs, num_samples=1)
                
                # Append to sequence
                input_ids = torch.cat([input_ids, idx_next], dim=1)
                
                # Safely truncate input_ids inside the loop too
                if input_ids.size(1) > config.context_length:
                    input_ids = input_ids[:, -config.context_length:]
                    
                new_token_id = idx_next.item()
                generated_ids.append(new_token_id)
                
                # Safe UTF-8 decoding by checking difference on generated_ids only
                full_decoded = tokenizer.decode(generated_ids, skip_special_tokens=False)
                new_text = full_decoded[last_decoded_len:]
                last_decoded_len = len(full_decoded)
                
                full_reply_text += new_text

                # --- Live Colorization Logic (count-based, works across turns) ---
                think_open  = full_reply_text.count("[THINK]")
                think_close = full_reply_text.count("[/THINK]")
                search_open       = full_reply_text.count("<search>")
                search_close_count = full_reply_text.count("</search>")

                in_think  = think_open > think_close
                in_search = search_open > search_close_count

                if in_think:
                    if current_color != GREY:
                        sys.stdout.write(GREY)
                        current_color = GREY
                elif in_search:
                    if current_color != YELLOW:
                        sys.stdout.write(YELLOW)
                        current_color = YELLOW
                else:
                    if current_color != RESET:
                        sys.stdout.write(RESET)
                        current_color = RESET

                # Stream out (Hide the invisible EOT token from UI)
                if new_token_id != eot_id:
                    sys.stdout.write(new_text)
                    sys.stdout.flush()

                # --- Stop Conditions ---
                if new_token_id == eot_id:
                    break
                
                # String Intercept for Repetition Bug (Fixes <search> loop)
                current_search_count = full_reply_text.count("</search>")
                if current_search_count > search_count:
                    search_count = current_search_count
                    triggered_search = True
                    break

            # Handle post-stream Agentic Actions (Search)
            if triggered_search:
                # Guarantee color resets
                sys.stdout.write(RESET)
                print(f"\n\n{CYAN}🔍 [Agent is Searching the Internet...]{RESET}")
                
                # Get the LAST search block in the text
                search_matches = re.findall(r'<search>(.*?)</search>', full_reply_text, re.DOTALL)
                query = search_matches[-1].strip() if search_matches else ""
                
                if query:
                    print(f"{CYAN}   Query: {query}{RESET}")
                    result_text = do_search(query, timeout=10)
                else:
                    result_text = "No query provided."
                
                print(f"{CYAN}   Result length: {len(result_text)} chars{RESET}\n")
                
                # Inject knowledge (Silent to UI)
                sys.stdout.write(f"{GREY}[Context Injected. Resuming...]{RESET}\n")
                sys.stdout.flush()
                
                injection = f"\n<search_result>\n{result_text}\n</search_result>\n"
                full_reply_text += injection
                
                # Convert injection to tokens and append
                inj_ids = tokenizer.encode(injection, add_special_tokens=False, return_tensors="pt").to(device)
                input_ids = torch.cat([input_ids, inj_ids], dim=1)
                
                # Continue outer loop to resume generation based on injected knowledge
                continue
            else:
                # Finished generating naturally
                print("\n")
                break

if __name__ == "__main__":
    main()