#!/usr/bin/env python3
"""
Medical AI Assistant - FastAPI Only Version
Simplified endpoints for backend integration with Swagger UI

This file is Hugging Face Spaces compatible: the FastAPI app is exposed as 'app' at the module level.
"""

from fastapi import FastAPI, HTTPException, File, UploadFile, BackgroundTasks
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from fastapi.openapi.docs import get_swagger_ui_html
from fastapi.openapi.utils import get_openapi
from pydantic import BaseModel, Field
from typing import List, Optional, Dict, Any, Union
import logging
import uuid
import os
import json
import asyncio
from contextlib import asynccontextmanager
import time

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

# Initialize models globally
pipeline = None
whisper_model = None

async def load_models():
    """Load ML models asynchronously"""
    global pipeline, whisper_model
    try:
        logger.info("Loading Medical AI models...")
        
        from medical_ai import CompetitionMedicalAIPipeline
        pipeline = CompetitionMedicalAIPipeline()
        logger.info("✅ Medical pipeline loaded successfully")
        
        try:
            from faster_whisper import WhisperModel
            model_cache = os.getenv('HF_HOME', '/tmp/models')
            whisper_model = WhisperModel(
                "medium",
                device="cpu",
                compute_type="int8",
                download_root=model_cache
            )
            logger.info("✅ Whisper model loaded successfully")
        except Exception as e:
            logger.warning(f"⚠️ Could not load Whisper model: {str(e)}")
            whisper_model = None
            
        logger.info("🚀 All models loaded successfully")
        
    except Exception as e:
        logger.error(f"❌ Error loading models: {str(e)}", exc_info=True)
        raise

@asynccontextmanager
async def lifespan(app: FastAPI):
    """Application lifespan management (robust for Hugging Face Spaces)"""
    try:
        await load_models()
        logger.info("✅ Models loaded in lifespan.")
    except Exception as e:
        logger.error(f"❌ Error during startup: {str(e)}", exc_info=True)
        # Do not raise, just log. App will start but endpoints will return 503 if models are missing.
    yield
    logger.info("🔄 Shutting down...")

# Custom OpenAPI schema
def custom_openapi():
    if app.openapi_schema:
        return app.openapi_schema
    
    openapi_schema = get_openapi(
        title="🩺 Medical AI Assistant API",
        version="2.0.0",
        description="""
        ## 🎯 Advanced Medical AI Assistant
        
        **Multilingual medical consultation API** supporting:
        - 🌍 French, English, and local African languages
        - 🎤 Audio processing with speech-to-text
        - 🧠 Advanced medical knowledge retrieval
        - ⚡ Real-time medical consultations
        
        ### 🔧 Main Endpoints:
        - **POST /medical/ask** - Text-based medical consultation
        - **POST /medical/audio** - Audio-based medical consultation
        - **GET /health** - System health check
        - **POST /feedback** - Submit user feedback
        
        ### 🔒 Important Medical Disclaimer:
        This API provides educational medical information only. Always consult qualified healthcare professionals for medical advice.
        """,
        routes=app.routes,
        contact={
            "name": "Medical AI Support",
            "email": "support@medicalai.com"
        },
        license_info={
            "name": "MIT License",
            "url": "https://opensource.org/licenses/MIT"
        }
    )
    
    # Add custom tags
    openapi_schema["tags"] = [
        {
            "name": "medical",
            "description": "Medical consultation endpoints"
        },
        {
            "name": "audio",
            "description": "Audio processing endpoints"
        },
        {
            "name": "system",
            "description": "System monitoring and health"
        },
        {
            "name": "feedback",
            "description": "User feedback and analytics"
        }
    ]
    
    app.openapi_schema = openapi_schema
    return app.openapi_schema

# Initialize FastAPI app
app = FastAPI(
    title="🩺 Medical AI Assistant",
    description="Advanced multilingual medical consultation API",
    version="2.0.0",
    lifespan=lifespan,
    docs_url="/docs",
    redoc_url="/redoc",
    openapi_url="/openapi.json"
)

# Set custom OpenAPI
app.openapi = custom_openapi

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

# ============================================================================
# PYDANTIC MODELS FOR REQUEST/RESPONSE VALIDATION
# ============================================================================

