# app.py - Hugging Face Spaces version with auto-extract
import os
import zipfile
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
import chromadb
from sentence_transformers import SentenceTransformer
import gradio as gr

# Database path
DB_PATH = "./medqa_db"
ZIP_PATH = "./medqa_db.zip"

# Extract database if needed
if not os.path.exists(DB_PATH) and os.path.exists(ZIP_PATH):
    print("Extracting database from zip file...")
    with zipfile.ZipFile(ZIP_PATH, 'r') as zip_ref:
        zip_ref.extractall(".")
    print("Database extracted successfully!")

# Initialize
print(f"Loading database from: {DB_PATH}")
client = chromadb.PersistentClient(path=DB_PATH)
collection = client.get_collection("medqa")
print(f"Collection loaded with {collection.count()} items")
print(f"Loading MedCPT model...")
model = SentenceTransformer('ncbi/MedCPT-Query-Encoder')
print("Initialization complete!")

# FastAPI app
app = FastAPI(title="MedQA Search API")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

class SearchRequest(BaseModel):
    query: str
    num_results: int = 3

class SearchResponse(BaseModel):
    results: list