"""Integration for Google Gemini AI Assistant.

This integration provides AI-powered analysis and chat capabilities for XSOAR users.
Supports both Google AI Studio (API key) and Vertex AI (service account) authentication.
"""

import demistomock as demisto
from CommonServerPython import *  # noqa # pylint: disable=unused-wildcard-import
from CommonServerUserPython import *  # noqa

""" IMPORTS """
import json
from typing import Any
from uuid import uuid4

from google.oauth2 import service_account
from google.auth.transport.requests import Request

""" CONSTANTS """
DATE_FORMAT = "%Y-%m-%dT%H:%M:%SZ"  # ISO8601
SUPPORTED_MODELS = [
    # Stable models
    "gemini-2.5-pro",
    "gemini-2.5-flash",
    # Deprecated (Legacy Only) - operational until June 1, 2026
    "gemini-2.0-flash",
    "gemini-2.0-flash-lite",
    # Preview models
    "gemini-3.1-pro-preview",
    "gemini-3.1-flash-preview",
    "gemini-3.1-flash-lite",
]
AUTH_TYPE_AI_STUDIO = "AI Studio API Key"
AUTH_TYPE_VERTEX_AI = "Vertex AI Service Account"
VERTEX_AI_BASE_URL = "https://aiplatform.googleapis.com"
GOOGLE_AUTH_SCOPE = "https://www.googleapis.com/auth/cloud-platform"


class Client(BaseClient):
    """Client to interact with the Google Gemini API.

    Supports both AI Studio (API key) and Vertex AI (service account) authentication.
    It inherits from BaseClient which handles proxy, SSL verification, etc.
    """

    def __init__(
        self,
        base_url: str,
        verify: bool,
        proxy: bool,
        auth_type: str,
        model: str = "gemini-2.5-flash",
        max_tokens: int = 1024,
        temperature: float | None = None,
        top_p: float | None = None,
        top_k: int | None = None,
        api_key: str | None = None,
        service_account_json: str | None = None,
        project_id: str | None = None,
        location: str = "global",
    ):
        """Initialize Client class.

        :param base_url: The base URL of the API.
        :param verify: Whether to verify SSL certificate.
        :param proxy: Whether to use system proxy settings.
        :param auth_type: Authentication type - AI Studio API Key or Vertex AI Service Account.
        :param model: The default Gemini model to use for requests.
        :param max_tokens: Default maximum tokens for responses.
        :param temperature: Default temperature for response generation.
        :param top_p: Default top-p value for response generation.
        :param top_k: Default top-k value for response generation.
        :param api_key: API key for AI Studio authentication.
        :param service_account_json: Service account JSON key for Vertex AI authentication.
        :param project_id: Google Cloud project ID for Vertex AI.
        :param location: Google Cloud location for Vertex AI (default: global).
        """
        super().__init__(base_url=base_url, verify=verify, proxy=proxy)
        self.auth_type = auth_type
        self.model = model
        self.max_tokens = max_tokens
        self.temperature = temperature
        self.top_p = top_p
        self.top_k = top_k

        if auth_type == AUTH_TYPE_AI_STUDIO:
            self.api_key = api_key
            self._headers = {
                "Content-Type": "application/json",
                "Accept": "application/json",
                "x-goog-api-key": self.api_key,
            }
        else:
            try:
                self.service_account_info: dict[str, Any] = json.loads(service_account_json) if service_account_json else {}
            except json.JSONDecodeError as e:
                raise ValueError(f"Invalid Service Account JSON provided: {e}")
            self.project_id = project_id
            self.location = location or "global"
            self._credentials = service_account.Credentials.from_service_account_info(
                self.service_account_info,
                scopes=[GOOGLE_AUTH_SCOPE],
            )
            self._headers = {
                "Content-Type": "application/json",
                "Accept": "application/json",
            }

    def _get_access_token(self) -> str:
        """Get a valid access token for Vertex AI using google-auth credentials.

        Refreshes the token automatically when expired.

        :return: Valid OAuth2 access token string.
        """
        if not self._credentials.valid:
            demisto.debug("Refreshing Vertex AI access token")
            self._credentials.refresh(Request())
        return self._credentials.token

    def _get_request_headers(self) -> dict[str, str] | None:
        """Get the appropriate request headers based on auth type.

        For AI Studio, returns None to use the default self._headers (with API key).
        For Vertex AI, returns headers with a fresh Bearer token.

        :return: Headers dict for Vertex AI, or None for AI Studio.
        """
        if self.auth_type == AUTH_TYPE_AI_STUDIO:
            return None
        access_token = self._get_access_token()
        return {
            "Content-Type": "application/json",
            "Accept": "application/json",
            "Authorization": f"Bearer {access_token}",
        }

    def _get_url_suffix(self, model: str) -> str:
        """Get the appropriate URL suffix based on auth type and model.

        :param model: The model name to use.
        :return: URL suffix string for the generateContent endpoint.
        """
        if self.auth_type == AUTH_TYPE_AI_STUDIO:
            return f"/v1beta/models/{model}:generateContent"
        return f"/v1/projects/{self.project_id}/locations/{self.location}/publishers/google/models/{model}:generateContent"

    def send_chat_message(
        self,
        prompt: str,
        model: str | None = None,
        history: list[dict[str, Any]] | None = None,
    ) -> dict[str, Any]:
        """Send a chat message to the Gemini API with optional conversation history.

        Conversation history format:
        [
            {
                "role": "user",
                "parts": [{"text": "Previous user message"}]
            },
            {
                "role": "model",
                "parts": [{"text": "Previous AI response"}]
            }
        ]
        :param prompt: The user's prompt/question.
        :param model: The Gemini model to use (defaults to instance default).
        :param history: Optional conversation history in Gemini format.
        :return: Dictionary containing the API response.
        """
        selected_model = model or self.model
        contents = []

        if history:
            contents.extend(history)

        # Add current user prompt
        contents.append({"role": "user", "parts": [{"text": prompt}]})

        # Build generation config using instance defaults
        generation_config = assign_params(
            maxOutputTokens=self.max_tokens, temperature=self.temperature, topP=self.top_p, topK=self.top_k
        )

        request_body = {"contents": contents, "generationConfig": generation_config}

        return self._http_request(
            method="POST",
            url_suffix=self._get_url_suffix(selected_model),
            json_data=request_body,
            headers=self._get_request_headers(),
        )


