"""
AirClaw Server — RabbitLLM/AirLLM wrapped as an OpenAI-compatible API.
Runs 70B+ models on a single 4GB GPU. Zero API cost.
Tries RabbitLLM first (newer models), falls back to AirLLM.
"""

import os
import time
import uuid
import threading
import argparse
from typing import Optional, List

# ── FastAPI ───────────────────────────────────────────────────────────────────
try:
    from fastapi import FastAPI, HTTPException
    from pydantic import BaseModel
    import uvicorn
except ImportError:
    os.system("pip install fastapi uvicorn pydantic -q")
    from fastapi import FastAPI, HTTPException
    from pydantic import BaseModel
    import uvicorn

# ── Torch ─────────────────────────────────────────────────────────────────────
try:
    import torch
    HAS_CUDA = torch.cuda.is_available()
except ImportError:
    HAS_CUDA = False

# ── LLM Backend (RabbitLLM preferred, AirLLM fallback) ───────────────────────
LLM_BACKEND = None
try:
    from rabbitllm import AutoModel as RabbitModel
    LLM_BACKEND = "rabbitllm"
    print("[AirClaw] ✅ Using RabbitLLM backend (newer models supported)")
except ImportError:
    try:
        from airllm import AutoModel as RabbitModel
        LLM_BACKEND = "airllm"
        print("[AirClaw] ✅ Using AirLLM backend")
    except ImportError:
        RabbitModel = None
        print("[AirClaw] ⚠️  No LLM backend found. Run: pip install rabbitllm or pip install airllm")

# ── Config ────────────────────────────────────────────────────────────────────
DEFAULT_MODEL  = os.environ.get("AIRCLAW_MODEL", "mistralai/Mistral-7B-Instruct-v0.2")
DEFAULT_PORT   = int(os.environ.get("AIRCLAW_PORT", "4096"))
MAX_NEW_TOKENS = int(os.environ.get("AIRCLAW_MAX_TOKENS", "512"))
MAX_LENGTH     = 512

app = FastAPI(title="AirClaw — Local LLM API", version="2.0.0")

# ── State ─────────────────────────────────────────────────────────────────────
_model       = None
_model_name  = DEFAULT_MODEL
_model_lock  = threading.Lock()
_model_error = None


def load_model(model_id: str):
    global _model, _model_name, _model_error
    if RabbitModel is None:
        _model_error = "No LLM backend installed. Run: pip install rabbitllm"
        print(f"[AirClaw] ❌ {_model_error}")
        return
    print(f"\n[AirClaw] Loading: {model_id}")
    print("[AirClaw] First run downloads model weights (~4-14 GB). Subsequent starts are instant.\n")
    try:
        with _model_lock:
            _model = RabbitModel.from_pretrained(model_id)
            _model_name = model_id
        print(f"[AirClaw] ✅ Model ready: {model_id}\n")
    except Exception as e:
        _model_error = str(e)
        print(f"[AirClaw] ❌ Failed to load model: {e}\n")


# ── Schemas ───────────────────────────────────────────────────────────────────
class Message(BaseModel):
    role: str
    content: str

class ChatRequest(BaseModel):
    model: Optional[str] = None
    messages: List[Message]
    max_tokens: Optional[int] = None
    temperature: Optional[float] = 0.7
    stream: Optional[bool] = False

class CompletionRequest(BaseModel):
    model: Optional[str] = None
    prompt: str
    max_tokens: Optional[int] = None
    temperature: Optional[float] = 0.7


# ── Helpers ───────────────────────────────────────────────────────────────────
def messages_to_prompt(messages: List[Message]) -> str:
    parts = []
    system = ""
    for m in messages:
        if m.role == "system":
            system = m.content
        elif m.role == "user":
            if system:
                parts.append(f"<s>[INST] <<SYS>>\n{system}\n<</SYS>>\n\n{m.content} [/INST]")
                system = ""
            else:
                parts.append(f"<s>[INST] {m.content} [/INST]")
        elif m.role == "assistant":
            parts.append(f" {m.content} </s>")
    return "".join(parts)


