import os
import re
import time
import traceback
from pathlib import Path
from typing import Any, Dict, List, Tuple
from config import *
import torch
from fastapi import Request
from fastapi.responses import HTMLResponse, JSONResponse
from fastapi.staticfiles import StaticFiles
from gradio import Server
from transformers import AutoModelForCausalLM, AutoTokenizer

# ── ZeroGPU Detection ───────────────────────────────────────────────────
IS_ZEROGPU = False
try:
    import spaces
    from spaces.config import Config as _ZeroGPUConfig
    IS_ZEROGPU = bool(_ZeroGPUConfig.zero_gpu)
except ImportError:
    class _DummySpaces:
        @staticmethod
        def GPU(fn):
            return fn
    spaces = _DummySpaces()

# ── Configuration ───────────────────────────────────────────────────────
DEFAULT_SYSTEM_PROMPT = (
    "You are ChatUI, a helpful local AI assistant. "
    "You provide accurate, detailed, and thoughtful responses."
)

STATIC_DIR = Path(__file__).parent / "static"
os.makedirs(STATIC_DIR, exist_ok=True)

# ── Fallback chat templates for tokenizers missing .chat_template ───────
_FALLBACK_TEMPLATES = {
    "chatml": (
        "{% for message in messages %}"
        "{{'<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n'}}"
        "{% endfor %}"
        "{% if add_generation_prompt %}"
        "{{'<|im_start|>assistant\n'}}"
        "{% endif %}"
    ),
    "llama2": (
        "{% if messages[0]['role'] == 'system' %}"
        "{{'<<SYS>>\n' + messages[0]['content'] + '\n<</SYS>>\n\n'}}"
        "{% set loop_messages = messages[1:] %}"
        "{% else %}"
        "{% set loop_messages = messages %}"
        "{% endif %}"
        "{% for message in loop_messages %}"
        "{% if message['role'] == 'user' %}"
        "{{'[INST] ' + message['content'] + ' [/INST]'}}"
        "{% elif message['role'] == 'assistant' %}"
        "{{message['content'] + '</s>'}}"
        "{% endif %}"
        "{% endfor %}"
    ),
    "phi3": (
        "{% for message in messages %}"
        "{% if message['role'] == 'system' %}"
        "{{'<|system|>\n' + message['content'] + '<|end|>\n'}}"
        "{% elif message['role'] == 'user' %}"
        "{{'<|user|>\n' + message['content'] + '<|end|>\n'}}"
        "{% elif message['role'] == 'assistant' %}"
        "{{'<|assistant|>\n' + message['content'] + '<|end|>\n'}}"
        "{% endif %}"
        "{% endfor %}"
        "{% if add_generation_prompt %}"
        "{{'<|assistant|>\n'}}"
        "{% endif %}"
    ),
    "qwen": (
        "{% for message in messages %}"
        "{% if message['role'] == 'system' %}"
        "{{'<|im_start|>system\n' + message['content'] + '<|im_end|>\n'}}"
        "{% elif message['role'] == 'user' %}"
        "{{'<|im_start|>user\n' + message['content'] + '<|im_end|>\n'}}"
        "{% elif message['role'] == 'assistant' %}"
        "{{'<|im_start|>assistant\n' + message['content'] + '<|im_end|>\n'}}"
        "{% endif %}"
        "{% endfor %}"
        "{% if add_generation_prompt %}"
        "{{'<|im_start|>assistant\n'}}"
        "{% endif %}"
    ),
}

def _pick_fallback_template(model_id: str) -> str:
    mid = model_id.lower()
    if "phi" in mid:
        return _FALLBACK_TEMPLATES["phi3"]
    if "qwen" in mid:
        return _FALLBACK_TEMPLATES["qwen"]
    if "tinyllama" in mid or "llama" in mid:
        return _FALLBACK_TEMPLATES["llama2"]
    return _FALLBACK_TEMPLATES["chatml"]