class MedicalQuestion(BaseModel):
    """Medical question request model"""
    question: str = Field(..., description="The medical question", min_length=3, max_length=1000)
    language: str = Field("auto", description="Preferred language (auto, en, fr)", pattern="^(auto|en|fr)$")
    conversation_id: Optional[str] = Field(None, description="Optional conversation ID for context")
    
    class Config:
        schema_extra = {
            "example": {
                "question": "What are the symptoms of malaria and how is it treated?",
                "language": "en",
                "conversation_id": "conv_123"
            }
        }

class MedicalResponse(BaseModel):
    """Medical response model"""
    success: bool = Field(..., description="Whether the request was successful")
    response: str = Field(..., description="The medical response")
    detected_language: str = Field(..., description="Detected or used language")
    conversation_id: str = Field(..., description="Conversation identifier")
    context_used: List[str] = Field(default_factory=list, description="Medical contexts used")
    processing_time: float = Field(..., description="Response time in seconds")
    confidence: str = Field(..., description="Response confidence level")
    
    class Config:
        schema_extra = {
            "example": {
                "success": True,
                "response": "Malaria symptoms include high fever, chills, headache...",
                "detected_language": "en",
                "conversation_id": "conv_123",
                "context_used": ["Malaria treatment protocols", "Symptom guidelines"],
                "processing_time": 2.5,
                "confidence": "high"
            }
        }

class AudioResponse(BaseModel):
    """Audio processing response model"""
    success: bool = Field(..., description="Whether the request was successful")
    transcription: str = Field(..., description="Transcribed text from audio")
    response: str = Field(..., description="The medical response")
    detected_language: str = Field(..., description="Detected audio language")
    conversation_id: str = Field(..., description="Conversation identifier")
    context_used: List[str] = Field(default_factory=list, description="Medical contexts used")
    processing_time: float = Field(..., description="Response time in seconds")
    audio_duration: Optional[float] = Field(None, description="Audio duration in seconds")
    
    class Config:
        schema_extra = {
            "example": {
                "success": True,
                "transcription": "What are the symptoms of malaria?",
                "response": "Malaria symptoms include high fever, chills...",
                "detected_language": "en",
                "conversation_id": "conv_456",
                "context_used": ["Malaria diagnosis"],
                "processing_time": 3.2,
                "audio_duration": 4.5
            }
        }

class FeedbackRequest(BaseModel):
    """Feedback request model"""
    conversation_id: str = Field(..., description="Conversation ID")
    rating: int = Field(..., description="Rating from 1-5", ge=1, le=5)
    feedback: Optional[str] = Field(None, description="Optional text feedback", max_length=500)
    
    class Config:
        schema_extra = {
            "example": {
                "conversation_id": "conv_123",
                "rating": 5,
                "feedback": "Very helpful and accurate medical information"
            }
        }

class HealthStatus(BaseModel):
    """System health status model"""
    status: str = Field(..., description="Overall system status")
    models_loaded: bool = Field(..., description="Whether ML models are loaded")
    audio_available: bool = Field(..., description="Whether audio processing is available")
    uptime: float = Field(..., description="System uptime in seconds")
    version: str = Field(..., description="API version")
    
    class Config:
        schema_extra = {
            "example": {
                "status": "healthy",
                "models_loaded": True,
                "audio_available": True,
                "uptime": 3600.0,
                "version": "2.0.0"
            }
        }

class ErrorResponse(BaseModel):
    """Error response model"""
    success: bool = Field(False, description="Always false for errors")
    error: str = Field(..., description="Error message")
    error_code: str = Field(..., description="Error code")
    conversation_id: Optional[str] = Field(None, description="Conversation ID if available")

# ============================================================================
# UTILITY FUNCTIONS
# ============================================================================

def generate_conversation_id() -> str:
    """Generate a unique conversation ID"""
    return f"conv_{uuid.uuid4().hex[:8]}"

def validate_models():
    """Check if models are loaded"""
    if pipeline is None:
        raise HTTPException(
            status_code=503,
            detail="Medical AI models are not loaded yet. Please try again in a moment."
        )

# ============================================================================
# API ENDPOINTS
# ============================================================================

@app.get("/", tags=["system"])
async def root():
    """Root endpoint with API information"""
    return {
        "message": "🩺 Medical AI Assistant API",
        "version": "2.0.0",
        "status": "running",
        "docs": "/docs",
        "redoc": "/redoc",
        "endpoints": {
            "medical_consultation": "/medical/ask",
            "audio_consultation": "/medical/audio",
            "health_check": "/health",
            "feedback": "/feedback"
        }
    }

