from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from huggingface_hub import login
import os
from transformers import pipeline
from PIL import Image
import io
import base64
from datetime import datetime
from typing import Optional

# Login no Hugging Face
#login(new_session=False)

app = FastAPI(
    title="CareLink Medical API",
    description="API para análise de exames médicos usando IA",
    version="1.0.0"
)

# Configurar CORS para permitir requisições do frontend
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # Em produção, especifique o domínio do frontend
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Inicializar o pipeline do modelo (fazer isso uma vez no startup)
print("Carregando modelo MedGemma...")
pipe = None

import os
from transformers import pipeline

def get_hf_token():
        return (
                os.environ.get("HUGGINGFACE_HUB_TOKEN")
                or os.environ.get("HF_TOKEN")
                )

@app.on_event("startup")
async def startup_event():
    global pipe
    token = get_hf_token()
    if not token:
        raise RuntimeError("HF_TOKEN não encontrado no ambiente")
    pipe = pipeline(
        "image-text-to-text",
        model="google/medgemma-4b-it",
        token=token   # 👈 ESTE é o ponto crítico
    )
    print("Modelo MedGemma carregado com sucesso!")
    import logging
    logging.basicConfig(level=logging.INFO)
    logging.info("===================================")
    logging.info("🚀 Application ready")
    logging.info(f"🌐 Space URL: {SPACE_URL}")
    logging.info(f"📡 Gradio endpoint: {SPACE_URL}/run/predict")
    logging.info("===================================")


@app.get("/")
def root():
    return {
        "message": "CareLink Medical API",
        "status": "online",
        "version": "1.0.0",
        "endpoints": {
            "analyze_exam": "/api/analyze-exam (POST)",
            "health": "/health (GET)"
        }
    }


@app.get("/health")
def health_check():
    return {
        "status": "healthy",
        "model_loaded": pipe is not None,
        "timestamp": datetime.now().isoformat()
    }


@app.post("/api/analyze-exam")
async def analyze_exam(
    file: UploadFile = File(...),
    patient_name: Optional[str] = None
):
    """
    Analisa uma imagem de exame médico usando o modelo MedGemma.
    
    - **file**: Imagem do exame (JPEG, PNG, etc)
    - **patient_name**: Nome do paciente (opcional)
    """
    
    if pipe is None:
        raise HTTPException(
            status_code=503,
            detail="Modelo ainda não foi carregado. Tente novamente em alguns instantes."
        )
    
    # Validar tipo de arquivo
    if not file.content_type.startswith('image/'):
        raise HTTPException(
            status_code=400,
            detail="O arquivo enviado não é uma imagem válida"
        )
    
    try:
        # Ler a imagem
        contents = await file.read()
        image = Image.open(io.BytesIO(contents))
        
        # Converter imagem para base64 para retornar ao frontend (opcional)
        buffered = io.BytesIO()
        image.save(buffered, format=image.format or "PNG")
        img_base64 = base64.b64encode(buffered.getvalue()).decode()
        
        # Preparar mensagem para o modelo
        messages = [
            {
                "role": "user",
                "content": [
                    {"type": "image", "image": image},
                    {
                        "type": "text",
                        "text": "Analise esta imagem de exame médico. Descreva o que você observa, "
                                "identifique possíveis achados clínicos e forneça uma análise detalhada. "
                                "Se possível, indique se há algo que requer atenção médica."
                    }
                ]
            }
        ]
        
        # Processar com o modelo
        print(f"Analisando exame para paciente: {patient_name or 'não especificado'}")
        result = pipe(messages)
        
        # Extrair a resposta do modelo
        analysis_text = ""
        if isinstance(result, list) and len(result) > 0:
            if isinstance(result[0], dict) and "generated_text" in result[0]:
                analysis_text = result[0]["generated_text"]
            else:
                analysis_text = str(result[0])
        else:
            analysis_text = str(result)
        
        # Montar resposta estruturada
        response = {
            "success": True,
            "timestamp": datetime.now().isoformat(),
            "patient_name": patient_name,
            "file_info": {
                "filename": file.filename,
                "content_type": file.content_type,
                "size_bytes": len(contents)
            },
            "analysis": {
                "model": "google/medgemma-4b-it",
                "result": analysis_text,
                "confidence": "high"  # Você pode adicionar lógica para calcular confiança
            },
            "image_preview": f"data:image/png;base64,{img_base64[:100]}..."  # Preview truncado
        }
        
        print(f"Análise concluída com sucesso para: {file.filename}")
        return response
        
    except Exception as e:
        print(f"Erro ao processar imagem: {str(e)}")
        raise HTTPException(
            status_code=500,
            detail=f"Erro ao processar a imagem: {str(e)}"
        )


@app.get("/api/test-model")
def test_model():
    """Endpoint para testar se o modelo está funcionando"""
    if pipe is None:
        return {"status": "error", "message": "Modelo não carregado"}
    return {
        "status": "ok",
        "model": "google/medgemma-4b-it",
        "message": "Modelo carregado e pronto para uso"
    }