"""
gemini.py

LeafDoctor AI
Google Gemini Vision Client

Part 1
"""

from __future__ import annotations

import json
import mimetypes
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Any

from google import genai
from google.genai import types
import json
from pydantic import BaseModel
from config import settings


@dataclass(slots=True)
class DiagnosisResult:
    crop: str
    label: str
    confidence: float
    raw_response: str
    model: str

class DiagnosisSchema(BaseModel):
        crop: str
        label: str
        confidence: float

class GeminiClient:

    def __init__(self):

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

        self.model = settings.GEMINI_MODEL

        self.crop_mapping = self._load_crop_mapping()

        self.allowed_crops = sorted(
            self.crop_mapping.keys()
        )

    # --------------------------------------------------
    # Mapping
    # --------------------------------------------------

    def _load_crop_mapping(self):

        mapping_file = (
            settings.data_path /
            "crop_label_mapping.json"
        )

        if not mapping_file.exists():
            raise FileNotFoundError(
                mapping_file
            )

        with open(
            mapping_file,
            "r",
            encoding="utf-8"
        ) as fp:

            mapping = json.load(fp)

        return mapping

    # --------------------------------------------------
    # Image Validation
    # --------------------------------------------------

    @staticmethod
    def validate_image(
        image_path: str | Path
    ):

        image_path = Path(image_path)

        if not image_path.exists():
            raise FileNotFoundError(
                image_path
            )

        if not image_path.is_file():
            raise ValueError(
                "Invalid image path."
            )

        extension = (
            image_path.suffix.lower()
        )

        if extension not in (
            ".jpg",
            ".jpeg",
            ".png",
            ".webp",
        ):
            raise ValueError(
                f"Unsupported image: {extension}"
            )

        size_mb = (
            image_path.stat().st_size
            / 1024
            / 1024
        )

        if (
            size_mb >
            settings.MAX_UPLOAD_SIZE_MB
        ):
            raise ValueError(
                "Image size exceeds "
                f"{settings.MAX_UPLOAD_SIZE_MB} MB."
            )

    # --------------------------------------------------
    # Prompt
    # --------------------------------------------------

    def _build_prompt(self):

        crops = "\n".join(
            f"- {crop}"
            for crop in self.allowed_crops
        )

        return f"""
You are an expert agricultural crop disease
diagnosis AI.

Your ONLY responsibility is diagnosis.

DO NOT recommend

- medicine
- treatment
- dosage
- fertilizer
- pesticide
- chemical

Those are handled separately.

Allowed Crops

{crops}

Instructions

1.
Identify crop.

2.
Identify disease/problem.

3.
Return confidence.

4.
Return ONLY JSON.

JSON Format

{{
    "crop":"Rice",
    "label":"Blast",
    "confidence":0.98
}}

Rules

Crop must be from allowed crops.

Return only JSON.

Do not use markdown.

Do not explain.
"""

    # --------------------------------------------------
    # Helpers
    # --------------------------------------------------

    @staticmethod
    def _mime_type(
        image_path: Path
    ) -> str:

        mime, _ = mimetypes.guess_type(
            image_path
        )

        if mime is None:
            mime = "image/jpeg"

        return mime

    @staticmethod
    def _extract_json(
        text: str
    ) -> dict[str, Any]:

        text = text.strip()

        if text.startswith("```"):

            lines = text.splitlines()

            if lines[0].startswith("```"):
                lines = lines[1:]

            if lines[-1].startswith("```"):
                lines = lines[:-1]

            text = "\n".join(lines)

        start = text.find("{")

        end = text.rfind("}")

        if start == -1 or end == -1:
            raise ValueError(
                "Gemini returned invalid JSON."
            )

        return json.loads(
            text[start:end + 1]
        )

    def _validate_result(
        self,
        data: dict
    ) -> DiagnosisResult:

        crop = (
            str(
                data.get("crop", "")
            )
            .strip()
        )

        label = (
            str(
                data.get("label", "")
            )
            .strip()
        )

        confidence = float(
            data.get("confidence", 0)
        )

        if crop not in self.crop_mapping:

            raise ValueError(
                f"Unsupported crop: {crop}"
            )

        allowed = (
            self.crop_mapping[crop]
            ["labels"]
        )

        if label not in allowed:

            raise ValueError(
                f"Invalid label '{label}' "
                f"for crop '{crop}'."
            )

        if (
            confidence < 0
            or confidence > 1
        ):
            raise ValueError(
                "Confidence must be "
                "between 0 and 1."
            )

        return DiagnosisResult(
            crop=crop,
            label=label,
            confidence=confidence,
            raw_response="",
            model=self.model,
        )
    
    # --------------------------------------------------
    # Gemini API
    # --------------------------------------------------
    def _call_gemini(
        self,
        image_path: str | Path,
    ) -> DiagnosisResult:

        image_path = Path(image_path)

        prompt = self._build_prompt()

        mime_type = self._mime_type(image_path)

        with open(image_path, "rb") as fp:
            image_bytes = fp.read()

        response = self.client.models.generate_content(
            model=self.model,
            contents=[
                types.Part.from_text(
                    text=prompt,
                ),
                types.Part.from_bytes(
                    data=image_bytes,
                    mime_type=mime_type,
                ),
            ],
            config=types.GenerateContentConfig(
                temperature=settings.GEMINI_TEMPERATURE,
                max_output_tokens=settings.GEMINI_MAX_OUTPUT_TOKENS,
                 thinking_config=types.ThinkingConfig(
                        thinking_budget=0
                    ),
                response_mime_type="application/json",
                response_schema=DiagnosisSchema,
            ),
        )

        if response.parsed is None:
            raise RuntimeError(
                f"""
        Gemini returned an invalid response.

        TEXT:
        {response.text}

        RESPONSE:
        {response}
        """
            )

        data = response.parsed.model_dump()

        crop = data["crop"].strip()

        label = data["label"].strip()

        confidence = float(data["confidence"])

        if crop not in self.crop_mapping:
            raise ValueError(
                f"Unsupported crop '{crop}'."
            )

        allowed_labels = self.crop_mapping[crop]["labels"]

        if label not in allowed_labels:
            raise ValueError(
                f"Invalid label '{label}' for crop '{crop}'."
            )

        if not 0.0 <= confidence <= 1.0:
            raise ValueError(
                "Confidence must be between 0 and 1."
            )

        return DiagnosisResult(
            crop=crop,
            label=label,
            confidence=confidence,
            raw_response=json.dumps(
                data,
                ensure_ascii=False,
                indent=2,
            ),
            model=self.model,
        )

    # --------------------------------------------------
    # Retry
    # --------------------------------------------------

    def _call_with_retry(
        self,
        image_path: str | Path,
        retries: int = 3,
    ) -> DiagnosisResult:

        last_exception = None

        for attempt in range(retries):

            try:

                return self._call_gemini(
                    image_path
                )

            except Exception as ex:

                last_exception = ex

                if attempt < retries - 1:

                    wait = (
                        attempt + 1
                    ) * 2

                    time.sleep(wait)

        raise RuntimeError(
            f"Gemini diagnosis failed: "
            f"{last_exception}"
        )

    # --------------------------------------------------
    # Logging
    # --------------------------------------------------

    @staticmethod
    def _log_result(
        result: DiagnosisResult
    ):

        print()

        print("=" * 60)

        print("LeafDoctor AI Diagnosis")

        print("=" * 60)

        print(
            f"Crop       : {result.crop}"
        )

        print(
            f"Problem    : {result.label}"
        )

        print(
            f"Confidence : "
            f"{result.confidence:.2%}"
        )

        print(
            f"Model      : "
            f"{result.model}"
        )

        print("=" * 60)

        print()

    # --------------------------------------------------
    # Health Check
    # --------------------------------------------------

    def test_connection(self):

        response = self.client.models.generate_content(
            model=self.model,
            contents="Reply with OK",
        )

        if not response.text:
            raise RuntimeError(
                "Gemini connection failed."
            )

        return True
    

        # --------------------------------------------------
    # Public API
    # --------------------------------------------------

    def diagnose(
        self,
        image_path: str | Path,
    ) -> DiagnosisResult:
        """
        Complete diagnosis workflow.

        1. Validate image
        2. Call Gemini
        3. Validate response
        4. Return DiagnosisResult
        """

        self.validate_image(image_path)

        result = self._call_with_retry(
            image_path=image_path
        )

        self._log_result(result)

        return result

    # --------------------------------------------------
    # Utility Methods
    # --------------------------------------------------

    def supported_crops(self) -> list[str]:
        """
        Returns all supported crops.
        """

        return self.allowed_crops.copy()

    def supported_labels(
        self,
        crop: str,
    ) -> list[str]:
        """
        Returns valid labels for a crop.
        """

        if crop not in self.crop_mapping:
            return []

        return self.crop_mapping[crop]["labels"]

    def is_valid_crop(
        self,
        crop: str,
    ) -> bool:

        return crop in self.crop_mapping

    def is_valid_label(
        self,
        crop: str,
        label: str,
    ) -> bool:

        if crop not in self.crop_mapping:
            return False

        return (
            label
            in self.crop_mapping[crop]["labels"]
        )

    def reload_mapping(self):
        """
        Reload crop mapping without
        restarting the application.
        """

        self.crop_mapping = (
            self._load_crop_mapping()
        )

        self.allowed_crops = sorted(
            self.crop_mapping.keys()
        )

    # --------------------------------------------------
    # Singleton
    # --------------------------------------------------

    _instance = None

    @classmethod
    def instance(cls):

        if cls._instance is None:

            cls._instance = cls()

        return cls._instance


gemini_client = GeminiClient()