import asyncio
import json
import base64
from app.config import GEMINI_API_KEY, GEMINI_MODEL


class GeminiClient:

    def __init__(self):
        self.api_key = GEMINI_API_KEY
        self.model = GEMINI_MODEL or "gemini-1.5-flash"

    async def extract_from_document(
        self,
        file_bytes: bytes,
        mime_type: str,
        patient_context: dict,
    ) -> dict:
        if not self.api_key or self.api_key in ["your_gemini_key", "your_gemini_api_key", ""]:
            return {
                "document_type": "medical_report",
                "extracted_text": "Sample HbA1c 7.2%, Blood Pressure 130/85 mmHg",
                "findings": [
                    {"test": "HbA1c", "value": "7.2", "unit": "%", "interpretation": "Elevated"},
                    {"test": "Blood Pressure", "value": "130/85", "unit": "mmHg", "interpretation": "Stage 1"},
                ],
                "medications_found": ["Metformin 500mg"],
                "allergies_found": [],
                "summary": "Report indicates elevated HbA1c and prehypertension.",
            }

        try:
            from google import genai
            from google.genai import types

            client = genai.Client(api_key=self.api_key)

            prompt = f"""
You are a medical document extraction AI.
Analyze this medical report / prescription / lab result for patient {patient_context.get('name', 'Patient')}.
Extract all lab tests, numeric values, units, diagnosed conditions, medications, and allergies into structured JSON.

Return ONLY JSON:
{{
    "document_type": "lab_report | prescription | discharge_summary | other",
    "extracted_text": "brief summary text",
    "findings": [
        {{
            "test": "test name",
            "value": "numeric or text value",
            "unit": "unit or null",
            "interpretation": "Normal | High | Low | etc."
        }}
    ],
    "medications_found": ["medication 1", "medication 2"],
    "allergies_found": ["allergy 1"],
    "summary": "short clinical summary of the document"
}}
"""
            part = types.Part.from_bytes(data=file_bytes, mime_type=mime_type)
            response = await asyncio.wait_for(
                client.aio.models.generate_content(
                    model=self.model,
                    contents=[part, prompt]
                ),
                timeout=4.0
            )
            cleaned_text = response.text.replace("```json", "").replace("```", "").strip()
            return json.loads(cleaned_text)
        except Exception as e:
            print(f"[Gemini Extraction Warning]: {e}")
            return {
                "document_type": "unknown",
                "findings": [],
                "medications_found": [],
                "allergies_found": [],
                "summary": f"Could not extract document details: {e}",
            }

    async def generate_response(
        self,
        system_prompt: str,
        user_message: str,
    ) -> str:
        if not self.api_key or self.api_key in ["your_gemini_key", "your_gemini_api_key", ""]:
            raise ValueError("Gemini API key not configured")

        from google import genai
        from google.genai import types

        client = genai.Client(api_key=self.api_key)
        full_prompt = f"{system_prompt}\n\nPATIENT SAID: {user_message}"
        config = types.GenerateContentConfig(response_mime_type="application/json")
        response = await asyncio.wait_for(
            client.aio.models.generate_content(
                model=self.model,
                contents=full_prompt,
                config=config,
            ),
            timeout=3.5
        )
        return response.text