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

    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 == "your_gemini_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:
            import google.generativeai as genai
            genai.configure(api_key=self.api_key)
            model = genai.GenerativeModel(self.model)

            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 = {
                "mime_type": mime_type,
                "data": file_bytes,
            }
            response = await model.generate_content_async([part, prompt])
            cleaned_text = response.text.replace("```json", "").replace("```", "").strip()
            return json.loads(cleaned_text)
        except Exception as e:
            print(f"[Gemini Extraction Error]: {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")

        import google.generativeai as genai
        genai.configure(api_key=self.api_key)
        model = genai.GenerativeModel(
            model_name=self.model or "gemini-1.5-flash",
            generation_config={"response_mime_type": "application/json"},
        )
        full_prompt = f"{system_prompt}\n\nPATIENT SAID: {user_message}"
        response = await model.generate_content_async(full_prompt)
        return response.text