from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse
from pydantic import BaseModel
import torch
import time
import json
import os
from typing import List, Dict

import JiRackTernaryPyTorch_1b_inf as base
from transformers import AutoTokenizer

app = FastAPI(title="JiRack Packed Server")

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

# ========================= CONFIG =========================
MODEL_PATH = "orca_packed.safetensors"
DEVICE = torch.device("cpu")

print(f"🔄 Loading model: {MODEL_PATH} on {DEVICE}")

config = base.TernaryConfig()
model = base.TernaryTransformer1B(config)

try:
    model.load_prod_weights(MODEL_PATH, DEVICE)
    print("✅ Model loaded via load_prod_weights()")
except Exception as e:
    print(f"❌ Load error: {e}")
    exit(1)

tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.2-1B-Instruct")
print("✅ Server ready on http://0.0.0.0:7869")

@torch.no_grad()
async def chat_generate(prompt: str, temperature: float = 0.7, max_tokens: int = 512):
    p = f"<|begin_of_text|><|start_header_id|>user<|end_header_id|>\n\n{prompt}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
    ids = tokenizer.encode(p, return_tensors="pt").to(DEVICE)
    generated = []
    for _ in range(max_tokens):
        logits, _ = model(ids[:, -1024:])
        next_token_logits = logits[:, -1, :]

        probs = torch.softmax(next_token_logits / temperature, dim=-1)
        nxt = torch.multinomial(probs, 1)
        token_id = nxt.item()

        if token_id in [tokenizer.eos_token_id, 128001, 128009]:
            break

        generated.append(token_id)
        token_str = tokenizer.decode([token_id], skip_special_tokens=True)
        yield token_str

        ids = torch.cat([ids, nxt], dim=-1)

    torch.cuda.empty_cache()


class ChatRequest(BaseModel):
    messages: List[Dict]
    model: str = "jirack"
    temperature: float = 0.7
    max_tokens: int = 512
    stream: bool = False

@app.post("/v1/chat/completions")
async def chat_completions(request: ChatRequest):
    try:
        user_content = request.messages[-1]["content"]

        if request.stream:
            async def stream():
                async for token in chat_generate(user_content, request.temperature, request.max_tokens):
                    yield f'data: {json.dumps({"choices": [{"delta": {"content": token}}]})}\n\n'
                yield "data: [DONE]\n\n"
            return StreamingResponse(stream(), media_type="text/event-stream")
        else:
            response_text = ""
            async for token in chat_generate(user_content, request.temperature, request.max_tokens):
                response_text += token

            return {
                "id": f"chatcmpl-{int(time.time())}",
                "object": "chat.completion",
                "created": int(time.time()),
                "model": request.model,
                "choices": [{
                    "index": 0,
                    "message": {"role": "assistant", "content": response_text},
                    "finish_reason": "stop"
                }]
            }

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


@app.get("/health")
async def health():
    return {"status": "ok"}

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(
        app,
        host="0.0.0.0",
        port=7869,
        timeout_keep_alive=600,
        workers=1,
        log_level="info"
    )