import os

import torch
from fastapi import FastAPI, Header, HTTPException
from pydantic import BaseModel
from peft import PeftConfig, PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

ADAPTER_MODEL_ID = os.environ.get(
    "ADAPTER_MODEL_ID",
    "maimd/Maimd-HPI-SFT-Behavioral-MedGemma-4B-v001-20260608",
).strip()
MODEL_ID = os.environ.get("MODEL_ID", "").strip()
HF_TOKEN = os.environ.get("HF_TOKEN")
REMOTE_API_KEY = os.environ.get("REMOTE_API_KEY", "")

app = FastAPI(title="Agentic Dr Inference")

model = None
tokenizer = None
resolved_model_id = None
resolved_base_model_id = None


class GenerateRequest(BaseModel):
    prompt: str
    max_new_tokens: int = 300


def load_model():
    global model, tokenizer, resolved_model_id, resolved_base_model_id

    if model is not None and tokenizer is not None:
        return

    device = "cuda" if torch.cuda.is_available() else "cpu"
    quantization_config = None

    if device == "cuda":
        quantization_config = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_compute_dtype=torch.bfloat16,
            bnb_4bit_use_double_quant=True,
            bnb_4bit_quant_type="nf4",
        )

    target_model_id = ADAPTER_MODEL_ID or MODEL_ID
    if not target_model_id:
        raise RuntimeError("Either ADAPTER_MODEL_ID or MODEL_ID must be configured.")

    resolved_model_id = target_model_id
    adapter_model_id = None

    try:
        peft_config = PeftConfig.from_pretrained(
            target_model_id,
            token=HF_TOKEN,
        )
        adapter_model_id = target_model_id
        resolved_base_model_id = peft_config.base_model_name_or_path
    except Exception:
        resolved_base_model_id = None

    tokenizer_source = resolved_base_model_id or target_model_id
    tokenizer = AutoTokenizer.from_pretrained(
        tokenizer_source,
        token=HF_TOKEN,
    )
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    base_model = AutoModelForCausalLM.from_pretrained(
        resolved_base_model_id or target_model_id,
        token=HF_TOKEN,
        torch_dtype=torch.bfloat16 if device == "cuda" else torch.float32,
        device_map="auto" if device == "cuda" else None,
        low_cpu_mem_usage=True,
        quantization_config=quantization_config,
    )

    if adapter_model_id:
        model = PeftModel.from_pretrained(
            base_model,
            adapter_model_id,
            token=HF_TOKEN,
        ).merge_and_unload()
    else:
        model = base_model

    model.eval()


@app.get("/health")
def health():
    return {
        "status": "ok",
        "model_id": resolved_model_id or MODEL_ID or ADAPTER_MODEL_ID,
        "base_model_id": resolved_base_model_id,
    }


@app.post("/generate")
def generate(request: GenerateRequest, authorization: str | None = Header(default=None)):
    if REMOTE_API_KEY:
        expected = f"Bearer {REMOTE_API_KEY}"
        if authorization != expected:
            raise HTTPException(status_code=401, detail="unauthorized")

    if not request.prompt.strip():
        raise HTTPException(status_code=400, detail="prompt is required")

    load_model()

    messages = [
        {
            "role": "user",
            "content": request.prompt,
        }
    ]

    prompt_text = tokenizer.apply_chat_template(
        messages,
        add_generation_prompt=True,
        tokenize=False,
    )
    inputs = tokenizer(prompt_text, return_tensors="pt")

    model_input_device = next(model.parameters()).device
    inputs = {
        key: value.to(model_input_device) if isinstance(value, torch.Tensor) else value
        for key, value in inputs.items()
    }

    input_len = inputs["input_ids"].shape[-1]

    with torch.inference_mode():
        outputs = model.generate(
            **inputs,
            max_new_tokens=request.max_new_tokens,
            do_sample=False,
            pad_token_id=tokenizer.eos_token_id,
        )

    generated_tokens = outputs[0][input_len:]
    text = tokenizer.decode(generated_tokens, skip_special_tokens=True).strip()

    return {"text": text}