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

# Hugging Face model ID you just uploaded
MODEL_ID = os.getenv("MODEL_ID", "milaadesign/helper-qwen-1_5b")
MODEL_MAX_LEN = int(os.getenv("MODEL_MAX_LEN", "512"))

app = FastAPI(title="Patient Helper API")

# CORS so your DreamHost site can call it from the browser
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # later you can restrict to your domain
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

class GenerateIn(BaseModel):
    prompt: str
    max_new_tokens: int = 128
    temperature: float = 0.2
    top_p: float = 0.9
    repetition_penalty: float = 1.05

device = torch.device("cpu")  # Spaces CPU Basic

torch.set_num_threads(int(os.getenv("TORCH_NUM_THREADS", "2")))
os.environ.setdefault("OMP_NUM_THREADS", "2")

print(f"Loading tokenizer from {MODEL_ID}...")
tok = AutoTokenizer.from_pretrained(MODEL_ID)
if tok.pad_token is None:
    tok.pad_token = tok.eos_token
tok.truncation_side = "left"
tok.model_max_length = MODEL_MAX_LEN

print(f"Loading model from {MODEL_ID}...")
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.float32,
    low_cpu_mem_usage=True,
)
model.to(device)
model.eval()
print("Model loaded.")

SYSTEM = (
    "You are a supportive, non-clinical assistant. "
    "Offer gentle, practical tips on how to support a patient. "
    "Do not diagnose or give unsafe advice. Encourage professional help when needed."
)

@app.post("/generate")
def generate(body: GenerateIn):
    messages = [
        {"role": "system", "content": SYSTEM},
        {"role": "user", "content": body.prompt.strip()},
    ]
    text = tok.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = tok(text, return_tensors="pt", truncation=True).to(device)

    with torch.no_grad():
        out = model.generate(
            **inputs,
            max_new_tokens=body.max_new_tokens,
            temperature=body.temperature,
            top_p=body.top_p,
            repetition_penalty=body.repetition_penalty,
            do_sample=body.temperature > 0,
            pad_token_id=tok.eos_token_id,
        )

    resp = tok.decode(out[0], skip_special_tokens=True)
    return {"text": resp}