"""
aquin serve — local inference server for on-prem model diff.

Loads a base model + fine-tuned checkpoint locally, exposes two streaming
endpoints that the Aquin dashboard proxies to for the prompt tester UI.

Usage:
    aquin serve --base meta-llama/Llama-3.2-1B --ft ./my-checkpoint --port 7862

Then in the dashboard set:  Local server → http://localhost:7862
"""

from __future__ import annotations

import argparse
import json
import os
import sys
import threading
from typing import Iterator


def _load_models(base_id: str, ft_path: str):
    try:
        import torch
        from transformers import AutoModelForCausalLM, AutoTokenizer
    except ImportError:
        print("[aquin serve] requires torch + transformers: pip install aquin[ml]")
        sys.exit(1)

    print(f"[aquin serve] loading tokenizer from {base_id} …")
    tokenizer = AutoTokenizer.from_pretrained(base_id)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    dtype = torch.float16 if torch.cuda.is_available() else torch.float32
    device = "cuda" if torch.cuda.is_available() else "cpu"

    print(f"[aquin serve] loading base model {base_id} on {device} …")
    base_model = AutoModelForCausalLM.from_pretrained(
        base_id, device_map=device, torch_dtype=dtype
    )
    base_model.eval()

    # Resolve checkpoint file
    if os.path.isdir(ft_path):
        candidates = []
        for ext in (".safetensors", ".bin", ".pt"):
            candidates += [
                os.path.join(ft_path, f)
                for f in os.listdir(ft_path)
                if f.endswith(ext)
            ]
        if not candidates:
            print(f"[aquin serve] no checkpoint found in {ft_path}")
            sys.exit(1)
        ft_file = candidates[0]
    else:
        ft_file = ft_path

    print(f"[aquin serve] loading fine-tuned checkpoint {ft_file} …")

    # Try loading as a full HF model directory first, fall back to state dict patch
    ft_model = None
    if os.path.isdir(ft_path) and os.path.exists(os.path.join(ft_path, "config.json")):
        ft_model = AutoModelForCausalLM.from_pretrained(
            ft_path, device_map=device, torch_dtype=dtype
        )
        ft_model.eval()
    else:
        import torch
        ft_model = AutoModelForCausalLM.from_pretrained(
            base_id, device_map=device, torch_dtype=dtype
        )
        state = torch.load(ft_file, map_location=device, weights_only=True)
        if isinstance(state, dict) and "state_dict" in state:
            state = state["state_dict"]
        missing, unexpected = ft_model.load_state_dict(state, strict=False)
        if missing:
            print(f"[aquin serve] {len(missing)} keys not in checkpoint (using base weights)")
        ft_model.eval()

    print("[aquin serve] both models loaded — ready")
    return tokenizer, base_model, ft_model


def _generate_stream(model, tokenizer, prompt: str, max_new_tokens: int, temperature: float) -> Iterator[str]:
    import torch

    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    input_len = inputs["input_ids"].shape[-1]

    with torch.no_grad():
        output_ids = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            temperature=temperature if temperature > 0 else 1.0,
            do_sample=temperature > 0,
            pad_token_id=tokenizer.pad_token_id,
        )

    new_ids = output_ids[0][input_len:]
    # Decode token by token for streaming effect
    buffer = ""
    for i in range(len(new_ids)):
        partial = tokenizer.decode(new_ids[: i + 1], skip_special_tokens=True)
        new_text = partial[len(buffer):]
        buffer = partial
        if new_text:
            yield new_text


def run_server(base_id: str, ft_path: str, port: int, host: str = "0.0.0.0"):
    try:
        from fastapi import FastAPI, Request
        from fastapi.middleware.cors import CORSMiddleware
        from fastapi.responses import StreamingResponse
        import uvicorn
    except ImportError:
        print("[aquin serve] requires fastapi + uvicorn: pip install fastapi uvicorn")
        sys.exit(1)

    tokenizer, base_model, ft_model = _load_models(base_id, ft_path)
    _lock = threading.Lock()

    app = FastAPI(title="aquin-serve")
    app.add_middleware(
        CORSMiddleware,
        allow_origins=["*"],
        allow_methods=["*"],
        allow_headers=["*"],
    )

    @app.get("/health")
    def health():
        return {"status": "ok", "base": base_id, "ft": ft_path}

    @app.post("/generate/stream")
    async def generate_stream(request: Request):
        body = await request.json()
        prompt         = body.get("prompt", "")
        max_new_tokens = int(body.get("max_new_tokens", 256))
        temperature    = float(body.get("temperature", 0.7))

        def event_stream():
            with _lock:
                # Stream base model tokens
                base_tokens = []
                for chunk in _generate_stream(base_model, tokenizer, prompt, max_new_tokens, temperature):
                    base_tokens.append(chunk)
                    yield f"data: {json.dumps({'type': 'base', 'text': chunk})}\n\n"

                yield f"data: {json.dumps({'type': 'base_done', 'full': ''.join(base_tokens)})}\n\n"

                # Stream ft model tokens
                ft_tokens = []
                for chunk in _generate_stream(ft_model, tokenizer, prompt, max_new_tokens, temperature):
                    ft_tokens.append(chunk)
                    yield f"data: {json.dumps({'type': 'ft', 'text': chunk})}\n\n"

                yield f"data: {json.dumps({'type': 'done', 'full': ''.join(ft_tokens)})}\n\n"

        return StreamingResponse(event_stream(), media_type="text/event-stream")

    print(f"[aquin serve] listening on http://{host}:{port}")
    uvicorn.run(app, host=host, port=port, log_level="warning")


def main():
    parser = argparse.ArgumentParser(
        prog="aquin serve",
        description="Local inference server for Aquin model diff prompt tester",
    )
    parser.add_argument("--base",   required=True, help="HuggingFace base model ID or local path")
    parser.add_argument("--ft",     required=True, help="Fine-tuned checkpoint path (.pt/.bin/.safetensors or dir)")
    parser.add_argument("--port",   type=int, default=7862, help="Port to listen on (default: 7862)")
    parser.add_argument("--host",   default="0.0.0.0", help="Host to bind (default: 0.0.0.0)")
    args = parser.parse_args()

    run_server(base_id=args.base, ft_path=args.ft, port=args.port, host=args.host)


if __name__ == "__main__":
    main()