@app.get("/health", response_model=HealthStatus, tags=["system"])
async def health_check():
    """
    ## System Health Check
    
    Returns the current status of the Medical AI system including:
    - Overall system health
    - Model loading status
    - Audio processing availability
    - System uptime
    """
    global pipeline, whisper_model
    
    # Calculate uptime (simplified)
    uptime = time.time() - getattr(health_check, 'start_time', time.time())
    if not hasattr(health_check, 'start_time'):
        health_check.start_time = time.time()
    
    return HealthStatus(
        status="healthy" if pipeline is not None else "loading",
        models_loaded=pipeline is not None,
        audio_available=whisper_model is not None,
        uptime=uptime,
        version="2.0.0"
    )

@app.post("/medical/ask", response_model=MedicalResponse, tags=["medical"])
async def medical_consultation(request: MedicalQuestion):
    """
    ## Text-based Medical Consultation
    
    Process a medical question and return expert medical guidance.
    
    **Features:**
    - 🌍 Multilingual support (auto-detect or specify language)
    - 🧠 AI-powered medical knowledge retrieval
    - ⚡ Fast response generation
    - 🔒 Medical disclaimers included
    
    **Supported Languages:** English (en), French (fr), Auto-detect (auto)
    """
    start_time = time.time()
    validate_models()
    
    conversation_id = request.conversation_id or generate_conversation_id()
    
    try:
        logger.info(f"🩺 Processing medical question: {request.question[:50]}...")
        
        # Process with medical AI pipeline
        result = pipeline.process(
            question=request.question,
            user_lang=request.language,
            conversation_history=[]
        )
        
        processing_time = time.time() - start_time
        
        return MedicalResponse(
            success=True,
            response=result["response"],
            detected_language=result["source_lang"],
            conversation_id=conversation_id,
            context_used=result.get("context_used", []),
            processing_time=round(processing_time, 2),
            confidence=result.get("confidence", "medium")
        )
        
    except Exception as e:
        logger.error(f"❌ Error in medical consultation: {str(e)}", exc_info=True)
        processing_time = time.time() - start_time
        
        raise HTTPException(
            status_code=500,
            detail={
                "success": False,
                "error": "Internal processing error occurred",
                "error_code": "MEDICAL_PROCESSING_ERROR",
                "conversation_id": conversation_id,
                "processing_time": round(processing_time, 2)
            }
        )

@app.post("/medical/audio", response_model=AudioResponse, tags=["audio", "medical"])
async def audio_medical_consultation(
    file: UploadFile = File(..., description="Audio file (WAV, MP3, M4A, etc.)")
):
    """
    ## Audio-based Medical Consultation
    
    Process an audio medical question and return expert medical guidance.
    
    **Features:**
    - 🎤 Speech-to-text conversion
    - 🌍 Language detection from audio
    - 🧠 Medical AI processing of transcribed text
    - 📝 Full transcription provided
    
    **Supported Audio Formats:** WAV, MP3, M4A, FLAC, OGG
    **Max File Size:** 25MB
    **Max Duration:** 5 minutes
    """
    start_time = time.time()
    validate_models()
    
    if whisper_model is None:
        raise HTTPException(
            status_code=503,
            detail="Audio processing is currently unavailable"
        )
    
    conversation_id = generate_conversation_id()
    
    try:
        logger.info(f"🎤 Processing audio file: {file.filename}")
        
        # Read audio file
        file_bytes = await file.read()
        
        # Process audio
        from audio_utils import preprocess_audio
        processed_audio = preprocess_audio(file_bytes)
        
        if len(processed_audio) == 0:
            raise HTTPException(
                status_code=400,
                detail="Could not process audio file. Please check the format and try again."
            )
        
        # Transcribe audio
        segments, info = whisper_model.transcribe(
            processed_audio,
            beam_size=5,
            language=None,
            task='transcribe',
            vad_filter=True
        )
        
        transcription = "".join([seg.text for seg in segments])
        detected_language = info.language
        
        if not transcription.strip():
            raise HTTPException(
                status_code=400,
                detail="Could not transcribe audio. Please ensure clear speech and try again."
            )
        
        logger.info(f"🔤 Transcription: {transcription[:100]}...")
        
        # Process transcribed text with medical AI
        result = pipeline.process(
            question=transcription,
            user_lang=detected_language,  
            conversation_history=[]
        )
        
        processing_time = time.time() - start_time
        
        return AudioResponse(
            success=True,
            transcription=transcription,
            response=result["response"],
            detected_language=detected_language,
            conversation_id=conversation_id,
            context_used=result.get("context_used", []),
            processing_time=round(processing_time, 2),
            audio_duration=len(processed_audio) / 16000  # Assuming 16kHz sample rate
        )
        
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"❌ Error in audio processing: {str(e)}", exc_info=True)
        processing_time = time.time() - start_time
        
        raise HTTPException(
            status_code=500,
            detail={
                "success": False,
                "error": "Audio processing error occurred",
                "error_code": "AUDIO_PROCESSING_ERROR",
                "conversation_id": conversation_id,
                "processing_time": round(processing_time, 2)
            }
        )

