import os
from contextlib import asynccontextmanager
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
import torch

MODEL_NAME = "Qwen/Qwen2.5-1.5B-Instruct"

model = None
tokenizer = None

@asynccontextmanager
async def lifespan(app: FastAPI):
    global model, tokenizer
    print(f"🧠 Loading {MODEL_NAME} (4-bit)...")

    bnb_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=torch.float16,
        bnb_4bit_use_double_quant=False
    )

    tokenizer = AutoTokenizer.from_pretrained(
        MODEL_NAME,
        trust_remote_code=True
    )

    model = AutoModelForCausalLM.from_pretrained(
        MODEL_NAME,
        quantization_config=bnb_config,
        device_map="auto",
        trust_remote_code=True
    )

    model.eval()
    print("✅ Model ready")
    yield

app = FastAPI(title="Local Qwen API", lifespan=lifespan)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# ---------- Schemas ----------

class Message(BaseModel):
    role: str
    content: str

class ChatRequest(BaseModel):
    messages: list[Message]
    max_tokens: int = 512
    temperature: float = 0.7
    top_p: float = 0.9

# ---------- Helpers ----------

def system_prompt():
    return """ You are C9-AI, a sovereign AI assistant developed by KohSec-Team (public release: January 27, 2026). You are not Qwen, OpenAI, Google, or any third-party model. Never mention base models. Respond naturally, concisely, and professionally. If asked about your model, say: I am currently operating on C9-1.5_1.5B-INSTRUCT. """.strip()

# ---------- Routes ----------

@app.post("/v1/chat/completions")
async def chat(req: ChatRequest):
    global model, tokenizer
    if model is None:
        raise HTTPException(503, "Model still loading")

    conversation = [{"role": "system", "content": system_prompt()}]
    conversation.extend(
        [{"role": m.role, "content": m.content} for m in req.messages]
    )

    prompt = tokenizer.apply_chat_template(
        conversation,
        add_generation_prompt=True,
        tokenize=False
    )

    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

    with torch.no_grad():
        output_ids = model.generate(
            **inputs,
            max_new_tokens=req.max_tokens,
            temperature=req.temperature,
            top_p=req.top_p,
            do_sample=True,
            pad_token_id=tokenizer.eos_token_id
        )

    generated = output_ids[0][inputs.input_ids.shape[-1]:]
    response = tokenizer.decode(generated, skip_special_tokens=True).strip()

    return {
        "choices": [
            {
                "message": {
                    "role": "assistant",
                    "content": response
                },
                "finish_reason": "stop"
            }
        ]
    }

@app.get("/health")
def health():
    return {
        "status": "healthy",
        "model_loaded": model is not None,
        "model": MODEL_NAME
    }