server.py

import os
import uuid
import json
import logging
from typing import List, Optional
from datetime import datetime, timedelta
from pathlib import Path

import jwt
from fastapi import FastAPI, File, UploadFile, HTTPException, Depends, Form, Body, Header, Query
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from fastapi.security import OAuth2PasswordBearer, OAuth2PasswordRequestForm
from pydantic import BaseModel, EmailStr, Field
from passlib.context import CryptContext
import ai_models
import label_bridge
from pymongo import MongoClient
from bson.objectid import ObjectId
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

# Setup logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Initialize FastAPI app
app = FastAPI(title="ZemedicAI API", description="API for ZemedicAI medical image analysis")

# Configure CORS
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # In production, replace with specific origins
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# MongoDB connection
MONGO_URL = os.environ.get("MONGO_URL")
DB_NAME = os.environ.get("DB_NAME", "zemedic_db")

try:
    client = MongoClient(MONGO_URL)
    # Test the connection
    client.admin.command('ping')
    db = client[DB_NAME]
    logger.info("Successfully connected to MongoDB")
except Exception as e:
    logger.error(f"Failed to connect to MongoDB: {str(e)}")
    # For demo purposes, we'll still initialize the app but won't use the database
    client = None
    db = None

# JWT Configuration
SECRET_KEY = os.environ.get("SECRET_KEY", "your-secret-key")  # Should be properly secured in production
ALGORITHM = "HS256"
ACCESS_TOKEN_EXPIRE_MINUTES = 60 * 24 * 7  # 1 week

# Password hashing
pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto")

# OAuth2 setup
oauth2_scheme = OAuth2PasswordBearer(tokenUrl="api/auth/token")

# File upload directory
UPLOAD_DIR = Path("./uploads")
UPLOAD_DIR.mkdir(parents=True, exist_ok=True)

# ----------------------------------------
# Models
# ----------------------------------------

class UserBase(BaseModel):
    email: EmailStr
    name: str

class UserCreate(UserBase):
    password: str

class User(UserBase):
    id: str
    created_at: datetime

    class Config:
        from_attributes = True

class UserInDB(User):
    hashed_password: str

class Token(BaseModel):
    access_token: str
    token_type: str
    user_id: str
    email: str
    name: str

class TokenData(BaseModel):
    user_id: Optional[str] = None

class LoginData(BaseModel):
    email: EmailStr
    password: str

class RegisterData(UserCreate):
    pass

class Prediction(BaseModel):
    label: str
    confidence: float
    description: Optional[str] = None

class AnalysisResult(BaseModel):
    id: str
    user_id: str
    type: str
    date: datetime
    image_url: str
    predictions: List[Prediction]
    recommendations: Optional[List[str]] = None

# ----------------------------------------
# Auth functions
# ----------------------------------------

def verify_password(plain_password, hashed_password):
    return pwd_context.verify(plain_password, hashed_password)

def get_password_hash(password):
    return pwd_context.hash(password)

def get_user(email: str):
    user_doc = db.users.find_one({"email": email})
    if user_doc:
        user_doc["id"] = str(user_doc["_id"])
        del user_doc["_id"]
        return UserInDB(**user_doc)
    return None

def authenticate_user(email: str, password: str):
    user = get_user(email)
    if not user:
        return False
    if not verify_password(password, user.hashed_password):
        return False
    return user

def create_access_token(data: dict, expires_delta: Optional[timedelta] = None):
    to_encode = data.copy()
    if expires_delta:
        expire = datetime.utcnow() + expires_delta
    else:
        expire = datetime.utcnow() + timedelta(minutes=15)
    to_encode.update({"exp": expire})
    encoded_jwt = jwt.encode(to_encode, SECRET_KEY, algorithm=ALGORITHM)
    return encoded_jwt

