from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from clinical_ner import ClinicalNER
import uvicorn

# Initialize FastAPI app
app = FastAPI(
    title="Clinical NER API",
    description="Named Entity Recognition API using Bio_ClinicalBERT",
    version="1.0.0"
)

# Initialize the NER model (singleton pattern)
ner_model = None

@app.on_event("startup")
async def startup_event():
    """Load the NER model on startup"""
    global ner_model
    ner_model = ClinicalNER()
    print("NER model loaded successfully!")

# Request model
class TextRequest(BaseModel):
    text: str
    
    class Config:
        json_schema_extra = {
            "example": {
                "text": "Patient presents with hypertension and diabetes. Prescribed metformin 500mg."
            }
        }

# Response models
class Entity(BaseModel):
    entity_group: str
    score: float
    word: str
    start: int
    end: int

class BasicNERResponse(BaseModel):
    entities: list[Entity]
    count: int

class PrologNERResponse(BaseModel):
    prolog_facts: str
    count: int

@app.get("/")
async def root():
    """Root endpoint with API information"""
    return {
        "message": "Clinical NER API",
        "endpoints": {
            "/ner/basic": "POST - Get basic NER annotations",
            "/ner/prolog": "POST - Get Prolog facts",
            "/docs": "GET - Interactive API documentation"
        }
    }

@app.post("/ner/basic", response_model=BasicNERResponse)
async def ner_basic(request: TextRequest):
    """
    Perform basic NER on the input text.
    
    Returns a list of detected entities with their types, positions, and confidence scores.
    """
    try:
        if not request.text.strip():
            raise HTTPException(status_code=400, detail="Text cannot be empty")
        
        entities = ner_model.basic_ner(request.text)
        
        return {
            "entities": entities,
            "count": len(entities)
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error processing text: {str(e)}")

@app.post("/ner/prolog", response_model=PrologNERResponse)
async def ner_prolog(request: TextRequest):
    """
    Perform NER and return results as Prolog facts.
    
    Returns Prolog facts in the format: entity(Id, Type, Word, Start, End, Score).
    """
    try:
        if not request.text.strip():
            raise HTTPException(status_code=400, detail="Text cannot be empty")
        
        prolog_facts = ner_model.prolog_ner(request.text)
        
        # Count the number of facts (lines)
        count = len(prolog_facts.split('\n')) if prolog_facts else 0
        
        return {
            "prolog_facts": prolog_facts,
            "count": count
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error processing text: {str(e)}")

@app.get("/health")
async def health_check():
    """Health check endpoint"""
    return {
        "status": "healthy",
        "model_loaded": ner_model is not None
    }

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)