def _ensure_chat_template(tokenizer, model_id: str):
    if getattr(tokenizer, "chat_template", None):
        return
    tpl = _pick_fallback_template(model_id)
    tokenizer.chat_template = tpl
    print(f"[chat_template] injected fallback template for {model_id}")


# ── Model cache (lazy) ──────────────────────────────────────────────────
_MODEL_CACHE: Dict[str, Dict[str, Any]] = {}

def _device() -> str:
    if IS_ZEROGPU:
        return "cuda"
    if torch.cuda.is_available():
        return "cuda"
    if getattr(torch.backends, "mps", None) and torch.backends.mps.is_available():
        return "mps"
    return "cpu"

def _dtype():
    if IS_ZEROGPU or torch.cuda.is_available():
        return torch.float16
    return torch.float32

def get_model(model_id: str) -> Dict[str, Any]:
    if model_id in _MODEL_CACHE:
        return _MODEL_CACHE[model_id]
    if model_id not in MODEL_REGISTRY:
        raise ValueError(f"Unknown model: {model_id}")

    device = _device()
    dtype = _dtype()

    tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
    if tokenizer.pad_token_id is None:
        tokenizer.pad_token = tokenizer.eos_token

    _ensure_chat_template(tokenizer, model_id)

    model = AutoModelForCausalLM.from_pretrained(
        model_id,
        torch_dtype=dtype,
        trust_remote_code=True,
        low_cpu_mem_usage=True,
    ).to(device)
    model.eval()

    _MODEL_CACHE[model_id] = {
        "tokenizer": tokenizer,
        "model": model,
        "device": device,
        "loaded_at": time.time(),
    }
    return _MODEL_CACHE[model_id]

# ── Generation core ─────────────────────────────────────────────────────
def _flatten_content(content) -> str:
    if isinstance(content, str):
        return content
    parts: List[str] = []
    for c in content or []:
        if isinstance(c, dict) and c.get("type") == "text":
            parts.append(c.get("text", ""))
    return "\n".join(parts).strip()

def _build_prompt(model_id: str, messages: List[Dict[str, Any]],
                  system_prompt: str) -> str:
    bundle = get_model(model_id)
    tokenizer = bundle["tokenizer"]

    flat_messages: List[Dict[str, str]] = []
    if system_prompt and system_prompt.strip():
        flat_messages.append({"role": "system", "content": system_prompt.strip()})

    for m in messages:
        role = m.get("role", "user")
        if role == "system":
            continue
        text = _flatten_content(m.get("content", ""))
        if text:
            flat_messages.append({"role": role, "content": text})

    if getattr(tokenizer, "chat_template", None):
        try:
            return tokenizer.apply_chat_template(
                flat_messages, tokenize=False, add_generation_prompt=True
            )
        except Exception as e:
            print(f"[chat_template] apply failed for {model_id}: {e}; using manual fallback")

    # Manual fallback
    lines = [f"{m['role']}: {m['content']}" for m in flat_messages]
    lines.append("assistant:")
    return "\n".join(lines) + "\n"

@spaces.GPU
def generate(model_id: str, messages: List[Dict[str, Any]],
             system_prompt: str, max_tokens: int, temperature: float,
             top_p: float = 0.95) -> Tuple[str, str]:
    bundle = get_model(model_id)
    tokenizer = bundle["tokenizer"]
    model = bundle["model"]
    device = bundle["device"]

    prompt = _build_prompt(model_id, messages, system_prompt)
    inputs = tokenizer(prompt, return_tensors="pt").to(device)

    do_sample = temperature > 0.01
    with torch.inference_mode():
        out = model.generate(
            **inputs,
            max_new_tokens=max_tokens,
            do_sample=do_sample,
            temperature=max(temperature, 1e-3) if do_sample else 1.0,
            top_p=top_p if do_sample else 1.0,
            pad_token_id=tokenizer.pad_token_id,
            eos_token_id=tokenizer.eos_token_id,
        )

    new_tokens = out[0][inputs["input_ids"].shape[1]:]
    text = tokenizer.decode(new_tokens, skip_special_tokens=True).strip()

    reasoning = ""
    think_match = re.search(r"<think[^>]*>(.*?)</think\s*>", text, re.DOTALL | re.IGNORECASE)
    if think_match:
        reasoning = think_match.group(1).strip()
        text = re.sub(r"<think[^>]*>.*?</think\s*>", "", text, flags=re.DOTALL | re.IGNORECASE).strip()

    return text, reasoning

