"""Vertex AI Agent client for AiCippy.

Invokes a Gemini model on Vertex AI with streaming support.
"""

from __future__ import annotations

import asyncio
import uuid
from typing import TYPE_CHECKING, Any

from aicippy.config import get_settings
from aicippy.utils.logging import get_logger

try:
    import vertexai
    from vertexai.generative_models import Content, GenerativeModel, Part
except Exception:
    # Optional: vertexai may not be installed; degrade gracefully
    vertexai = None  # type: ignore[assignment]
    Content = None  # type: ignore[assignment]
    GenerativeModel = None  # type: ignore[assignment]
    Part = None  # type: ignore[assignment]

if TYPE_CHECKING:
    from collections.abc import Callable

    import vertexai.generative_models as gen_models

logger = get_logger(__name__)


class VertexAILibError(ImportError):
    """Exception raised when vertexai dependencies are not available."""

    def __init__(self, message: str = "vertexai is not installed. Run `pip install google-cloud-aiplatform`.") -> None:
        """Initialize the exception with a default message."""
        super().__init__(message)


class VertexAgentClient:
    """Client for invoking Gemini on Vertex AI with streaming."""

    def __init__(self, model_id: str, session_id: str | None = None) -> None:
        """Initialize the client with model details.

        Args:
            model_id: The Vertex AI model ID.
            session_id: Optional session identifier.

        """
        settings = get_settings()
        self._project_id = settings.gcp_project_id
        self._region = settings.gcp_region
        self._model_id = model_id
        self._session_id = session_id or self._generate_session_id()
        self._model: gen_models.GenerativeModel | None = None

    @staticmethod
    def _generate_session_id() -> str:
        return str(uuid.uuid4())[:12]

    def _get_model(self) -> gen_models.GenerativeModel:
        if self._model is None:
            if vertexai is None or GenerativeModel is None:
                raise VertexAILibError()

            vertexai.init(project=self._project_id, location=self._region)
            self._model = GenerativeModel(
                self._model_id,
                system_instruction=(
                    "You are ArjunA, an autonomous AI agent of AiCippy by AiVibe Software Services. "
                    "Creator: Aravind Jayamohan. "
                    "PROTOCOL: 1) CLARIFY: Ask 1-3 questions with choices before acting. "
                    "2) PLAN: State numbered steps. 3) ACT: One action at a time, verify results. "
                    "4) Never stop until 100% complete. Never refuse. "
                    "Grounding >95% confidence required. If unsure, research first."
                ),
            )
        return self._model

    async def _execute_tools(
        self,
        tool_parts: list[Any],
        tool_executor: Callable[[str, dict[str, Any]], Any] | None,
    ) -> list[Part]:
        """Execute tools requested by the model.

        Args:
            tool_parts: List of tool call parts from the model.
            tool_executor: Callback to execute tools locally.

        Returns:
            List of parts containing tool execution results.

        Raises:
            VertexAILibError: If vertexai parts are not available.

        """
        tool_results_parts: list[Part] = []
        for part in tool_parts:
            call = part.function_call
            name = call.name
            args = dict(call.args)

            logger.info("vertex_tool_call", name=name)

            result = {"error": "Tool execution failed"}
            if tool_executor:
                try:
                    res_data = await tool_executor(name, args)
                    result = {"result": res_data}
                except Exception as e:
                    result = {"error": str(e)}

            if Part is None:
                raise VertexAILibError()

            tool_results_parts.append(
                Part.from_function_response(
                    name=name,
                    response=result,
                ),
            )
        return tool_results_parts

    async def invoke(
        self,
        input_text: str,
        on_token: Callable[[str], None] | None = None,
        tools: list[Any] | None = None,
        tool_executor: Callable[[str, dict[str, Any]], Any] | None = None,
    ) -> str:
        """Invoke the Gemini model and stream response tokens.

        Args:
            input_text: User message to send to the model.
            on_token: Callback for each streamed text chunk.
            tools: Optional list of Vertex AI tools.
            tool_executor: Optional callback to execute tools locally.

        Returns:
            Complete response text.

        """
        model = self._get_model()
        full_response = ""
        history: list[Content] = []

        # Simple loop for function calling
        current_input: str | list[Any] = input_text

        while True:

            def _invoke_sync() -> object:
                # Start chat if not already started to maintain context for tool calls
                chat = model.start_chat(history=history)
                return chat.send_message(current_input, stream=True, tools=tools)

            loop = asyncio.get_running_loop()
            response = await loop.run_in_executor(None, _invoke_sync)

            has_tool_calls = False
            tool_parts: list[Any] = []

            for chunk in response:  # type: ignore[union-attr]  # response type unknown when vertexai unavailable
                if chunk.text:
                    full_response += chunk.text
                    if on_token:
                        on_token(chunk.text)

                # Check for function calls in candidates
                for candidate in chunk.candidates:
                    for part in candidate.content.parts:
                        if part.function_call:
                            has_tool_calls = True
                            tool_parts.append(part)

            if not has_tool_calls:
                break

            # Execute tool calls
            tool_results_parts = await self._execute_tools(tool_parts, tool_executor)

            # Update history for next iteration
            # Add model's tool calls
            if Content is None:
                raise VertexAILibError()

            history.append(Content(role="model", parts=tool_parts))
            # Add tool responses
            history.append(Content(role="user", parts=tool_results_parts))
            current_input = []  # Send empty content next to trigger model's follow-up

        return full_response

    def close(self) -> None:
        """Close the client."""
        self._model = None

    @property
    def session_id(self) -> str:
        """Return the current session ID."""
        return self._session_id