async def get_current_user(token: str = Depends(oauth2_scheme)):
    credentials_exception = HTTPException(
        status_code=401,
        detail="Could not validate credentials",
        headers={"WWW-Authenticate": "Bearer"},
    )
    
    # Demo mode for MongoDB unavailability
    if client is None or db is None:
        try:
            payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
            user_id: str = payload.get("sub")
            if user_id is None:
                raise credentials_exception
                
            # Create a demo user
            return User(
                id=user_id,
                email="demo@example.com",
                name="Demo User",
                created_at=datetime.utcnow()
            )
        except jwt.PyJWTError:
            raise credentials_exception
    
    # Normal flow with MongoDB
    try:
        payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
        user_id: str = payload.get("sub")
        if user_id is None:
            raise credentials_exception
        token_data = TokenData(user_id=user_id)
    except jwt.PyJWTError:
        raise credentials_exception
    
    user_doc = db.users.find_one({"_id": ObjectId(token_data.user_id)})
    if user_doc is None:
        raise credentials_exception
    
    user_doc["id"] = str(user_doc["_id"])
    del user_doc["_id"]
    return UserInDB(**user_doc)

# ----------------------------------------
# Image Analysis Functions  (delegated to ai_models — real models, no mock/random)
# ----------------------------------------

def analyze_xray_image(image_data: bytes, patient_context: Optional[dict] = None):
    return ai_models.analyze_xray_image(image_data, patient_context)

def analyze_skin_image(image_data: bytes, patient_context: Optional[dict] = None):
    return ai_models.analyze_skin_image(image_data, patient_context)

def analyze_ct_scan(image_data: bytes, patient_context: Optional[dict] = None):
    # The legacy "CT" endpoint is repurposed as ultrasound; route to the
    # breast-ultrasound (BUSI) model, which is our current ultrasound model.
    return ai_models.analyze_breast_ultrasound(image_data, patient_context)

def analyze_breast_image(image_data: bytes, patient_context: Optional[dict] = None):
    return ai_models.analyze_breast_ultrasound(image_data, patient_context)

def _analysis_response(analysis_id: str, type_: str, analysis_result: dict, image_url: str) -> dict:
    """Shape the API response, passing through the richer ai_models fields
    (model_status, severity, risk, and the dual-language Assistant)."""
    return {
        "id": analysis_id,
        "type": type_,
        "model_status": analysis_result.get("model_status"),
        "predictions": analysis_result.get("predictions", []),
        "top_finding": analysis_result.get("top_finding"),
        "severity": analysis_result.get("severity"),
        "risk": analysis_result.get("risk"),
        "recommendations": analysis_result.get("recommendations", []),
        "assistant": analysis_result.get("assistant"),
        "assistant_name": analysis_result.get("assistant_name", "Zafya AI"),
        "region_of_interest": analysis_result.get("region_of_interest"),
        "disclaimer": analysis_result.get("disclaimer"),
        "image_url": image_url,
    }


def _build_patient_context(
    age: Optional[int], sex: Optional[str],
    is_known_diabetic: Optional[bool], diabetes_duration_years: Optional[int],
) -> Optional[dict]:
    """Collects optional patient demographics submitted alongside an image.

    None of our vision models take these as literal model inputs (they're
    pure image classifiers) — but age/sex/diabetes history materially affect
    clinical interpretation (e.g. DR risk scales with diabetes duration;
    melanoma risk rises with age), so we thread them into the Zafya AI
    assistant's clinician-facing text and into Gemini's foot-ulcer prompt.
    They're also stored with the analysis record to enrich future
    ZemedicLabel training data with real demographic metadata.
    """
    if age is None and not sex and is_known_diabetic is None and diabetes_duration_years is None:
        return None
    return {
        "age": age,
        "sex": sex,
        "is_known_diabetic": is_known_diabetic,
        "diabetes_duration_years": diabetes_duration_years,
    }

def _flag_for_review_async(modality: str, image_data: bytes, filename: str, analysis_result: dict):
    """Fire-and-forget: push low-confidence predictions to ZemedicLabel for
    human review, without blocking or risking the diagnostic response."""
    import threading
    threading.Thread(
        target=label_bridge.maybe_flag_for_review,
        args=(modality, image_data, filename, analysis_result),
        daemon=True,
    ).start()

# ----------------------------------------
# API Endpoints
# ----------------------------------------

@app.get("/api/health")
def health_check():
    """Health check endpoint."""
    return {"status": "ok", "timestamp": datetime.now().isoformat()}