@app.post("/feedback", tags=["feedback"])
async def submit_feedback(request: FeedbackRequest):
    """
    ## Submit User Feedback
    
    Submit feedback about a medical consultation to help improve the service.
    
    **Rating Scale:**
    - 1: Very Poor
    - 2: Poor  
    - 3: Average
    - 4: Good
    - 5: Excellent
    """
    try:
        logger.info(f"📊 Feedback received - ID: {request.conversation_id}, Rating: {request.rating}")
        
        # Here you could store feedback in a database
        # For now, just log it
        feedback_data = {
            "conversation_id": request.conversation_id,
            "rating": request.rating,
            "feedback": request.feedback,
            "timestamp": time.time()
        }
        
        return {
            "success": True,
            "message": "Thank you for your feedback! This helps us improve our medical AI service.",
            "feedback_id": f"fb_{uuid.uuid4().hex[:8]}"
        }
        
    except Exception as e:
        logger.error(f"❌ Error processing feedback: {str(e)}")
        raise HTTPException(
            status_code=500,
            detail="Error processing feedback"
        )

@app.get("/medical/specialties", tags=["medical"])
async def get_medical_specialties():
    """
    ## Get Supported Medical Specialties
    
    Returns a list of medical specialties and conditions supported by the AI.
    """
    return {
        "specialties": [
            {
                "name": "Primary Care",
                "description": "General medical consultations and health guidance",
                "conditions": ["General symptoms", "Preventive care", "Health maintenance"]
            },
            {
                "name": "Infectious Diseases", 
                "description": "Infectious disease diagnosis and treatment",
                "conditions": ["Malaria", "Tuberculosis", "HIV/AIDS", "Respiratory infections"]
            },
            {
                "name": "Emergency Medicine",
                "description": "Emergency protocols and urgent care guidance",
                "conditions": ["Stroke recognition", "Cardiac emergencies", "Trauma assessment"]
            },
            {
                "name": "Chronic Disease Management",
                "description": "Management of chronic conditions",
                "conditions": ["Diabetes", "Hypertension", "Gastritis"]
            }
        ],
        "languages_supported": ["English", "French", "Auto-detect"],
        "disclaimer": "This AI provides educational information only. Always consult healthcare professionals for medical advice."
    }

# ============================================================================
# ERROR HANDLERS
# ============================================================================

@app.exception_handler(404)
async def not_found_handler(request, exc):
    return JSONResponse(
        status_code=404,
        content={
            "success": False,
            "error": "Endpoint not found",
            "error_code": "NOT_FOUND", 
            "available_endpoints": [
                "/docs - API Documentation",
                "/medical/ask - Text consultation",
                "/medical/audio - Audio consultation", 
                "/health - System status",
                "/feedback - Submit feedback"
            ]
        }
    )

@app.exception_handler(422)
async def validation_exception_handler(request, exc):
    return JSONResponse(
        status_code=422,
        content={
            "success": False,
            "error": "Invalid request data", 
            "error_code": "VALIDATION_ERROR",
            "details": exc.errors()
        }
    )

# =========================================================================
# STARTUP MESSAGE
# =========================================================================

# The following block is removed for Hugging Face Spaces compatibility:
# if __name__ == "__main__":
#     import uvicorn
#     print("🩺 Starting Medical AI Assistant API...")
#     print("📚 Documentation available at: http://localhost:8000/docs")
#     print("🔄 Alternative docs at: http://localhost:8000/redoc")
#     uvicorn.run(
#         app,
#         host="0.0.0.0",
#         port=8000,
#         log_level="info",
#         reload=False
#     )