def run_inference(prompt: str, max_new_tokens: int) -> str:
    if _model is None:
        if _model_error:
            raise HTTPException(status_code=500, detail=f"Model error: {_model_error}")
        raise HTTPException(status_code=503, detail="Model still loading. Wait 30-60 seconds and retry.")

    try:
        input_tokens = _model.tokenizer(
            [prompt],
            return_tensors="pt",
            return_attention_mask=False,
            truncation=True,
            max_length=MAX_LENGTH,
            padding=False,
        )
        input_ids = input_tokens['input_ids']
        if HAS_CUDA:
            input_ids = input_ids.cuda()

        with _model_lock:
            output = _model.generate(
                input_ids,
                max_new_tokens=max_new_tokens,
                use_cache=True,
                return_dict_in_generate=True,
            )

        full = _model.tokenizer.decode(output.sequences[0])
        if prompt in full:
            result = full[len(prompt):].strip()
        else:
            new_tokens = output.sequences[0][input_tokens['input_ids'].shape[1]:]
            result = _model.tokenizer.decode(new_tokens, skip_special_tokens=True).strip()

        return result.replace("</s>", "").replace("<s>", "").strip()

    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Inference error: {str(e)}")


# ── Routes ────────────────────────────────────────────────────────────────────
@app.get("/")
def root():
    return {
        "service": "AirClaw",
        "version": "2.0.0",
        "backend": LLM_BACKEND,
        "model": _model_name,
        "ready": _model is not None,
        "cuda": HAS_CUDA,
        "error": _model_error,
    }

@app.get("/health")
def health():
    return {
        "status": "ready" if _model is not None else ("error" if _model_error else "loading"),
        "model_loaded": _model is not None,
        "backend": LLM_BACKEND,
        "error": _model_error,
    }

@app.get("/v1/models")
def list_models():
    return {
        "object": "list",
        "data": [{"id": _model_name, "object": "model", "created": int(time.time()), "owned_by": "airclaw"}]
    }

@app.post("/v1/chat/completions")
def chat(req: ChatRequest):
    prompt  = messages_to_prompt(req.messages)
    text    = run_inference(prompt, req.max_tokens or MAX_NEW_TOKENS)
    return {
        "id": f"chatcmpl-{uuid.uuid4().hex[:12]}",
        "object": "chat.completion",
        "created": int(time.time()),
        "model": _model_name,
        "choices": [{"index": 0, "message": {"role": "assistant", "content": text}, "finish_reason": "stop"}],
        "usage": {"prompt_tokens": len(prompt.split()), "completion_tokens": len(text.split()), "total_tokens": len(prompt.split()) + len(text.split())},
    }

@app.post("/v1/completions")
def complete(req: CompletionRequest):
    text = run_inference(req.prompt, req.max_tokens or MAX_NEW_TOKENS)
    return {
        "id": f"cmpl-{uuid.uuid4().hex[:12]}",
        "object": "text_completion",
        "created": int(time.time()),
        "model": _model_name,
        "choices": [{"text": text, "index": 0, "finish_reason": "stop"}],
    }


# ── Entry ─────────────────────────────────────────────────────────────────────
def start_server(model: str = DEFAULT_MODEL, port: int = DEFAULT_PORT, host: str = "127.0.0.1"):
    threading.Thread(target=load_model, args=(model,), daemon=True).start()
    print(f"""
╔══════════════════════════════════════════════════════╗
║              AirClaw v2.0 — Local LLM                ║
║  Backend : {LLM_BACKEND:<41} ║
║  Model   : {model:<41} ║
║  API     : http://{host}:{port}/v1                ║
╚══════════════════════════════════════════════════════╝
  ⏳ Model loading in background...
  ✅ Check status: http://{host}:{port}/health
""")
    uvicorn.run(app, host=host, port=port, log_level="warning")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--model", default=DEFAULT_MODEL)
    parser.add_argument("--port",  default=DEFAULT_PORT, type=int)
    parser.add_argument("--host",  default="127.0.0.1")
    args = parser.parse_args()
    start_server(args.model, args.port, args.host)