# Model readiness — warm models in a background thread so the server is
# immediately ready. Models are cached on first load (lru_cache), so requests
# that arrive before warmup finishes simply trigger a lazy load for that model.
def _do_warmup():
    logger.info("Warming up AI models (background)...")
    result = ai_models.warmup()
    logger.info("Model status: %s", result)

@app.on_event("startup")
def _warmup_models():
    import threading
    threading.Thread(target=_do_warmup, daemon=True).start()

@app.get("/api/models/status")
def models_status():
    """Reports which AI models are loaded and ready. Use before a live demo.
    Reads the live, incrementally-updated status (ready models show even while
    a slow one is still downloading)."""
    return {
        "models": ai_models.STATUS,
        "model_ids": {
            "xray":   ai_models.XRAY_WEIGHTS,
            "skin":   ai_models.SKIN_MODEL_ID,
            "retina": ai_models.RETINA_MODEL_ID,
            "foot":   "gemini-1.5-flash" if ai_models.GEMINI_API_KEY else "not-configured",
            "breast": ai_models.BREAST_MODEL_ID,
        },
        "model_source": ai_models.MODEL_SOURCE,
    }

@app.post("/api/models/reload")
def reload_model(modality: str, model_path: str):
    """Hot-swap a ViT model to a locally fine-tuned checkpoint produced by
    training/finetune.py — the Label->MVP accuracy improvement loop.
    Usage: POST /api/models/reload?modality=skin&model_path=/abs/path/to/checkpoint"""
    try:
        return ai_models.reload_model(modality, model_path)
    except (ValueError, FileNotFoundError) as e:
        raise HTTPException(status_code=400, detail=str(e))
    except Exception as e:  # noqa: BLE001
        raise HTTPException(status_code=500, detail=f"Reload failed: {e}")

@app.get("/api/active-learning/status")
def active_learning_status():
    """How many low-confidence predictions are queued in ZemedicLabel awaiting
    human review, per modality — the MVP -> Label side of the feedback loop."""
    return label_bridge.queue_status()

# Auth endpoints
@app.post("/api/auth/register", response_model=Token)
async def register(data: RegisterData):
    """Register a new user."""
    # Demo mode for MongoDB unavailability
    if client is None or db is None:
        # Generate a demo user ID and token for testing
        user_id = str(uuid.uuid4())
        access_token_expires = timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES)
        access_token = create_access_token(
            data={"sub": user_id}, expires_delta=access_token_expires
        )
        
        return {
            "access_token": access_token,
            "token_type": "bearer",
            "user_id": user_id,
            "email": data.email,
            "name": data.name
        }
    
    # Normal flow with MongoDB
    # Check if email already exists
    existing_user = db.users.find_one({"email": data.email})
    if existing_user:
        raise HTTPException(
            status_code=400,
            detail="Email already registered"
        )
    
    # Create new user document
    hashed_password = get_password_hash(data.password)
    user_id = str(ObjectId())
    
    user_data = {
        "_id": ObjectId(user_id),
        "email": data.email,
        "name": data.name,
        "hashed_password": hashed_password,
        "created_at": datetime.utcnow()
    }
    
    db.users.insert_one(user_data)
    
    # Create access token
    access_token_expires = timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES)
    access_token = create_access_token(
        data={"sub": user_id}, expires_delta=access_token_expires
    )
    
    return {
        "access_token": access_token,
        "token_type": "bearer",
        "user_id": user_id,
        "email": data.email,
        "name": data.name
    }

@app.post("/api/auth/login", response_model=Token)
async def login(data: LoginData):
    """Login and get access token."""
    # Demo mode for MongoDB unavailability
    if client is None or db is None:
        # Generate a demo user ID and token for testing
        user_id = str(uuid.uuid4())
        access_token_expires = timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES)
        access_token = create_access_token(
            data={"sub": user_id}, expires_delta=access_token_expires
        )
        
        return {
            "access_token": access_token,
            "token_type": "bearer",
            "user_id": user_id,
            "email": data.email,
            "name": "Demo User"
        }
    
    # Normal flow with MongoDB
    user = authenticate_user(data.email, data.password)
    if not user:
        raise HTTPException(
            status_code=401,
            detail="Incorrect email or password",
            headers={"WWW-Authenticate": "Bearer"},
        )
    
    access_token_expires = timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES)
    access_token = create_access_token(
        data={"sub": user.id}, expires_delta=access_token_expires
    )
    
    return {
        "access_token": access_token,
        "token_type": "bearer",
        "user_id": user.id,
        "email": user.email,
        "name": user.name
    }

