"""
For sending prompts to Google Generative AI (genai), receiving responses, and managing model lists and cache.

Features:
- Send prompts to GenAI and cache responses for reproducibility.
- List available models from the GenAI API and save as both JSON and TXT for validation and reference.
- Validate selected model names against the cached JSON model list.
- Unified workflow for model management, matching other LLM modules in the project.
"""
# Metadata
__author__ = ["Andy Turner <agdturner@gmail.com>"]
__version__ = "0.3.0"
__copyright__ = "Copyright (c) 2026 GABM contributors, University of Leeds"

# Standard library imports
import ast
# Google Generative AI client library
import google.genai as genai
# LLM service base class
from .llm_service import LLMService
# Shared utilities for caching and logging
from .utils import pre_send_check_and_cache, call_and_cache_response, cache_and_log, write_models_json_and_txt


class GenAIService(LLMService):
    """
    Service class for Google Generative AI LLM integration. 
    Handles prompt sending, response caching, logging, and model listing.
    """
    SERVICE_NAME = "genai"

    def send(self, api_key, message, model="models/gemini-2.5-pro"):
        """
        Send a prompt to Google Generative AI and return the response object.
        Caches and logs the response for reproducibility.
        Args:
            api_key (str): Google API key.
            message (str): Prompt to send.
            model (str): Model name (default: "models/gemini-2.5-pro").
        Returns:
            Response object (dict) or None on error.
        """
        cached = pre_send_check_and_cache(api_key, message, model, self.cache, self.logger, self.SERVICE_NAME, self.API_KEY_ENV_VAR)
        if cached is not None:
            return cached
        cache_key = (message, model)
        # Set API key for google-genai
        import os
        os.environ["GOOGLE_API_KEY"] = api_key
        def api_call():
            model_obj = genai.GenerativeModel(model)
            response = model_obj.generate_content(message)
            # Convert to dict if possible
            if hasattr(response, 'to_dict'):
                return response.to_dict()
            elif hasattr(response, '__dict__'):
                return dict(response.__dict__)
            else:
                return response
        return call_and_cache_response(
            api_call,
            cache_and_log,
            self.cache,
            cache_key,
            self.cache_path,
            self.jsonl_path,
            message,
            model,
            api_key,
            self.logger,
            self.SERVICE_NAME,
            self.list_available_models,
            extract_text_from_response=self.extract_text_from_response
        )

    def list_available_models(self, api_key):
        """
        Return a static list of supported Gemini model names for google-genai.
        The google-genai package does not provide a list_models() method.
        See: https://ai.google.dev/models/gemini
        """
        models = [
            {"name": "models/gemini-1.5-pro-latest", "description": "Gemini 1.5 Pro (latest)", "supported_generation_methods": ["generateContent"]},
            {"name": "models/gemini-1.0-pro", "description": "Gemini 1.0 Pro", "supported_generation_methods": ["generateContent"]},
            {"name": "models/gemini-1.0-pro-001", "description": "Gemini 1.0 Pro (001)", "supported_generation_methods": ["generateContent"]},
            {"name": "models/gemini-1.0-pro-latest", "description": "Gemini 1.0 Pro (latest)", "supported_generation_methods": ["generateContent"]},
            {"name": "models/gemini-1.0-pro-vision-latest", "description": "Gemini 1.0 Pro Vision (latest)", "supported_generation_methods": ["generateContent"]},
            {"name": "models/gemini-1.5-flash-latest", "description": "Gemini 1.5 Flash (latest)", "supported_generation_methods": ["generateContent"]},
        ]
        def formatter(model):
            name = model.get('name', 'N/A')
            desc = model.get('description', 'N/A')
            methods = model.get('supported_generation_methods', 'N/A')
            return (f"Model ID: {name}\n"
                    f"  Description: {desc}\n"
                    f"  Supported methods: {methods}\n")
        self.logger.info("Writing GenAI model list to JSON and TXT (static list, see https://ai.google.dev/models/gemini).")
        write_models_json_and_txt(
            models,
            self.cache_path.parent / "models.json",
            self.cache_path.parent / "models.txt",
            formatter,
            header=f"Available {self.SERVICE_NAME.capitalize()} models and supported methods (static list):\n"
        )
        return models

    def extract_text_from_response(self, response):
        """
        Extract the text content from a GenAI response object for logging.
        Recursively searches for the first 'text' value in any nested structure.
        Logs a warning if no text is found.
        """
        def find_text(obj):
            if isinstance(obj, dict):
                # Direct text field
                if "text" in obj and isinstance(obj["text"], str):
                    return obj["text"]
                # Recursively search all values
                for v in obj.values():
                    result = find_text(v)
                    if result:
                        return result
            elif isinstance(obj, list):
                for item in obj:
                    result = find_text(item)
                    if result:
                        return result
            return None

        # Debug: log the type and a preview of the response
        if hasattr(self, 'logger') and self.logger:
            self.logger.debug(f"[GenAI] extract_text_from_response type={type(response)} preview={str(response)[:200]}")

        # If response is a string that looks like a dict, try to eval/parse it
        if isinstance(response, str) and response.strip().startswith("{"):
            try:
                resp_dict = ast.literal_eval(response)
                text = find_text(resp_dict)
                if text:
                    return text
            except Exception:
                pass
        # If response is a dict or list, search recursively
        if isinstance(response, (dict, list)):
            text = find_text(response)
            if text:
                return text
        # If response is an object, try attribute access
        if hasattr(response, "text") and isinstance(response.text, str):
            return response.text
        if hasattr(response, "candidates") and response.candidates:
            text = find_text(response.candidates)
            if text:
                return text
        # Fallback
        if hasattr(self, 'logger') and self.logger:
            self.logger.warning(f"[GenAI] Could not extract text from response: {str(response)[:200]}")
        return str(response)