def test_module(client: Client):
    """Tests API connectivity and authentication.

    Uses a simple chat message to verify that the API is reachable and the provided token is valid.

    :param client: Google Gemini API client.
    :return: 'ok' if successful, or an error message string.
    """
    try:
        client.send_chat_message("Hello, please respond with 'OK' to test connectivity.")
        return "ok"
    except DemistoException as e:
        err_msg = e.message
        try:
            err_msg = demisto.get(e.res.json(), "error.message", err_msg)
        except Exception:
            pass
        return_error(f"An unexpected error occurred during connectivity test: {err_msg}")


def google_gemini_send_message_command(client: Client, args: dict[str, Any]):
    """Command function to send a chat message to the Google Gemini API with optional conversation history.

    :param client: Google Gemini API client.
    :param args: Dictionary of command arguments (prompt, model, history, save_conversation).
    :return: CommandResults object(s) with outputs and readable representation.
    """
    prompt = str(args.get("prompt", ""))
    model = args.get("model", None)
    history_arg = args.get("history", [])
    save_conversation = argToBoolean(args.get("save_conversation", False))

    if not prompt:
        raise ValueError("The 'prompt' argument is required.")

    history = []
    if history_arg:
        try:
            if isinstance(history_arg, str):
                history = json.loads(history_arg)
            elif isinstance(history_arg, list):
                history = history_arg
        except json.JSONDecodeError:
            raise ValueError("History must be valid JSON array of conversation objects.")

    conversation_id = None
    outputs_key_field = "prompt"
    if save_conversation:
        context = demisto.context()
        existing_history = None

        if google_gemini_context := demisto.get(context, "GoogleGemini.Chat"):
            if isinstance(google_gemini_context, dict) and "History" in google_gemini_context:
                existing_history = google_gemini_context["History"]
                conversation_id = google_gemini_context["ConversationId"]

            elif isinstance(google_gemini_context, list):
                for item in reversed(google_gemini_context):
                    if isinstance(item, dict) and "History" in item:
                        existing_history = item["History"]
                        conversation_id = item["ConversationId"]
                        break

        # trying to take the last 2 entries
        if existing_history and isinstance(existing_history, list):
            if len(existing_history) >= 2:
                history = existing_history[-2:]
            else:
                history = existing_history

    response = client.send_chat_message(prompt, model, history)

    content = ""
    finish_reason = ""
    if (candidates := response.get("candidates")) and len(candidates) > 0:
        parts = demisto.get(candidates[0], "content.parts")
        if parts and isinstance(parts, list) and len(parts) > 0 and isinstance(parts[0], dict):
            content = parts[0].get("text")  # type: ignore[assignment]
        else:
            finish_reason = demisto.get(candidates[0], "finishReason")

    if not content:
        content = "No response generated."
        if finish_reason:
            return_warning(f"The model finished before completing the full response, due to {finish_reason}")

    outputs = {"Prompt": prompt, "Response": content, "Model": model or client.model, "Temperature": client.temperature}
    if save_conversation:
        current_conversation = history.copy() if history else []
        current_conversation.append({"role": "user", "parts": [{"text": prompt}]})

        if content and content != "No response generated.":
            current_conversation.append({"role": "model", "parts": [{"text": content}]})

        outputs["History"] = current_conversation
        outputs["ConversationId"] = conversation_id or str(uuid4())
        outputs_key_field = "ConversationId"

    return CommandResults(
        outputs_prefix="GoogleGemini.Chat",
        outputs_key_field=outputs_key_field,
        outputs=outputs,
        raw_response=response,
        readable_output=content,
    )


