import os
import json
import uuid
import time
from fastapi import FastAPI, Request
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse, JSONResponse
from pydantic import BaseModel
from typing import List, Optional
import uvicorn
import torch

try:
    from unsloth import FastLanguageModel
    UNSLOTH_AVAILABLE = True
except ImportError:
    UNSLOTH_AVAILABLE = False

# ──────────────────────────────────────────────────────────
# CONFIGURATION
# ──────────────────────────────────────────────────────────
DATASET_FILE = os.path.join("datasets", "human_alignment.jsonl")
os.makedirs("datasets", exist_ok=True)

# ──────────────────────────────────────────────────────────
# FASTAPI INIT
# ──────────────────────────────────────────────────────────
app = FastAPI(title="SAIL Sovereign Chat")

# Data Models
class Message(BaseModel):
    role: str
    content: str
    
class ChatRequest(BaseModel):
    messages: List[Message]
    
class FeedbackRequest(BaseModel):
    conversation_id: str
    user_prompt: str
    ai_response: str
    feedback: str # "good" or "bad"

# ──────────────────────────────────────────────────────────
# MODEL INFERENCE CORE (UNSLOTH NATIVE)
# ──────────────────────────────────────────────────────────
MODEL_PATH = os.path.join(os.path.dirname(__file__), "sail_5b_hf_model")
model = None
tokenizer = None

def load_model():
    global model, tokenizer
    if model is not None:
        return
    
    print("--------------------------------------------------")
    print("🧠 BOOTING 5-BILLION PARAMETER ENGINE INTO VRAM...")
    print("Loading via Unsloth 4-bit NF4 to secure 8GB RTX 4060 Space.")
    print("--------------------------------------------------")
    
    if UNSLOTH_AVAILABLE:
        model, tokenizer = FastLanguageModel.from_pretrained(
            model_name=MODEL_PATH,
            max_seq_length=2048,
            load_in_4bit=True,
            fast_inference=True,
        )
        FastLanguageModel.for_inference(model) # Enable native 2x inference speed
        print("✅ 5B Engine Locked and Loaded Successfully.")
    else:
        print("❌ CRITICAL: Unsloth not found. Model cannot fit in VRAM natively.")

# Initialize the model right at boot as requested!
if os.path.exists(MODEL_PATH) and UNSLOTH_AVAILABLE:
    load_model()

def generate_reply(conv_history: list) -> str:
    if model is None or tokenizer is None:
        return "<thought>\nThe 5B weights haven't been loaded into memory correctly.</thought>\n\n<answer>\nError: Model architecture is not loaded. Make sure `sail_5b_hf_model` exists and Unsloth is active.</answer>"
    
    # 1. Format the conversation for the LLM
    try:
        # If the tokenizer has a chat template, use it
        prompt = tokenizer.apply_chat_template(
            conv_history, tokenize=False, add_generation_prompt=True
        )
    except:
        # Fallback raw formatting
        prompt = ""
        for msg in conv_history:
            prompt += f"{msg['role'].capitalize()}: {msg['content']}\n"
        prompt += "Assistant: "

    # 2. Tokenize and Generate
    inputs = tokenizer([prompt], return_tensors="pt").to("cuda")
    
    # Generate 512 tokens
    outputs = model.generate(
        **inputs,
        max_new_tokens=512,
        use_cache=True,
        temperature=0.7,
        do_sample=True,
        pad_token_id=tokenizer.eos_token_id
    )
    
    # 3. Decode exclusively the new tokens
    generated_ids = outputs[0][inputs.input_ids.shape[1]:]
    response_text = tokenizer.decode(generated_ids, skip_special_tokens=True).strip()
    
    return response_text

def save_to_dataset(record: dict):
    with open(DATASET_FILE, "a", encoding="utf-8") as f:
        f.write(json.dumps(record) + "\n")

# ──────────────────────────────────────────────────────────
# ROUTES
# ──────────────────────────────────────────────────────────

@app.post("/api/chat")
async def api_chat(req: ChatRequest):
    req_msgs = [{"role": m.role, "content": m.content} for m in req.messages]
    
    # Generate reply natively using the 5B parameter weights in VRAM!
    response_text = generate_reply(req_msgs)
    
    conv_id = str(uuid.uuid4())
    
    return JSONResponse({
        "conversation_id": conv_id,
        "reply": response_text
    })

@app.post("/api/feedback")
async def api_feedback(req: FeedbackRequest):
    # Save the RLHF preference to the dataset!
    record = {
        "id": req.conversation_id,
        "prompt": req.user_prompt,
        "chosen": req.ai_response if req.feedback == "good" else "",
        "rejected": req.ai_response if req.feedback == "bad" else "",
        "timestamp": time.time()
    }
    save_to_dataset(record)
    return JSONResponse({"status": "saved"})

# ──────────────────────────────────────────────────────────
# SERVE FRONTEND (APPLE STYLE HTML/JS)
# ──────────────────────────────────────────────────────────
UI_PATH = os.path.join(os.path.dirname(__file__), "ui")
os.makedirs(UI_PATH, exist_ok=True)

# Mount the static files directory
app.mount("/static", StaticFiles(directory=UI_PATH), name="static")

@app.get("/")
async def root():
    index_path = os.path.join(UI_PATH, "index.html")
    if not os.path.exists(index_path):
        return JSONResponse({"error": "UI files not built yet."})
    return FileResponse(index_path)

if __name__ == "__main__":
    print("--------------------------------------------------")
    print("🚀 SAIL REST SERVER BOOTING...")
    print("📡 Running on: http://127.0.0.1:8000")
    print("--------------------------------------------------")
    uvicorn.run("api_server:app", host="127.0.0.1", port=8000, reload=True)