# Analysis endpoints
@app.post("/api/analyze/xray")
async def analyze_xray(
    file: UploadFile = File(...),
    age: Optional[int] = Form(None),
    sex: Optional[str] = Form(None),
    is_known_diabetic: Optional[bool] = Form(None),
    diabetes_duration_years: Optional[int] = Form(None),
    current_user: User = Depends(get_current_user)
):
    """Analyze a chest X-ray image."""
    if not file.content_type.startswith("image/"):
        raise HTTPException(
            status_code=400,
            detail="File must be an image"
        )

    # Read the image file
    image_data = await file.read()
    patient_context = _build_patient_context(age, sex, is_known_diabetic, diabetes_duration_years)

    # Analyze the image
    analysis_result = analyze_xray_image(image_data, patient_context)
    
    # Generate a unique ID
    analysis_id = str(uuid.uuid4())
    
    # Save the image if possible
    try:
        file_extension = file.filename.split(".")[-1]
        unique_filename = f"{uuid.uuid4()}.{file_extension}"
        file_path = UPLOAD_DIR / unique_filename
        with open(file_path, "wb") as f:
            f.write(image_data)
        image_url = f"/uploads/{unique_filename}"
    except Exception as e:
        logger.error(f"Error saving image: {str(e)}")
        # Use a placeholder if saving fails
        image_url = "https://images.unsplash.com/photo-1584555684040-bad07f46a21f"
    
    # Store in MongoDB if available
    if client is not None and db is not None:
        try:
            analysis_doc = {
                "_id": ObjectId(analysis_id),
                "user_id": current_user.id,
                "type": "xray",
                "date": datetime.utcnow(),
                "image_url": image_url,
                "predictions": analysis_result["predictions"],
                "recommendations": analysis_result["recommendations"],
                "patient_context": patient_context,
                # Real region-of-interest if the model produced one; X-ray's
                # DenseNet backbone has no Grad-CAM wired yet, so this is None.
                "region_of_interest": analysis_result.get("region_of_interest"),
            }

            db.analyses.insert_one(analysis_doc)
        except Exception as e:
            logger.error(f"Error storing analysis in MongoDB: {str(e)}")

    _flag_for_review_async("xray", image_data, file.filename, analysis_result)
    # Return the result
    return _analysis_response(analysis_id, "xray", analysis_result, image_url)

@app.post("/api/analyze/skin")
async def analyze_skin_lesion(
    file: UploadFile = File(...),
    age: Optional[int] = Form(None),
    sex: Optional[str] = Form(None),
    is_known_diabetic: Optional[bool] = Form(None),
    diabetes_duration_years: Optional[int] = Form(None),
    current_user: User = Depends(get_current_user)
):
    """Analyze a skin lesion image."""
    if not file.content_type.startswith("image/"):
        raise HTTPException(
            status_code=400,
            detail="File must be an image"
        )

    # Read the image file
    image_data = await file.read()
    patient_context = _build_patient_context(age, sex, is_known_diabetic, diabetes_duration_years)

    # Analyze the image
    analysis_result = analyze_skin_image(image_data, patient_context)
    
    # Generate a unique ID
    analysis_id = str(uuid.uuid4())
    
    # Save the image if possible
    try:
        file_extension = file.filename.split(".")[-1]
        unique_filename = f"{uuid.uuid4()}.{file_extension}"
        file_path = UPLOAD_DIR / unique_filename
        with open(file_path, "wb") as f:
            f.write(image_data)
        image_url = f"/uploads/{unique_filename}"
    except Exception as e:
        logger.error(f"Error saving image: {str(e)}")
        # Use a placeholder if saving fails
        image_url = "https://images.unsplash.com/photo-1606501190025-f3ad6d3ea6ae"
    
    # Store in MongoDB if available
    if client is not None and db is not None:
        try:
            analysis_doc = {
                "_id": ObjectId(analysis_id),
                "user_id": current_user.id,
                "type": "skin",
                "date": datetime.utcnow(),
                "image_url": image_url,
                "predictions": analysis_result["predictions"],
                "recommendations": analysis_result["recommendations"],
                "patient_context": patient_context,
                "region_of_interest": analysis_result.get("region_of_interest"),
            }

            db.analyses.insert_one(analysis_doc)
        except Exception as e:
            logger.error(f"Error storing analysis in MongoDB: {str(e)}")

    _flag_for_review_async("skin", image_data, file.filename, analysis_result)
    # Return the result
    return _analysis_response(analysis_id, "skin", analysis_result, image_url)