# ── Agent tools ─────────────────────────────────────────────────────────
def _tool_calculator(expression: str) -> str:
    try:
        if not re.fullmatch(r"[\d+\-*/().%\s]+", expression):
            return "Error: only numeric expressions are allowed."
        result = eval(expression, {"__builtins__": {}}, {})
        return f"= {result}"
    except Exception as e:
        return f"Error: {e}"

def _tool_web_search(query: str) -> str:
    return (
        f"[stub search] I cannot reach the web (no API key configured), "
        f"but for query '{query}' you might want to check Wikipedia, "
        f"Stack Overflow, or arXiv. Use the calculator for any maths."
    )

def _tool_file_read(path: str) -> str:
    safe_root = Path(__file__).parent.resolve()
    try:
        target = (safe_root / path).resolve()
        target.relative_to(safe_root)
        if not target.exists() or not target.is_file():
            return f"Error: file not found: {path}"
        text = target.read_text(encoding="utf-8", errors="ignore")
        if len(text) > 4000:
            text = text[:4000] + "\n…[truncated]"
        return text
    except Exception as e:
        return f"Error: {e}"

def _tool_now(_arg: str = "") -> str:
    return time.strftime("%Y-%m-%d %H:%M:%S %Z")

AGENT_TOOLS = {
    "calculator": _tool_calculator,
    "web_search": _tool_web_search,
    "file_read":  _tool_file_read,
    "now":        _tool_now,
}

AGENT_TOOL_DESCRIPTION = """You have access to the following tools. To use one,
reply with a single line that starts with `TOOL:` followed by the tool name and
its argument in parentheses, e.g.:

TOOL: calculator(2 * (3 + 4))
TOOL: web_search(latest news about mars)
TOOL: file_read(static/index.html)
TOOL: now()

After receiving the tool output, you may use another tool or give your final
answer. When you are done, start your final message with `FINAL:`.
"""

@spaces.GPU
def run_agent(model_id: str, user_text: str, system_prompt: str,
              max_tokens: int, temperature: float,
              max_iters: int = 4) -> Dict[str, Any]:
    history = [
        {"role": "system", "content": system_prompt + "\n\n" + AGENT_TOOL_DESCRIPTION},
        {"role": "user", "content": user_text},
    ]
    trace: List[Dict[str, str]] = []

    for i in range(max_iters):
        answer, _reason = generate(
            model_id, history,
            system_prompt="",
            max_tokens=max_tokens, temperature=temperature,
        )
        answer_stripped = answer.strip()

        if answer_stripped.upper().startswith("FINAL:"):
            final = answer_stripped[6:].strip()
            trace.append({"step": i + 1, "type": "final", "content": final})
            return {"final": final, "trace": trace}

        m = re.match(r"\s*TOOL:\s*([A-Za-z_]+)\s*\((.*)\)\s*$",
                     answer_stripped, re.DOTALL)
        if not m:
            trace.append({"step": i + 1, "type": "final", "content": answer_stripped})
            return {"final": answer_stripped, "trace": trace}

        tool_name = m.group(1).strip().lower()
        tool_arg = m.group(2).strip()
        trace.append({
            "step": i + 1,
            "type": "tool_call",
            "tool": tool_name,
            "argument": tool_arg,
        })

        fn = AGENT_TOOLS.get(tool_name)
        if fn is None:
            tool_result = (f"Error: unknown tool '{tool_name}'. "
                           f"Available: {list(AGENT_TOOLS.keys())}")
        else:
            try:
                tool_result = fn(tool_arg)
            except Exception as e:
                tool_result = f"Error: {e}"

        trace.append({
            "step": i + 1,
            "type": "tool_result",
            "tool": tool_name,
            "output": tool_result,
        })

        history.append({"role": "assistant", "content": answer_stripped})
        history.append({
            "role": "user",
            "content": (f"TOOL OUTPUT: {tool_result}\nNow either call another "
                        f"tool or finish with FINAL: <your answer>."),
        })

    return {
        "final": "I ran out of steps before producing a final answer.",
        "trace": trace,
    }