def main():
    """Main execution function for the integration.

    Parses integration parameters and command arguments, initializes the client,
    and calls the appropriate command function.
    """
    params = demisto.params()
    auth_type = params.get("auth_type", AUTH_TYPE_AI_STUDIO)
    verify_certificate = not argToBoolean(params.get("insecure", False))
    proxy = argToBoolean(params.get("proxy", False))
    model = params.get("model", ["gemini-2.5-flash"])  # use multi select to enable adding custom val
    max_tokens = arg_to_number(params.get("max_tokens", 1024)) or 1024

    # Handle optional parameters - use defaults if empty or not provided
    temperature = arg_to_number(params.get("temperature", "").strip())
    top_p = arg_to_number(params.get("top_p", "").strip())
    top_k = arg_to_number(params.get("top_k", "").strip())

    # Auth-specific parameters
    api_key: str | None = None
    service_account_json: str | None = None
    project_id: str | None = None
    location: str = "global"

    if auth_type == AUTH_TYPE_VERTEX_AI:
        base_url = params.get("url", VERTEX_AI_BASE_URL)
        # Auto-switch from AI Studio default URL to Vertex AI URL
        if base_url == "https://generativelanguage.googleapis.com":
            base_url = VERTEX_AI_BASE_URL
        service_account_json = params.get("service_account_key", {}).get("password")
        project_id = params.get("project_id")
        location = params.get("location", "global") or "global"
        if not service_account_json:
            return_error("Service Account Key JSON is required for Vertex AI authentication.")
            return
        if not project_id:
            return_error("Project ID is required for Vertex AI authentication.")
            return
    else:
        base_url = params.get("url", "https://generativelanguage.googleapis.com")
        api_key = params.get("api_key", {}).get("password")
        if not api_key:
            return_error("API key is not configured. Please configure it in the instance settings.")
            return

    command = demisto.command()
    demisto.debug(f"Command being called is {command}")

    try:
        if len(model) > 1:
            raise DemistoException("Please select one model only.")
        client = Client(
            base_url=base_url,
            verify=verify_certificate,
            proxy=proxy,
            auth_type=auth_type,
            model=model[0],
            max_tokens=max_tokens,
            temperature=temperature,
            top_p=top_p,
            top_k=top_k,
            api_key=api_key,
            service_account_json=service_account_json,
            project_id=project_id,
            location=location,
        )
        args = demisto.args()

        if command == "test-module":
            result = test_module(client)
        elif command == "google-gemini-send-message":
            result = google_gemini_send_message_command(client, args)
        else:
            raise NotImplementedError(f"Command {command} is not implemented")

        return_results(result)

    except Exception as e:
        return_error(f"Failed to execute {command} command.\nError:\n{str(e)}")


if __name__ in ("__main__", "__builtin__", "builtins"):  # pragma: no cover
    main()