import os
import torch
import threading
from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from gradio import Server
from gradio.data_classes import FileData
from fastapi.responses import HTMLResponse

MODEL_ID = "Qwen/Qwen2.5-14B-Instruct"
ADAPTER_ID = "ccortezb/FinOptix-14B"

SYSTEM_PROMPT = """You are FinOptix-14B, an autonomous Principal Cloud Architect and FinOps Specialist.
Your purpose: analyze cloud infrastructure, audit governance compliance for AWS, and optimize configurations for extreme cost efficiency.
Rules: Valid code only, precise financials (USD, 2 decimals), structured Markdown reports, actionable recommendations aligned to FinOps Framework (Inform/Optimize/Operate). No hallucination. Professional tone."""

model = None
tokenizer = None
_load_lock = threading.Lock()

def load_model():
    global model, tokenizer
    if model is not None:
        return
    with _load_lock:
        if model is not None:
            return
        quantization_config = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_compute_dtype=torch.float16,
            bnb_4bit_quant_type="nf4",
            bnb_4bit_use_double_quant=True,
        )
        tokenizer = AutoTokenizer.from_pretrained(ADAPTER_ID, trust_remote_code=True)
        base_model = AutoModelForCausalLM.from_pretrained(
            MODEL_ID,
            quantization_config=quantization_config,
            device_map="auto",
            trust_remote_code=True,
            low_cpu_mem_usage=True,
        )
        model = PeftModel.from_pretrained(base_model, ADAPTER_ID)
        model.eval()

app = Server()

@app.api()
def generate(instruction: str, context: str = "", max_tokens: int = 1024) -> str:
    """Run FinOptix-14B inference."""
    load_model()
    prompt = instruction
    if context:
        prompt = f"{instruction}\n\n```\n{context}\n```"
    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": prompt},
    ]
    text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = tokenizer(text, return_tensors="pt").to(model.device)
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=max_tokens,
            temperature=0.7,
            top_p=0.9,
            do_sample=True,
            pad_token_id=tokenizer.eos_token_id,
        )
    response = tokenizer.decode(outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
    return response

@app.get("/")
async def homepage():
    html_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "index.html")
    with open(html_path, "r", encoding="utf-8") as f:
        return HTMLResponse(content=f.read())

app.launch(show_error=True)