"""Gemini AI client for KPI analysis."""

from __future__ import annotations

import json
import re
from typing import Any

import google.generativeai as genai

from .prompts import PromptLibrary


class GeminiClient:
    """Client for Gemini AI API."""

    def __init__(
        self,
        api_key: str,
        model: str = "gemini-2.0-flash-exp",
        temperature: float = 0.3,
    ):
        """Initialize Gemini client.

        Args:
            api_key: Gemini API key.
            model: Model name to use.
            temperature: Temperature for generation (lower = more consistent).
        """
        genai.configure(api_key=api_key)
        self.model = genai.GenerativeModel(model)
        self.temperature = temperature

        # Generation config
        self.generation_config = genai.GenerationConfig(
            temperature=temperature,
            top_p=0.8,
            top_k=40,
            max_output_tokens=4096,
        )

        # Safety settings (allow all for business analysis)
        self.safety_settings = [
            {"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_NONE"},
            {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_NONE"},
            {"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "threshold": "BLOCK_NONE"},
            {"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "BLOCK_NONE"},
        ]

    def generate(self, prompt: str, system_prompt: str | None = None) -> str:
        """Generate response from Gemini.

        Args:
            prompt: User prompt.
            system_prompt: Optional system prompt.

        Returns:
            Generated text response.
        """
        full_prompt = prompt
        if system_prompt:
            full_prompt = f"{system_prompt}\n\n{prompt}"

        response = self.model.generate_content(
            full_prompt,
            generation_config=self.generation_config,
            safety_settings=self.safety_settings,
        )

        return response.text

    def generate_json(
        self,
        prompt: str,
        system_prompt: str | None = None,
    ) -> dict[str, Any]:
        """Generate JSON response from Gemini.

        Args:
            prompt: User prompt.
            system_prompt: Optional system prompt.

        Returns:
            Parsed JSON response.
        """
        text = self.generate(prompt, system_prompt)
        return self._parse_json_response(text)

    def _parse_json_response(self, text: str) -> dict[str, Any]:
        """Parse JSON from response text.

        Handles responses with markdown code blocks.

        Args:
            text: Response text.

        Returns:
            Parsed JSON dictionary.
        """
        # Try to extract JSON from markdown code block
        json_match = re.search(r"```(?:json)?\s*([\s\S]*?)\s*```", text)

        if json_match:
            json_str = json_match.group(1)
        else:
            # Try to find JSON object directly
            json_match = re.search(r"\{[\s\S]*\}", text)
            if json_match:
                json_str = json_match.group(0)
            else:
                # Return raw text if no JSON found
                return {"raw_response": text, "parsed": False}

        try:
            return json.loads(json_str)
        except json.JSONDecodeError:
            # Try to fix common JSON issues
            fixed = self._fix_json(json_str)
            try:
                return json.loads(fixed)
            except json.JSONDecodeError:
                return {"raw_response": text, "parsed": False}

    def _fix_json(self, json_str: str) -> str:
        """Try to fix common JSON formatting issues.

        Args:
            json_str: Potentially malformed JSON string.

        Returns:
            Fixed JSON string.
        """
        # Remove trailing commas
        json_str = re.sub(r",\s*}", "}", json_str)
        json_str = re.sub(r",\s*\]", "]", json_str)

        # Fix unquoted keys
        json_str = re.sub(r"(\s*)(\w+)(\s*):", r'\1"\2"\3:', json_str)

        return json_str


class MockGeminiClient:
    """Mock Gemini client for testing."""

    def __init__(self, responses: dict[str, dict[str, Any]] | None = None):
        """Initialize mock client.

        Args:
            responses: Dictionary mapping prompt keywords to responses.
        """
        self.responses = responses or {}
        self.call_history: list[str] = []

    def set_response(self, keyword: str, response: dict[str, Any]) -> None:
        """Set response for a keyword.

        Args:
            keyword: Keyword to match in prompt.
            response: Response to return.
        """
        self.responses[keyword] = response

    def generate(self, prompt: str, system_prompt: str | None = None) -> str:
        """Generate mock response."""
        self.call_history.append(prompt)

        for keyword, response in self.responses.items():
            if keyword.lower() in prompt.lower():
                return json.dumps(response, ensure_ascii=False)

        return json.dumps({"summary": "Mock response", "parsed": True})

    def generate_json(
        self,
        prompt: str,
        system_prompt: str | None = None,
    ) -> dict[str, Any]:
        """Generate mock JSON response."""
        self.call_history.append(prompt)

        for keyword, response in self.responses.items():
            if keyword.lower() in prompt.lower():
                return response

        return {"summary": "Mock response", "parsed": True}