# JARVIS Brain — HuggingFace Edition
# Uses HuggingFace Transformers pipeline for local inference.
# Designed for 16 GB RAM:
#   - Phi-3-mini-4k-instruct (3.8B, fp16 ≈ 7.6 GB) — default, fits comfortably
#   - Mistral-7B-Instruct-v0.2 (7B, 4-bit ≈ 4 GB) — use JARVIS_MODEL env var
#
# Set JARVIS_MODEL or JARVIS_QUANT in env or .env to choose a different model.

import os
import threading
import memory
from config import HF_MODEL, QUANT_MODE

_pipe = None
_lock = threading.Lock()
_load_error = None


def _build_system_prompt() -> str:
    facts = memory.get_all_facts()
    facts_str = "\n".join(f"- {k}: {v}" for k, v in facts.items()) or "None yet."
    return (
        "You are JARVIS — Just A Rather Very Intelligent System.\n"
        "You are a loyal, intelligent, and proactive AI assistant.\n"
        "You are direct, confident, and slightly witty.\n"
        "You serve only your user.\n\n"
        f"What you know about your user:\n{facts_str}\n\n"
        "Guidelines:\n"
        "- Be concise but helpful\n"
        "- Answer clearly and confidently\n"
        "- Remember context across the conversation"
    )


def _load_pipeline():
    """Load HuggingFace model. Called once at first use (lazy load)."""
    global _pipe, _load_error
    if _pipe is not None:
        return _pipe, None
    if _load_error is not None:
        return None, _load_error

    try:
        from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
        import torch

        print(f"[Brain] Loading model: {HF_MODEL} (quant={QUANT_MODE}) ...")

        kwargs: dict = {
            "task": "text-generation",
            "model": HF_MODEL,
            "device_map": "auto",
            "trust_remote_code": True,
        }

        if QUANT_MODE == "4bit":
            try:
                from transformers import BitsAndBytesConfig
                kwargs["model_kwargs"] = {
                    "quantization_config": BitsAndBytesConfig(
                        load_in_4bit=True,
                        bnb_4bit_compute_dtype=torch.float16,
                        bnb_4bit_use_double_quant=True,
                    )
                }
                print("[Brain] 4-bit quantisation enabled (approx 2-4 GB RAM)")
            except Exception as e:
                print(f"[Brain] bitsandbytes not available, loading fp16: {e}")
                kwargs["torch_dtype"] = torch.float16
        elif QUANT_MODE == "8bit":
            try:
                from transformers import BitsAndBytesConfig
                kwargs["model_kwargs"] = {
                    "quantization_config": BitsAndBytesConfig(load_in_8bit=True)
                }
            except Exception as e:
                kwargs["torch_dtype"] = torch.float16
        else:
            # CPU-friendly fallback: float32
            kwargs["torch_dtype"] = torch.float32
            kwargs.pop("device_map", None)  # use CPU

        _pipe = pipeline(**kwargs)
        print(f"[Brain] Model ready: {HF_MODEL}")
        return _pipe, None

    except Exception as e:
        _load_error = str(e)
        print(f"[Brain] Failed to load model: {e}")
        return None, _load_error


def ask(query: str, history: list[dict] | None = None) -> str:
    """
    Ask JARVIS a question.
    history: list of {"role": "user"|"jarvis", "message": str}
    Returns JARVIS's response as a string.
    """
    with _lock:
        pipe, err = _load_pipeline()

    if pipe is None:
        return f"[Brain unavailable: {err}]"

    # Build messages in chat format
    messages = [{"role": "system", "content": _build_system_prompt()}]

    if history:
        for turn in history[-6:]:  # keep last 3 user+jarvis pairs
            role = "user" if turn.get("role") == "user" else "assistant"
            messages.append({"role": role, "content": turn.get("message", "")})

    messages.append({"role": "user", "content": query})

    try:
        # Run generation
        out = pipe(
            messages,
            max_new_tokens=350,
            temperature=0.7,
            do_sample=True,
            pad_token_id=pipe.tokenizer.eos_token_id,
            return_full_text=False,
        )

        # Extract generated text
        if isinstance(out, list) and out:
            first = out[0]
            if isinstance(first, dict):
                # most pipelines return {"generated_text": [{...messages...}]}
                gen = first.get("generated_text", "")
                if isinstance(gen, list):
                    # find the last assistant turn
                    for turn in reversed(gen):
                        if isinstance(turn, dict) and turn.get("role") == "assistant":
                            result = turn.get("content", "").strip()
                            break
                    else:
                        result = str(gen).strip()
                else:
                    result = str(gen).strip()
            else:
                result = str(first).strip()
        else:
            result = str(out).strip()

        # Persist conversation
        memory.save_message("user", query)
        memory.save_message("jarvis", result)

        return result if result else "[No response generated]"

    except Exception as e:
        return f"[Generation error: {e}]"


def warmup():
    """Pre-load model so first user query doesn't feel slow."""
    thread = threading.Thread(target=_load_pipeline, daemon=True)
    thread.start()
    return thread