"""Bedrock Agent client for AiCippy.

Invokes the odin Bedrock Agent via invoke_agent API with streaming support.
All AI responses flow through the Bedrock Agent (not raw invoke_model).
"""

from __future__ import annotations

import asyncio
import re
import uuid
from typing import TYPE_CHECKING, Any, Final

import boto3
from botocore.config import Config

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

if TYPE_CHECKING:
    from collections.abc import Callable

    import botocore.client

logger = get_logger(__name__)

# Bedrock Agent IDs — odin supervisor agent
ODIN_AGENT_ID: Final[str] = "CACLS9AC1J"
ODIN_AGENT_ALIAS_ID: Final[str] = "Y32URLZGJE"  # prod alias

# Timeouts
# No timeout limits — let the agent complete fully
AGENT_READ_TIMEOUT: Final[int] = 900
AGENT_CONNECT_TIMEOUT: Final[int] = 30
_TOOL_BUFFER_LIMIT: Final[int] = 200


def _clean_agent_response(text: str) -> str:
    """Strip tool call artifacts leaked by the Bedrock agent orchestration.

    Removes patterns like:
    - AgentCommunication__sendMessage(recipient="User", content="...")
    - get__codeinterpreteraction__execute(...)
    - ", "recipient": "User"}}
    - </sources> tags
    - Trailing ") or "}) artifacts
    """
    # Remove AgentCommunication__sendMessage wrapper
    text = re.sub(
        r'AgentCommunication__sendMessage\s*\(\s*recipient\s*=\s*"User"\s*,\s*content\s*=\s*"',
        "",
        text,
    )

    # Remove tool call references like get__codeinterpreteraction__execute
    text = re.sub(r"get__\w+__\w+\s*\([^)]*\)", "", text)
    text = re.sub(
        r"(I will call the|To get the final answer, I will call the)\s+`?\w+__\w+`?\s+(function|tool)[^.]*\.",
        "",
        text,
    )

    # Remove trailing JSON artifacts from tool calls
    text = re.sub(r',\s*"recipient"\s*:\s*"User"\s*\}\s*\}?\s*$', "", text)

    # Remove </sources> tags
    text = re.sub(r"</sources>\s*", "", text)

    # Remove trailing ") or "}) from tool call closures
    text = re.sub(r'"\s*\)\s*$', "", text)
    text = re.sub(r'"\s*\}\s*\)\s*$', "", text)

    return text.strip()


class BedrockAgentClient:
    """Client for invoking odin Bedrock Agent with streaming."""

    def __init__(
        self,
        agent_id: str | None = None,
        agent_alias_id: str | None = None,
        session_id: str | None = None,
    ) -> None:
        """Initialize the Bedrock Agent client.

        Args:
            agent_id: Bedrock agent ID. Defaults to configured value.
            agent_alias_id: Bedrock agent alias ID. Defaults to configured value.
            session_id: Session identifier. Auto-generated if not provided.

        """
        settings = get_settings()
        self._agent_id = agent_id or settings.bedrock_agent_id
        self._agent_alias_id = agent_alias_id or settings.bedrock_agent_alias_id
        self._session_id = session_id or self._generate_session_id()
        self._client: botocore.client.BaseClient | None = None

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

    def _get_client(self) -> botocore.client.BaseClient:
        if self._client is None:
            settings = get_settings()
            config = Config(
                read_timeout=AGENT_READ_TIMEOUT,
                connect_timeout=AGENT_CONNECT_TIMEOUT,
                retries={"max_attempts": 2, "mode": "adaptive"},
            )
            session_kwargs: dict[str, object] = {
                "region_name": settings.aws_region,
            }
            if settings.aws_profile:
                session_kwargs["profile_name"] = settings.aws_profile

            session = boto3.Session(**session_kwargs)  # type: ignore[arg-type]  # boto3 stubs incomplete for dict[str, object]
            self._client = session.client("bedrock-agent-runtime", config=config)
        return self._client

    async def invoke(  # noqa: PLR0912, PLR0915 — Complexity justified: streaming response parsing with tool call handling
        self,
        input_text: str,
        on_token: Callable[[str], None] | None = None,
        tool_executor: Callable[[str, dict[str, Any]], Any] | None = None,
    ) -> str:
        """Invoke the odin agent and stream response tokens.

        Args:
            input_text: User message to send to the agent.
            on_token: Callback for each streamed text chunk.
            tool_executor: Optional callback to execute tools locally.

        Returns:
            Complete response text.

        """
        client = self._get_client()
        full_response = ""
        current_input = input_text
        session_state: dict[str, Any] = {
            "promptSessionAttributes": {
                "max_output_tokens": "8192",
                "response_style": "complete_and_detailed",
            },
        }

        while True:

            def _invoke_sync() -> dict:
                return client.invoke_agent(  # type: ignore[attr-defined]  # boto3 stubs incomplete for bedrock-agent-runtime
                    agentId=self._agent_id,
                    agentAliasId=self._agent_alias_id,
                    sessionId=self._session_id,
                    inputText=current_input,
                    sessionState=session_state,
                )

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

            _buffer = ""
            _tool_prefix_seen = False
            tool_calls: list[dict[str, Any]] = []

            for event in response.get("completion", []):
                if "chunk" in event:
                    chunk_bytes = event["chunk"].get("bytes", b"")
                    if chunk_bytes:
                        text = chunk_bytes.decode("utf-8")
                        _buffer += text

                        if not _tool_prefix_seen:
                            if 'content="' in _buffer:
                                _buffer = re.sub(
                                    r'^.*?AgentCommunication__sendMessage\s*\(\s*recipient\s*=\s*"User"\s*,\s*content\s*=\s*"',
                                    "",
                                    _buffer,
                                    count=1,
                                )
                                _tool_prefix_seen = True
                            elif len(_buffer) > _TOOL_BUFFER_LIMIT:
                                _tool_prefix_seen = True

                            if not _tool_prefix_seen:
                                continue

                        to_emit = _buffer
                        _buffer = ""
                        full_response += to_emit
                        if on_token and to_emit:
                            on_token(to_emit)

                elif "invocationInput" in event:
                    tool_calls.append(event["invocationInput"])

            if not tool_calls:
                break

            # Execute tool calls and continue the conversation
            tool_results = []
            for call in tool_calls:
                if "actionGroupInvocationInput" in call:
                    inv = call["actionGroupInvocationInput"]
                    action_group = inv.get("actionGroupName")
                    api_path = inv.get("apiPath")
                    params = {p["name"]: p["value"] for p in inv.get("parameters", [])}

                    logger.info(
                        "bedrock_tool_call",
                        action_group=action_group,
                        api_path=api_path,
                    )

                    result = "Tool execution failed"
                    if tool_executor:
                        try:
                            # Map Bedrock call to local tool execution
                            tool_name = f"{action_group}_{api_path}".lower().replace("/", "_")
                            res_data = await tool_executor(tool_name, params)
                            result = str(res_data)
                        except Exception as e:
                            result = f"Error: {e}"

                    tool_results.append(
                        {
                            "actionGroupInvocationOutput": {
                                "actionGroupName": action_group,
                                "apiPath": api_path,
                                "text": result,
                            },
                        },
                    )

            # Update session state with tool results for the next iteration
            session_state["returnControlInvocationResults"] = tool_results
            current_input = ""  # Continue without new user input

        return _clean_agent_response(full_response)

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

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