@app.post("/api/analyze/ct-scan")
async def analyze_ct_scan_image(
    file: UploadFile = File(...),
    age: Optional[int] = Form(None),
    sex: Optional[str] = Form(None),
    is_known_diabetic: Optional[bool] = Form(None),
    diabetes_duration_years: Optional[int] = Form(None),
    current_user: User = Depends(get_current_user)
):
    """Analyze a Ultrasound image."""
    if not file.content_type.startswith("image/"):
        raise HTTPException(
            status_code=400,
            detail="File must be an image"
        )

    # Read the image file
    image_data = await file.read()
    patient_context = _build_patient_context(age, sex, is_known_diabetic, diabetes_duration_years)

    # Analyze the image
    analysis_result = analyze_ct_scan(image_data, patient_context)
    
    # Generate a unique ID
    analysis_id = str(uuid.uuid4())
    
    # Save the image if possible
    try:
        file_extension = file.filename.split(".")[-1]
        unique_filename = f"{uuid.uuid4()}.{file_extension}"
        file_path = UPLOAD_DIR / unique_filename
        with open(file_path, "wb") as f:
            f.write(image_data)
        image_url = f"/uploads/{unique_filename}"
    except Exception as e:
        logger.error(f"Error saving image: {str(e)}")
        # Use a placeholder if saving fails
        image_url = "https://images.unsplash.com/photo-1631563019676-dade0dbdb8fc"
    
    # Store in MongoDB if available
    if client is not None and db is not None:
        try:
            analysis_doc = {
                "_id": ObjectId(analysis_id),
                "user_id": current_user.id,
                "type": "ct-scan",
                "date": datetime.utcnow(),
                "image_url": image_url,
                "predictions": analysis_result["predictions"],
                "recommendations": analysis_result["recommendations"],
                "patient_context": patient_context,
                "region_of_interest": analysis_result.get("region_of_interest"),
            }

            db.analyses.insert_one(analysis_doc)
        except Exception as e:
            logger.error(f"Error storing analysis in MongoDB: {str(e)}")

    _flag_for_review_async("breast", image_data, file.filename, analysis_result)
    # Return the result
    return _analysis_response(analysis_id, "ultrasound", analysis_result, image_url)

@app.post("/api/analyze/breast")
async def analyze_breast(
    file: UploadFile = File(...),
    age: Optional[int] = Form(None),
    sex: Optional[str] = Form(None),
    is_known_diabetic: Optional[bool] = Form(None),
    diabetes_duration_years: Optional[int] = Form(None),
    current_user: User = Depends(get_current_user)
):
    """Analyze a breast ultrasound image (BUSI: benign / malignant / normal)."""
    if not file.content_type.startswith("image/"):
        raise HTTPException(status_code=400, detail="File must be an image")

    image_data = await file.read()
    patient_context = _build_patient_context(age, sex, is_known_diabetic, diabetes_duration_years)
    analysis_result = analyze_breast_image(image_data, patient_context)
    analysis_id = str(uuid.uuid4())

    try:
        file_extension = file.filename.split(".")[-1]
        unique_filename = f"{uuid.uuid4()}.{file_extension}"
        file_path = UPLOAD_DIR / unique_filename
        with open(file_path, "wb") as f:
            f.write(image_data)
        image_url = f"/uploads/{unique_filename}"
    except Exception as e:
        logger.error(f"Error saving image: {str(e)}")
        image_url = "https://images.unsplash.com/photo-1631563019676-dade0dbdb8fc"

    if client is not None and db is not None:
        try:
            db.analyses.insert_one({
                "_id": ObjectId(analysis_id),
                "user_id": current_user.id,
                "type": "breast",
                "date": datetime.utcnow(),
                "image_url": image_url,
                "predictions": analysis_result.get("predictions", []),
                "recommendations": analysis_result.get("recommendations", []),
                "patient_context": patient_context,
                "region_of_interest": analysis_result.get("region_of_interest"),
            })
        except Exception as e:
            logger.error(f"Error storing analysis in MongoDB: {str(e)}")

    _flag_for_review_async("breast", image_data, file.filename, analysis_result)
    return _analysis_response(analysis_id, "breast", analysis_result, image_url)


