app/api.py

"""
FastAPI REST API layer wrapping the existing medical chatbot logic.
Converts the Gradio interface to a REST API for React frontend integration.
"""

import hashlib
import logging
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from typing import Optional, Dict, Any
import json

from app.main import (
    chat_fn,
    _load_state,
    _save_state,
    _parse_response,
    _get_question_hint,
)
from app.services.llm_extractor import extract_features_from_text
from app.services.feature_builder import count_collected_features, is_ready_for_prediction, prepare_feature_vector
from app.services.predictor import get_predictor
from app.memory import initialize_state, update_state, get_missing_features
from app.config import DEFAULT_MODEL_FEATURES, MIN_FEATURES_FOR_PREDICTION, CLASS_NAMES
from app.utils.helpers import generate_question, prioritize_features
from app.services.session_manager import get_session_manager

logger = logging.getLogger(__name__)

# Initialize FastAPI
app = FastAPI(title="Medical Diagnosis AI", description="REST API for medical health prediction")

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

# Session manager
session_manager = get_session_manager()


# ===== Request/Response Models =====

class ChatRequest(BaseModel):
    """Chat message request"""
    session_id: Optional[str] = None
    message: str
    history: list = []


class ChatResponse(BaseModel):
    """Chat response with features and state"""
    session_id: str
    message: str
    features: Dict[str, Any]
    collected_count: int
    total_features: int = 16
    is_complete: bool
    prediction: Optional[Dict[str, Any]] = None
    hint: str = ""


class ResetRequest(BaseModel):
    """Reset session request"""
    session_id: str


class SessionStateResponse(BaseModel):
    """Session state response"""
    session_id: str
    features: Dict[str, Any]
    collected_count: int
    total_features: int = 16


# ===== Helper Functions =====

def _build_acknowledgment(extracted: dict) -> str:
    """Build acknowledgment message from extracted features"""
    extracted_items = []
    for feature, value in extracted.items():
        if value is not None and feature in DEFAULT_MODEL_FEATURES:
            extracted_items.append(f"{feature}: {value}")

    if extracted_items:
        return f"✓ Got your {', '.join(extracted_items[:2])}"
    return ""


# ===== API Endpoints =====

@app.get("/health")
def health_check():
    """Health check endpoint"""
    return {"status": "ok", "service": "Medical Diagnosis AI"}


@app.post("/api/chat", response_model=ChatResponse)
def chat_endpoint(req: ChatRequest):
    """
    Send a message and get AI response with updated features.

    Handles:
    - Session ID generation if not provided
    - Feature extraction from user message
    - State persistence
    - Prediction when all 16 features collected
    """
    try:
        # Generate or use session ID
        if req.session_id:
            session_id = req.session_id
        else:
            # Generate from first message hash
            session_id = "sess_" + hashlib.md5(req.message.encode()).hexdigest()[:8]
            logger.info(f"🔐 Created new session: {session_id}")

        # Load persisted state
        state = _load_state(session_id)

        # Extract features from user message
        extracted = extract_features_from_text(req.message)

        # Update state with extracted features
        state = update_state(state, extracted)

        # Count collected features
        collected = count_collected_features(state)
        missing = get_missing_features(state)

        # Check if ready for prediction
        if is_ready_for_prediction(state, MIN_FEATURES_FOR_PREDICTION):
            # All 16 features collected - make prediction
            feature_vector = prepare_feature_vector(state)
            predictor = get_predictor()
            pred_result = predictor.predict(feature_vector)

            pred_data = {
                "prediction_class": int(pred_result.prediction),
                "prediction_name": CLASS_NAMES[int(pred_result.prediction)],
                "confidence": float(pred_result.probability),
                "risk_level": pred_result.risk_level,
                "explanation": pred_result.explanation,
                "features": state
            }

            # Simple completion message - full details now in DiagnosisCard component
            response_msg = "✅ Assessment Complete! Your diagnosis is ready below."

            _save_state(session_id, state)

            return ChatResponse(
                session_id=session_id,
                message=response_msg,
                features=state,
                collected_count=collected,
                is_complete=True,
                prediction=pred_data,
                hint=""
            )

        # Not complete yet - ask for next missing feature
        prioritized_missing = prioritize_features(missing)
        next_question = generate_question(prioritized_missing[:1])
        ack = _build_acknowledgment(extracted)
        remaining = 16 - collected

        response_msg = f"""{ack}

{next_question}

**{remaining} more pieces of information needed.**""" if ack else f"""{next_question}

**{remaining} more pieces of information needed.**"""

        # Get hint for the next question
        hint = _get_question_hint(next_question)

        _save_state(session_id, state)

        return ChatResponse(
            session_id=session_id,
            message=response_msg,
            features=state,
            collected_count=collected,
            is_complete=False,
            hint=hint
        )

    except Exception as e:
        logger.error(f"❌ Error in chat endpoint: {e}", exc_info=True)
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/api/reset")
def reset_endpoint(req: ResetRequest):
    """Reset a session - clear all features and start fresh"""
    try:
        session_manager = get_session_manager()
        success = session_manager.reset_session(req.session_id)

        if success:
            logger.info(f"✅ Reset session {req.session_id}")
            return {
                "success": True,
                "message": "Session reset successfully",
                "session_id": req.session_id
            }
        else:
            raise HTTPException(status_code=404, detail="Session not found")

    except Exception as e:
        logger.error(f"❌ Error resetting session: {e}")
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/api/session/{session_id}", response_model=SessionStateResponse)
def get_session_endpoint(session_id: str):
    """Get current session state"""
    try:
        state = _load_state(session_id)
        collected = count_collected_features(state)

        return SessionStateResponse(
            session_id=session_id,
            features=state,
            collected_count=collected
        )

    except Exception as e:
        logger.error(f"❌ Error getting session: {e}")
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/api/features")
def get_features_list():
    """Get list of all 16 features with their metadata"""
    from app.config import FEATURE_RANGES

    features_info = {}
    for feature in DEFAULT_MODEL_FEATURES:
        if feature in FEATURE_RANGES:
            min_val, max_val, _ = FEATURE_RANGES[feature]
            features_info[feature] = {
                "min": min_val,
                "max": max_val,
                "type": "numeric" if feature not in ["Smoking", "Alcohol", "Family History"] else "binary"
            }
        else:
            features_info[feature] = {"min": None, "max": None, "type": "unknown"}

    return {
        "total": len(DEFAULT_MODEL_FEATURES),
        "features": DEFAULT_MODEL_FEATURES,
        "metadata": features_info
    }


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