from fastapi import FastAPI
from pydantic import BaseModel
from .llm_utils import simulate_search
from .umls_linker import link_umls
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
import functools

ANSWER_MODEL = "sunhaonlp/SearchSimulation_14B"

@functools.lru_cache(maxsize=1)
def _load_answer_pipe():
    tokenizer = AutoTokenizer.from_pretrained(ANSWER_MODEL)
    model = AutoModelForCausalLM.from_pretrained(
        ANSWER_MODEL,
        trust_remote_code=True,
        device_map="auto"
    )
    return pipeline(
        "text-generation",
        model=model,
        tokenizer=tokenizer,
        max_new_tokens=256,
        do_sample=False,
        temperature=0.0,
    )

class Query(BaseModel):
    question: str

app = FastAPI(
    title="ZeroSearch Medical Q&A API",
    description="Ask clinical questions; get answers with UMLS links, no external search APIs.",
    version="0.1.0",
)

@app.post("/ask")
def ask(query: Query):
    docs = simulate_search(query.question, k=5)
    context = "\n\n".join(docs)
    prompt = (
        "Answer the medical question strictly based on the provided context.\n\n"
        f"Context:\n{context}\n\n"
        f"Question: {query.question}\nAnswer:"
    )
    answer_pipe = _load_answer_pipe()
    answer = (
        answer_pipe(prompt, num_return_sequences=1)[0]["generated_text"]
        .split("Answer:")[-1].strip()
    )
    umls = link_umls(answer)
    return {"answer": answer, "docs": docs, "umls": umls}