@app.post("/api/analyze/retina")
async def analyze_retina(
    file: UploadFile = File(...),
    age: Optional[int] = Form(None),
    sex: Optional[str] = Form(None),
    is_known_diabetic: Optional[bool] = Form(None),
    diabetes_duration_years: Optional[int] = Form(None),
    current_user: User = Depends(get_current_user),
):
    """Diabetic retinopathy grading from fundus photography.

    Uses a ViT fine-tuned on APTOS 2019 / EyePACS (5-class DR grading:
    No DR → Proliferative DR). Returns structured predictions + dual-language
    clinician / patient assistant output. Age/diabetes-duration are clinically
    relevant here — DR risk and severity correlate strongly with how long a
    patient has had diabetes — so they're captured and passed to the assistant."""
    if not file.content_type.startswith("image/"):
        raise HTTPException(status_code=400, detail="File must be an image")

    image_data = await file.read()
    patient_context = _build_patient_context(age, sex, is_known_diabetic, diabetes_duration_years)
    analysis_result = ai_models.analyze_retina_image(image_data, patient_context)
    analysis_id = str(uuid.uuid4())

    try:
        ext = file.filename.rsplit(".", 1)[-1]
        unique_filename = f"{uuid.uuid4()}.{ext}"
        with open(UPLOAD_DIR / unique_filename, "wb") as f:
            f.write(image_data)
        image_url = f"/uploads/{unique_filename}"
    except Exception as e:
        logger.error("Error saving retina image: %s", e)
        image_url = ""

    if client is not None and db is not None:
        try:
            db.analyses.insert_one({
                "_id": ObjectId(analysis_id),
                "user_id": current_user.id,
                "type": "retina",
                "date": datetime.utcnow(),
                "image_url": image_url,
                "predictions": analysis_result.get("predictions", []),
                "recommendations": analysis_result.get("recommendations", []),
                "patient_context": patient_context,
                "region_of_interest": analysis_result.get("region_of_interest"),
            })
        except Exception as e:
            logger.error("MongoDB insert failed (retina): %s", e)

    _flag_for_review_async("retina", image_data, file.filename, analysis_result)
    return _analysis_response(analysis_id, "retina", analysis_result, image_url)


