import os
import gradio as gr
from llama_cpp import Llama
from huggingface_hub import list_bucket_tree, download_bucket_files

MODEL_DIR = "./model_files"

from huggingface_hub import snapshot_download

def download_phi3_vision_model():
    repo_id = "microsoft/Phi-3.5-vision-instruct-onnx"
    target_subdir = "cpu_and_mobile/cpu-int4-rtn-block-32-acc-level-4"
    local_dir = os.path.join(MODEL_DIR, target_subdir)
    
    print(f"[startup] Checking model files in repo {repo_id}...", flush=True)
    try:
        snapshot_download(
            repo_id=repo_id,
            allow_patterns=[f"{target_subdir}/*"],
            local_dir=MODEL_DIR,
        )
        print("[startup] Download/cache check complete.", flush=True)
    except Exception as e:
        print(f"[startup] Error checking or downloading model: {e}", flush=True)
        raise e
        
    return local_dir

_phi_model = None
_phi_processor = None
_phi_tokenizer = None

def _load_phi_vision():
    global _phi_model, _phi_processor, _phi_tokenizer
    if _phi_model is None:
        local_dir = download_phi3_vision_model()
        print(f"[startup] Loading Phi-3-Vision ONNX model from {local_dir}...", flush=True)
        import onnxruntime_genai as og
        _phi_model = og.Model(local_dir)
        _phi_processor = _phi_model.create_multimodal_processor()
        _phi_tokenizer = og.Tokenizer(_phi_model)
        print("[startup] Phi-3-Vision ONNX model loaded.", flush=True)
    return _phi_model, _phi_processor, _phi_tokenizer

def _clean_history_for_qwen(history):
    clean = []
    for turn in history:
        role = turn["role"]
        content = turn["content"]
        if isinstance(content, dict):
            text = content.get("text", "")
            clean.append({"role": role, "content": text})
        else:
            clean.append({"role": role, "content": str(content)})
    return clean

# ---------------------------------------------------------------------------
# Local reasoning models (Qwen3.5, quantized, via llama-cpp-python)
# ---------------------------------------------------------------------------

# "fast" is the default for everyday chat; "deep" trades latency for
# noticeably stronger reasoning (e.g. UPSC GS-style analysis) on the same
# 2 vCPU / no-GPU hardware — pick it per-message via the Response Mode radio.
LLM_VARIANTS = {
    "fast": {"repo_id": "unsloth/Qwen3.5-4B-GGUF", "filename": "Qwen3.5-4B-Q4_K_M.gguf"},
    "deep": {"repo_id": "unsloth/Qwen3.5-9B-GGUF", "filename": "Qwen3.5-9B-Q4_K_M.gguf"},
}

DEFAULT_SYSTEM_PROMPT = "You are a helpful, knowledgeable assistant."

_llms: dict[str, Llama] = {}


def _load_llm(variant: str) -> Llama:
    """Download (cached by huggingface_hub) and load a variant, once each."""
    if variant not in _llms:
        cfg = LLM_VARIANTS[variant]
        print(f"[startup] Loading {variant} model: {cfg['repo_id']}/{cfg['filename']} ...", flush=True)
        _llms[variant] = Llama.from_pretrained(
            repo_id=cfg["repo_id"],
            filename=cfg["filename"],
            n_ctx=8192,
            n_threads=os.cpu_count() or 2,
            verbose=False,
        )
        print(f"[startup] {variant} model loaded.", flush=True)
    return _llms[variant]


def _build_chatml_prompt(messages: list[dict]) -> str:
    """Manually render ChatML with an empty, pre-closed <think> block on the
    assistant turn, so the model never starts its own reasoning trace.

    Workaround for a currently-open llama.cpp bug where `enable_thinking:
    false` is silently ignored for Qwen3.5, causing every reply to pay for a
    full (often long) hidden reasoning trace even for trivial messages.
    See: https://github.com/ggml-org/llama.cpp/issues/20182
    """
    parts = [f"<|im_start|>{m['role']}\n{m['content']}<|im_end|>\n" for m in messages]
    parts.append("<|im_start|>assistant\n<think>\n\n</think>\n\n")
    return "".join(parts)