# ── Gradio Server App ───────────────────────────────────────────────────
app = Server(title="ChatUI Studio", version="2.0.0")

@app.get("/api/models")
async def api_list_models():
    return JSONResponse({
        "status": "success",
        "default": DEFAULT_MODEL,
        "models": [
            {"id": mid, **meta}
            for mid, meta in MODEL_REGISTRY.items()
        ],
        "cached": list(_MODEL_CACHE.keys()),
    })

@app.get("/api/tasks")
async def api_list_tasks():
    return JSONResponse({
        "status": "success",
        "tasks": [
            {"id": tid, **meta}
            for tid, meta in TASK_REGISTRY.items()
        ],
    })

@app.post("/api/load_model")
async def api_load_model(req: Request):
    body = await req.json()
    model_id = body.get("model_id") or DEFAULT_MODEL
    try:
        if model_id not in MODEL_REGISTRY:
            return JSONResponse({"status": "error", "message": "Unknown model."})
        get_model(model_id)
        return JSONResponse({
            "status": "success",
            "model": model_id,
            "device": _MODEL_CACHE[model_id]["device"],
        })
    except Exception as e:
        traceback.print_exc()
        return JSONResponse({"status": "error", "message": str(e)})

@app.post("/api/run")
async def api_run_inference(req: Request):
    try:
        body = await req.json()
        messages: List[Dict[str, Any]] = body.get("messages", [])
        model_id = body.get("model_id") or DEFAULT_MODEL
        task = body.get("task") or "chat"
        max_tokens_i = int(body.get("max_tokens", 512))
        temperature_f = float(body.get("temperature", 0.7))
        system_prompt = body.get("system_prompt", "") or ""

        if model_id not in MODEL_REGISTRY:
            return JSONResponse({"status": "error", "message": f"Unknown model: {model_id}"})

        task_meta = TASK_REGISTRY.get(task, TASK_REGISTRY["chat"])
        sys_prompt = (system_prompt.strip() or task_meta["system_prompt"])

        if task == "agent":
            user_text = _flatten_content(messages[-1].get("content", "")) if messages else ""
            result = run_agent(
                model_id=model_id,
                user_text=user_text,
                system_prompt=sys_prompt,
                max_tokens=max_tokens_i,
                temperature=temperature_f,
            )
            return JSONResponse({
                "status": "success",
                "task": "agent",
                "model": model_id,
                "content": result["final"],
                "trace": result["trace"],
                "reasoning_content": "",
            })

        content, reasoning = generate(
            model_id=model_id,
            messages=messages,
            system_prompt=sys_prompt,
            max_tokens=max_tokens_i,
            temperature=temperature_f,
        )
        return JSONResponse({
            "status": "success",
            "task": task,
            "model": model_id,
            "content": content,
            "reasoning_content": reasoning or "",
        })

    except Exception as e:
        traceback.print_exc()
        return JSONResponse({"status": "error", "message": str(e)})

@app.get("/")
async def homepage():
    html_path = STATIC_DIR / "index.html"
    if html_path.exists():
        with open(html_path, "r", encoding="utf-8") as f:
            return HTMLResponse(content=f.read(), status_code=200)
    return HTMLResponse(
        content="<h1>Frontend is building. Please refresh in a few seconds...</h1>",
        status_code=200,
    )