@app.post("/api/analyze/foot")
async def analyze_foot(
    file: UploadFile = File(...),
    age: Optional[int] = Form(None),
    sex: Optional[str] = Form(None),
    is_known_diabetic: Optional[bool] = Form(None),
    diabetes_duration_years: Optional[int] = Form(None),
    current_user: User = Depends(get_current_user),
):
    """Diabetic foot ulcer classification from clinical photography.

    Uses Gemini 1.5 Flash vision for wound classification (normal | DFU |
    infected | neuropathic | ischemic | callus | healed | other).
    Requires GEMINI_API_KEY in the backend .env. Diabetes duration is
    clinically relevant to foot ulcer risk, so it's passed into the Gemini
    prompt as patient context alongside the image."""
    if not file.content_type.startswith("image/"):
        raise HTTPException(status_code=400, detail="File must be an image")

    image_data = await file.read()
    patient_context = _build_patient_context(age, sex, is_known_diabetic, diabetes_duration_years)
    analysis_result = ai_models.analyze_foot_image(image_data, patient_context)
    analysis_id = str(uuid.uuid4())

    try:
        ext = file.filename.rsplit(".", 1)[-1]
        unique_filename = f"{uuid.uuid4()}.{ext}"
        with open(UPLOAD_DIR / unique_filename, "wb") as f:
            f.write(image_data)
        image_url = f"/uploads/{unique_filename}"
    except Exception as e:
        logger.error("Error saving foot image: %s", e)
        image_url = ""

    if client is not None and db is not None:
        try:
            db.analyses.insert_one({
                "_id": ObjectId(analysis_id),
                "user_id": current_user.id,
                "type": "foot",
                "date": datetime.utcnow(),
                "image_url": image_url,
                "predictions": analysis_result.get("predictions", []),
                "recommendations": analysis_result.get("recommendations", []),
                "patient_context": patient_context,
                "region_of_interest": analysis_result.get("region_of_interest"),
            })
        except Exception as e:
            logger.error("MongoDB insert failed (foot): %s", e)

    _flag_for_review_async("foot", image_data, file.filename, analysis_result)
    return _analysis_response(analysis_id, "foot", analysis_result, image_url)


# ----------------------------------------
# Zafya AI — conversational follow-up assistant
# ----------------------------------------
import zafya_chat  # noqa: E402


class ZafyaChatTurn(BaseModel):
    role: str  # "user" | "assistant"
    content: str


class ZafyaChatRequest(BaseModel):
    message: str
    history: List[ZafyaChatTurn] = []
    audience: str = "patient"  # "doctor" | "patient"
    context: Optional[dict] = None  # {modality, top_finding, severity, risk, age, sex}


@app.post("/api/zafya/chat")
async def zafya_chat_endpoint(
    payload: ZafyaChatRequest,
    current_user: User = Depends(get_current_user),
):
    """Follow-up conversational Q&A on a screening result, tailored to the
    asker (doctor gets clinical register, patient gets plain-language register).
    Powered by Gemini; distinct from ZafyaLM (the offline on-device model)."""
    history = [turn.dict() for turn in payload.history]
    result = zafya_chat.chat(payload.message, history, payload.audience, payload.context)
    return result


@app.get("/api/user/history")
async def get_user_history(
    current_user: User = Depends(get_current_user)
):
    """Get user's analysis history."""
    # Demo mode for MongoDB unavailability
    if client is None or db is None:
        # Return demo data
        demo_history = [
            {
                "id": "demo1",
                "user_id": current_user.id,
                "type": "xray",
                "date": datetime.utcnow() - timedelta(days=2),
                "findings": ["Pneumonia", "Pleural Effusion"],
                "image_url": "https://images.unsplash.com/photo-1584555684040-bad07f46a21f",
                "predictions": [
                    {"label": "Pneumonia", "confidence": 0.82},
                    {"label": "Pleural Effusion", "confidence": 0.67}
                ]
            },
            {
                "id": "demo2",
                "user_id": current_user.id,
                "type": "skin",
                "date": datetime.utcnow() - timedelta(days=5),
                "findings": ["Melanoma"],
                "image_url": "https://images.unsplash.com/photo-1606501190025-f3ad6d3ea6ae",
                "predictions": [
                    {"label": "Melanoma", "confidence": 0.75}
                ]
            },
            {
                "id": "demo3",
                "user_id": current_user.id,
                "type": "ct-scan",
                "date": datetime.utcnow() - timedelta(days=7),
                "findings": ["Brain Tumor"],
                "image_url": "https://images.unsplash.com/photo-1631563019676-dade0dbdb8fc",
                "predictions": [
                    {"label": "Brain Tumor", "confidence": 0.68}
                ]
            }
        ]
        return demo_history
    
    # Normal flow with MongoDB
    try:
        analyses = db.analyses.find({"user_id": current_user.id}).sort("date", -1)
        
        result = []
        for analysis in analyses:
            analysis["id"] = str(analysis["_id"])
            del analysis["_id"]
            result.append(analysis)
        
        return result
    except Exception as e:
        logger.error(f"Error retrieving history: {str(e)}")
        return []