def respond(message, history, system_prompt, max_tokens, temperature, disable_thinking, mode):
    """Streaming chat callback for gr.ChatInterface (type='messages')."""
    if mode == "vision":
        # Load Phi-3-Vision
        model, processor, tokenizer = _load_phi_vision()
        import onnxruntime_genai as og
        
        user_text = message.get("text", "") if isinstance(message, dict) else str(message)
        user_files = message.get("files", []) if isinstance(message, dict) else []
        
        # Format prompt and collect images
        prompt_parts = []
        if system_prompt.strip():
            prompt_parts.append(f"<|system|>\n{system_prompt.strip()}<|end|>\n")
            
        image_counter = 1
        image_paths = []
        
        for turn in history:
            role = turn["role"]
            content = turn["content"]
            
            if role == "user":
                if isinstance(content, dict):
                    text = content.get("text", "")
                    files = content.get("files", [])
                else:
                    text = str(content)
                    files = []
                
                img_tags = ""
                for f in files:
                    if f.lower().endswith((".png", ".jpg", ".jpeg", ".webp", ".bmp")):
                        img_tags += f"<|image_{image_counter}|>\n"
                        image_paths.append(f)
                        image_counter += 1
                prompt_parts.append(f"<|user|>\n{img_tags}{text}<|end|>\n")
            elif role == "assistant":
                prompt_parts.append(f"<|assistant|>\n{content}<|end|>\n")
                
        # Current message
        img_tags = ""
        for f in user_files:
            if f.lower().endswith((".png", ".jpg", ".jpeg", ".webp", ".bmp")):
                img_tags += f"<|image_{image_counter}|>\n"
                image_paths.append(f)
                image_counter += 1
                
        prompt_parts.append(f"<|user|>\n{img_tags}{user_text}<|end|>\n<|assistant|>\n")
        prompt = "".join(prompt_parts)
        
        try:
            # Prepare inputs
            if image_paths:
                print(f"[inference] Processing prompt with {len(image_paths)} images...", flush=True)
                images = og.Images.open(*image_paths)
                inputs = processor(prompt, images=images)
            else:
                print("[inference] Processing text-only prompt for Phi-3...", flush=True)
                try:
                    inputs = processor(prompt)
                except Exception:
                    inputs = tokenizer.encode(prompt)
                    
            # Set up generator parameters
            params = og.GeneratorParams(model)
            params.set_search_options(
                max_length=8192,
                temperature=temperature
            )
            
            # Stream output
            generator = og.Generator(model, params)
            
            # Pass inputs
            if type(inputs).__name__ == "NamedTensors":
                generator.set_inputs(inputs)
            else:
                import numpy as np
                generator.append_tokens(np.array(inputs, dtype=np.int32))
                
            partial = ""
            max_gen_tokens = int(max_tokens)
            tokens_generated = 0
            while not generator.is_done() and tokens_generated < max_gen_tokens:
                generator.generate_next_token()
                new_token = generator.get_next_tokens()[0]
                decoded = tokenizer.decode([new_token])
                partial += decoded
                yield partial
                tokens_generated += 1
        except Exception as e:
            print(f"[inference] Error during Phi-3 generation: {e}", flush=True)
            yield f"⚠️ **Error during model inference**: {str(e)}\n\n*This was caught gracefully to prevent the Space from crashing. Please try again or rephrase.*"
    else:
        # Qwen modes
        llm = _load_llm(mode)
        messages = [{"role": "system", "content": system_prompt.strip() or DEFAULT_SYSTEM_PROMPT}]
        messages.extend(_clean_history_for_qwen(history))
        
        user_text = message.get("text", "") if isinstance(message, dict) else str(message)
        messages.append({"role": "user", "content": user_text})
        
        partial = ""
        if disable_thinking:
            # Raw completion on a hand-built prompt (empty-think prefill trick),
            # since create_chat_completion's enable_thinking kwarg is a no-op here.
            stream = llm.create_completion(
                prompt=_build_chatml_prompt(messages),
                max_tokens=int(max_tokens),
                temperature=temperature,
                stop=["<|im_end|>", "<|im_start|>"],
                stream=True,
            )
            for chunk in stream:
                delta = chunk["choices"][0]["text"]
                if delta:
                    partial += delta
                    yield partial
        else:
            for chunk in llm.create_chat_completion(
                messages=messages,
                max_tokens=int(max_tokens),
                temperature=temperature,
                stream=True,
            ):
                delta = chunk["choices"][0]["delta"].get("content", "")
                if delta:
                    partial += delta
                    yield partial


# ---------------------------------------------------------------------------
# UI
# ---------------------------------------------------------------------------

with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 💬 Local Reasoning Chat")
    gr.Markdown(
        "General-purpose chat backed by quantized **Qwen3.5** models and **Phi-3-Vision**. "
        "Served locally via `llama-cpp-python` and `onnxruntime-genai` (CPU-only).\n\n"
        "**Image Uploading (Text Extraction)**: To ask the model to extract text or analyze an image, "
        "make sure you select the **Vision** mode below, then click the **+ (attachment icon)** "
        "inside the chat input box to upload your image."
    )

    with gr.Accordion("⚙️ Settings", open=True):
        mode_box = gr.Radio(
            choices=[
                ("⚡ Fast — Qwen3.5-4B (everyday chat)", "fast"),
                ("🧠 Deep reasoning — Qwen3.5-9B (slower, for UPSC-depth analysis)", "deep"),
                ("👁️ Vision — Phi-3-Vision (ONNX, supports text + images)", "vision"),
            ],
            value="fast",
            label="Response Mode",
        )
        system_prompt_box = gr.Textbox(
            label="System Prompt",
            value=DEFAULT_SYSTEM_PROMPT,
            lines=4,
        )
        max_tokens_box = gr.Slider(
            label="Max response tokens",
            minimum=128, maximum=4096, value=1024, step=128,
        )
        temperature_box = gr.Slider(
            label="Temperature",
            minimum=0.0, maximum=1.5, value=0.7, step=0.1,
        )
        disable_thinking_box = gr.Checkbox(
            label="Disable thinking (faster replies — works around a known llama.cpp/Qwen3.5 bug)",
            value=True,
        )

    gr.ChatInterface(
        fn=respond,
        additional_inputs=[
            system_prompt_box, max_tokens_box, temperature_box, disable_thinking_box, mode_box,
        ],
        type="messages",
        multimodal=True,
        examples=[
            [{"text": "Extract all text from this image exactly as written.", "files": []}, "You are a helpful, knowledgeable assistant.", 1024, 0.7, True, "vision"],
            [{"text": "Describe the contents of this image in detail.", "files": []}, "You are a helpful, knowledgeable assistant.", 1024, 0.7, True, "vision"]
        ]
    )

if __name__ == "__main__":
    # Pre-load the default ("fast") model now, so the cold download/load cost
    # is paid once during Space startup (visible in logs) instead of hanging
    # a user's first chat message with no feedback. "deep" stays lazy-loaded
    # on first use since picking it is an explicit opt-in to wait longer.
    _load_llm("fast")
    
    # Also pre-download Phi-3-Vision model files so they are cached during startup/build phase
    try:
        download_phi3_vision_model()
    except Exception as e:
        print(f"[startup] Failed to pre-download Phi-3-Vision model files: {e}", flush=True)
        
    demo.launch()