import json
from fastapi import APIRouter, Depends, File, HTTPException, UploadFile
from sqlalchemy.orm import Session

from app.ai.gemini_client import GeminiClient
from app.database.database import get_db
from app.schemas.documents import DocumentExtractionResponse
from app.services.onboarding_service import OnboardingService

from fastapi import APIRouter, Depends, File, HTTPException, Request, UploadFile
from app.core.security import get_auth_context, AuthContext, check_rate_limit, verify_session_ownership

router = APIRouter(
    prefix="/v1/documents",
    tags=["Medical Documents"],
)


@router.post(
    "/{session_id}/upload",
    response_model=DocumentExtractionResponse,
)
async def upload_medical_document(
    session_id: str,
    file: UploadFile = File(...),
    req_obj: Request = None,
    auth: AuthContext = Depends(get_auth_context),
    db: Session = Depends(get_db),
):
    if req_obj and req_obj.client:
        check_rate_limit(req_obj.client.host, limit=10, window_sec=60)

    onboarding_service = OnboardingService()
    gemini_client = GeminiClient()

    session = onboarding_service.get_session(db, session_id)
    if not session:
        raise HTTPException(status_code=404, detail="Onboarding session not found")

    # 1b. Verify Session Ownership (IDOR Prevention)
    verify_session_ownership(session.patient_id, auth, detail="Forbidden: You do not have authorization to upload documents to this session.")

    # Enforce MIME type and File Size Limits (10MB MAX)
    allowed_mimes = ["image/jpeg", "image/png", "image/webp", "application/pdf"]
    mime_type = file.content_type or "application/pdf"
    if mime_type not in allowed_mimes:
        raise HTTPException(
            status_code=400,
            detail=f"Unsupported document format: {mime_type}. Allowed formats: JPEG, PNG, WEBP, PDF.",
        )

    file_bytes = await file.read()
    if len(file_bytes) > 10 * 1024 * 1024:
        raise HTTPException(
            status_code=413,
            detail="File size exceeds maximum allowed limit of 10MB.",
        )

    patient_context = {
        "id": session.patient_id,
        "name": session.patient_name,
        "age": session.age,
        "gender": session.gender,
        "category": session.category_name,
    }

    # Extract clinical data with Gemini
    extracted = await gemini_client.extract_from_document(
        file_bytes=file_bytes,
        mime_type=mime_type,
        patient_context=patient_context,
    )

    # Merge document findings into onboarding state
    current_state = json.loads(state.state_json)

    # Append medications if found
    meds_found = extracted.get("medications_found", [])
    if meds_found:
        existing_meds = current_state.get("medications", {})
        items = existing_meds.get("items", []) if isinstance(existing_meds, dict) else (existing_meds or [])
        items = list(set(items + meds_found))
        current_state["medications"] = {"status": "reported", "items": items}

    # Append allergies if found
    allergies_found = extracted.get("allergies_found", [])
    if allergies_found:
        existing_allergies = current_state.get("allergies", {})
        items = existing_allergies.get("items", []) if isinstance(existing_allergies, dict) else (existing_allergies or [])
        items = list(set(items + allergies_found))
        current_state["allergies"] = {"status": "reported", "items": items}

    # Add document findings to state
    current_state["document_findings"] = extracted.get("findings", [])

    onboarding_service.update_state(db=db, state=state, data=current_state)

    return DocumentExtractionResponse(
        success=True,
        session_id=session_id,
        document_type=extracted.get("document_type", "medical_report"),
        summary=extracted.get("summary", "Document analyzed successfully"),
        findings=extracted.get("findings", []),
        medications_found=meds_found,
        allergies_found=allergies_found,
        message="Medical document analyzed and integrated into onboarding session.",
    )