@app.get("/api/analysis/{analysis_id}")
async def get_analysis_by_id(
    analysis_id: str,
    current_user: User = Depends(get_current_user)
):
    """Get a specific analysis by ID."""
    # Demo mode for MongoDB unavailability
    if client is None or db is None or analysis_id.startswith("demo"):
        # Return demo data
        if analysis_id == "demo1":
            return {
                "id": "demo1",
                "user_id": current_user.id,
                "type": "xray",
                "date": datetime.utcnow() - timedelta(days=2),
                "image_url": "https://images.unsplash.com/photo-1584555684040-bad07f46a21f",
                "predictions": [
                    {"label": "Pneumonia", "confidence": 0.82},
                    {"label": "Pleural Effusion", "confidence": 0.67}
                ],
                "recommendations": [
                    "Consult with a healthcare professional for proper diagnosis",
                    "Consider follow-up imaging to monitor any changes",
                    "Complete the full course of prescribed antibiotics",
                    "Rest and hydrate properly"
                ]
            }
        elif analysis_id == "demo2":
            return {
                "id": "demo2",
                "user_id": current_user.id,
                "type": "skin",
                "date": datetime.utcnow() - timedelta(days=5),
                "image_url": "https://images.unsplash.com/photo-1606501190025-f3ad6d3ea6ae",
                "predictions": [
                    {"label": "Melanoma", "confidence": 0.75},
                    {"label": "Seborrheic Keratosis", "confidence": 0.15}
                ],
                "recommendations": [
                    "Schedule a follow-up with a dermatologist",
                    "Protect your skin from sun exposure",
                    "Monitor any changes in size, shape, or color of the lesion",
                    "Apply prescribed topical treatments as directed"
                ]
            }
        elif analysis_id == "demo3":
            return {
                "id": "demo3",
                "user_id": current_user.id,
                "type": "ct-scan",
                "date": datetime.utcnow() - timedelta(days=7),
                "image_url": "https://images.unsplash.com/photo-1631563019676-dade0dbdb8fc",
                "predictions": [
                    {"label": "Brain Tumor", "confidence": 0.68},
                    {"label": "Normal Findings", "confidence": 0.22}
                ],
                "recommendations": [
                    "Immediate neurosurgery consultation",
                    "MRI with contrast for further characterization",
                    "Discuss biopsy options if appropriate",
                    "Consider second opinion from neuro-oncologist"
                ]
            }
        else:
            raise HTTPException(
                status_code=404,
                detail="Analysis not found"
            )
    
    # Normal flow with MongoDB
    try:
        analysis = db.analyses.find_one({"_id": ObjectId(analysis_id)})
        
        if not analysis:
            raise HTTPException(
                status_code=404,
                detail="Analysis not found"
            )
        
        if analysis["user_id"] != current_user.id:
            raise HTTPException(
                status_code=403,
                detail="Not authorized to access this analysis"
            )
        
        analysis["id"] = str(analysis["_id"])
        del analysis["_id"]
        
        return analysis
    except Exception as e:
        logger.error(f"Error retrieving analysis: {str(e)}")
        raise HTTPException(
            status_code=500,
            detail=f"Error retrieving analysis: {str(e)}"
        )

@app.get("/api/user/profile")
async def get_user_profile(
    current_user: User = Depends(get_current_user)
):
    """Get user's profile information."""
    return {
        "id": current_user.id,
        "email": current_user.email,
        "name": current_user.name,
        "created_at": current_user.created_at
    }

@app.put("/api/user/profile")
async def update_user_profile(
    name: str = Form(None),
    current_user: UserInDB = Depends(get_current_user)
):
    """Update user's profile information."""
    updates = {}
    
    if name:
        updates["name"] = name
    
    if updates:
        db.users.update_one(
            {"_id": ObjectId(current_user.id)},
            {"$set": updates}
        )
    
    updated_user = db.users.find_one({"_id": ObjectId(current_user.id)})
    updated_user["id"] = str(updated_user["_id"])
    del updated_user["_id"]
    del updated_user["hashed_password"]
    
    return updated_user

if __name__ == "__main__":
    import uvicorn
    
    # Run the FastAPI app
    uvicorn.run("server:app", host="0.0.0.0", port=8001, reload=True)