@app.get("/healthz")
async def healthz():
    return JSONResponse({
        "status": "ok",
        "zerogpu": IS_ZEROGPU,
        "device": _device(),
        "models_cached": list(_MODEL_CACHE.keys()),
        "models_registered": list(MODEL_REGISTRY.keys()),
    })

app.mount("/static", StaticFiles(directory=str(STATIC_DIR)), name="static")

if __name__ == "__main__":
    import argparse
    import inspect

    parser = argparse.ArgumentParser(
        description="ChatUI Studio — local AI chat server",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument(
        "--host", "--server-name", dest="server_name", default="0.0.0.0",
        help="Host/IP to bind to.",
    )
    parser.add_argument(
        "--port", "--server-port", dest="server_port", type=int, default=7860,
        help="Port to listen on.",
    )
    parser.add_argument(
        "--share", dest="share", action="store_true",
        help="Create a public Gradio share link.",
    )
    parser.add_argument(
        "--no-share", dest="share", action="store_false",
        help="Disable public share link.",
    )
    parser.set_defaults(share=False)
    parser.add_argument(
        "--show-error", dest="show_error", action="store_true",
        help="Show full error tracebacks in the UI.",
    )
    parser.add_argument(
        "--no-show-error", dest="show_error", action="store_false",
        help="Hide error tracebacks in the UI.",
    )
    parser.set_defaults(show_error=True)
    parser.add_argument(
        "--quiet", action="store_true",
        help="Suppress Gradio startup messages.",
    )
    parser.add_argument(
        "--auth", nargs=2, metavar=("USER", "PASS"), default=None,
        help="Enable basic auth with USER and PASS.",
    )
    parser.add_argument(
        "--ssl-keyfile", default=None,
        help="Path to SSL key file (enables HTTPS).",
    )
    parser.add_argument(
        "--ssl-certfile", default=None,
        help="Path to SSL certificate file.",
    )
    parser.add_argument(
        "--ssl-keyfile-password", default=None,
        help="Password for the SSL key file, if any.",
    )
    parser.add_argument(
        "--favicon-path", default=None,
        help="Path to a custom favicon.",
    )
    parser.add_argument(
        "--allowed-paths", nargs="*", default=None,
        help="Extra paths Gradio is allowed to serve from.",
    )
    parser.add_argument(
        "--blocked-paths", nargs="*", default=None,
        help="Paths Gradio must never serve.",
    )
    parser.add_argument(
        "--max-file-size", default=None,
        help="Max upload size (e.g. '10mb', '1gb').",
    )

    args = parser.parse_args()

    launch_kwargs = {
        "server_name": args.server_name,
        "server_port": args.server_port,
        "share": args.share,
        "show_error": args.show_error,
        "quiet": args.quiet,
        "auth": tuple(args.auth) if args.auth else None,
    }
    if args.ssl_keyfile:
        launch_kwargs["ssl_keyfile"] = args.ssl_keyfile
    if args.ssl_certfile:
        launch_kwargs["ssl_certfile"] = args.ssl_certfile
    if args.ssl_keyfile_password:
        launch_kwargs["ssl_keyfile_password"] = args.ssl_keyfile_password
    if args.favicon_path:
        launch_kwargs["favicon_path"] = args.favicon_path
    if args.allowed_paths:
        launch_kwargs["allowed_paths"] = args.allowed_paths
    if args.blocked_paths:
        launch_kwargs["blocked_paths"] = args.blocked_paths
    if args.max_file_size:
        launch_kwargs["max_file_size"] = args.max_file_size

    # Drop kwargs that aren't supported by this Gradio version.
    try:
        supported = set(inspect.signature(app.launch).parameters.keys())
        if "kwargs" not in supported:
            launch_kwargs = {k: v for k, v in launch_kwargs.items() if k in supported}
    except (TypeError, ValueError):
        pass

    app.launch(**launch_kwargs)