"""
Agent client for StackSpot AI.

This module provides a synchronous client for interacting with StackSpot AI Agents,
supporting single message requests and conversation context.
"""

import logging
from concurrent.futures import ThreadPoolExecutor, as_completed
from dataclasses import dataclass

import requests

from stkai._config import AgentConfig
from stkai._http import HttpClient
from stkai._retry import Retrying
from stkai.agents._conversation import ConversationContext, ConversationScope
from stkai.agents._handlers import ChatResultHandler, ChatResultHandlerError
from stkai.agents._models import ChatRequest, ChatResponse, ChatStatus
from stkai.agents._stream import ChatResponseStream

logger = logging.getLogger(__name__)


@dataclass(frozen=True)
class AgentOptions:
    """
    Configuration options for the Agent client.

    Fields set to None will use values from global config (STKAI.config.agent).

    Attributes:
        request_timeout: HTTP request timeout in seconds.
        retry_max_retries: Maximum number of retry attempts for failed chat calls.
            Use 0 to disable retries (single attempt only).
            Use 3 for 4 total attempts (1 original + 3 retries).
        retry_initial_delay: Initial delay in seconds for the first retry attempt.
            Subsequent retries use exponential backoff (delay doubles each attempt).
        max_workers: Maximum number of threads for batch execution (chat_many).

    Example:
        >>> # Use all defaults from config
        >>> agent = Agent(agent_id="my-agent")
        >>>
        >>> # Customize timeout and enable retry
        >>> options = AgentOptions(request_timeout=120, retry_max_retries=3)
        >>> agent = Agent(agent_id="my-agent", options=options)
    """
    request_timeout: int | None = None
    retry_max_retries: int | None = None
    retry_initial_delay: float | None = None
    max_workers: int | None = None

    def with_defaults_from(self, cfg: AgentConfig) -> "AgentOptions":
        """
        Returns a new AgentOptions with None values filled from config.

        User-provided values take precedence; None values use config defaults.
        This follows the Single Source of Truth principle where STKAI.config
        is the authoritative source for default values.

        Args:
            cfg: The AgentConfig to use for default values.

        Returns:
            A new AgentOptions with all fields resolved (no None values).

        Example:
            >>> options = AgentOptions(request_timeout=120)
            >>> resolved = options.with_defaults_from(STKAI.config.agent)
            >>> resolved.request_timeout  # 120 (user-defined)
        """
        return AgentOptions(
            request_timeout=self.request_timeout if self.request_timeout is not None else cfg.request_timeout,
            retry_max_retries=self.retry_max_retries if self.retry_max_retries is not None else cfg.retry_max_retries,
            retry_initial_delay=self.retry_initial_delay if self.retry_initial_delay is not None else cfg.retry_initial_delay,
            max_workers=self.max_workers if self.max_workers is not None else cfg.max_workers,
        )


class Agent:
    """
    Synchronous client for interacting with StackSpot AI Agents.

    This client provides a high-level interface for sending messages to Agents
    and receiving responses, with support for:

    - Single message requests (blocking)
    - Conversation context for multi-turn interactions
    - Knowledge source integration
    - Token usage tracking

    Example:
        >>> from stkai.agents import Agent, ChatRequest
        >>> agent = Agent(agent_id="my-agent-slug")
        >>> response = agent.chat(
        ...     request=ChatRequest(user_prompt="What is SOLID?")
        ... )
        >>> if response.is_success():
        ...     print(response.result)

    Attributes:
        agent_id: The Agent ID (slug) to interact with.
        base_url: The base URL for the StackSpot AI API.
        options: Configuration options for the client.
        http_client: HTTP client for API calls (default: EnvironmentAwareHttpClient).
    """

    def __init__(
        self,
        agent_id: str,
        base_url: str | None = None,
        options: AgentOptions | None = None,
        http_client: HttpClient | None = None,
    ):
        """
        Initialize the Agent client.

        Args:
            agent_id: The Agent ID (slug) to interact with.
            base_url: Base URL for the StackSpot AI API.
                If None, uses global config (STKAI.config.agent.base_url).
            options: Configuration options for the client.
                If None, uses defaults from global config (STKAI.config.agent).
                Partial options are merged with config defaults via with_defaults_from().
            http_client: Custom HTTP client implementation for API calls.
                If None, uses EnvironmentAwareHttpClient (auto-detects CLI or standalone).

        Raises:
            AssertionError: If agent_id is empty.
        """
        # Get global config for defaults
        from stkai._config import STKAI
        cfg = STKAI.config.agent

        # Resolve options with defaults from config (Single Source of Truth)
        resolved_options = (options or AgentOptions()).with_defaults_from(cfg)

        # Resolve base_url
        if base_url is None:
            base_url = cfg.base_url

        if not http_client:
            from stkai._http import EnvironmentAwareHttpClient
            http_client = EnvironmentAwareHttpClient()

        # Validations
        assert agent_id, "Agent ID cannot be empty."
        assert base_url, "Agent base_url cannot be empty."
        assert http_client is not None, "Agent http_client cannot be None."

        assert resolved_options.max_workers is not None, "Thread-pool max_workers can not be None."
        assert resolved_options.max_workers > 0, "Thread-pool max_workers must be greater than 0."

        self.agent_id = agent_id
        self.base_url = base_url.rstrip("/")
        self.options = resolved_options
        self.max_workers = resolved_options.max_workers
        self.executor = ThreadPoolExecutor(max_workers=self.max_workers)
        self.http_client: HttpClient = http_client

    def chat(
        self,
        request: ChatRequest,
        result_handler: ChatResultHandler | None = None,
    ) -> ChatResponse:
        """
        Send a message to the Agent and wait for the response (blocking).

        This method sends a user prompt to the Agent and blocks until
        a response is received or an error occurs.

        If retry is configured (retry_max_retries > 0), automatically retries on:
        - HTTP 5xx errors (500, 502, 503, 504)
        - Network errors (Timeout, ConnectionError)

        Does NOT retry on:
        - HTTP 4xx errors (client errors)

        Note:
            retry_max_retries=0 means 1 attempt (no retry).
            retry_max_retries=3 means 4 attempts (1 original + 3 retries).

        Args:
            request: The request containing the user prompt and options.
            result_handler: Optional handler to process the response message.
                If None, uses RawResultHandler (returns message as-is).
                Use JSON_RESULT_HANDLER to parse JSON responses.

        Returns:
            ChatResponse with the Agent's reply or error information.
            The `result` field contains the processed result from the handler.

        Raises:
            ChatResultHandlerError: If the result handler fails to process the response.

        Example:
            >>> # Single message (default RawResultHandler)
            >>> response = agent.chat(
            ...     request=ChatRequest(user_prompt="Hello!")
            ... )
            >>> print(response.result)  # Same as response.raw_result
            >>>
            >>> # Parse JSON response
            >>> from stkai.agents import JSON_RESULT_HANDLER
            >>> response = agent.chat(request, result_handler=JSON_RESULT_HANDLER)
            >>> print(response.result)  # Parsed dict
            >>>
            >>> # With conversation context
            >>> resp1 = agent.chat(
            ...     request=ChatRequest(
            ...         user_prompt="What is Python?",
            ...         use_conversation=True
            ...     )
            ... )
            >>> resp2 = agent.chat(
            ...     request=ChatRequest(
            ...         user_prompt="What are its main features?",
            ...         conversation_id=resp1.conversation_id,
            ...         use_conversation=True
            ...     )
            ... )
        """
        logger.info(f"{request.id[:26]:<26} | Agent | 🛜 Starting chat with agent '{self.agent_id}'.")
        logger.info(f"{request.id[:26]:<26} | Agent |    ├ base_url={self.base_url}")
        logger.info(f"{request.id[:26]:<26} | Agent |    └ agent_id='{self.agent_id}'")

        response = self._do_chat(
            request=request,
            result_handler=result_handler
        )

        logger.info(f"{request.id[:26]:<26} | Agent | 🛜 Chat finished.")
        if response.is_success():
            logger.info(f"{request.id[:26]:<26} | Agent |    └ with status = {response.status}")
        else:
            logger.info(f"{request.id[:26]:<26} | Agent |    ├ with status = {response.status}")
            logger.info(f"{request.id[:26]:<26} | Agent |    └ with error message = \"{response.error}\"")

        assert response.request is request, \
            "🌀 Sanity check | Unexpected mismatch: response does not reference its corresponding request."
        return response

    def chat_many(
        self,
        request_list: list[ChatRequest],
        result_handler: ChatResultHandler | None = None,
    ) -> list[ChatResponse]:
        """
        Send multiple chat messages concurrently, wait for all responses (blocking),
        and return them in the same order as `request_list`.

        Each request is executed in parallel threads using the internal thread-pool.
        Returns a list of ChatResponse objects in the same order as `request_list`.

        Args:
            request_list: List of ChatRequest objects to send.
            result_handler: Optional handler to process the response message.
                If None, uses RawResultHandler (returns message as-is).

        Returns:
            List[ChatResponse]: One response per request, in the same order.

        Example:
            >>> requests = [
            ...     ChatRequest(user_prompt="What is Python?"),
            ...     ChatRequest(user_prompt="What is Java?"),
            ... ]
            >>> responses = agent.chat_many(requests)
            >>> for resp in responses:
            ...     print(resp.result)
        """
        if not request_list:
            return []

        logger.info(
            f"{'Agent-Batch-Execution'[:26]:<26} | Agent | "
            f"🛜 Starting batch execution of {len(request_list)} requests."
        )
        logger.info(f"{'Agent-Batch-Execution'[:26]:<26} | Agent |    ├ base_url={self.base_url}")
        logger.info(f"{'Agent-Batch-Execution'[:26]:<26} | Agent |    ├ agent_id='{self.agent_id}'")
        logger.info(f"{'Agent-Batch-Execution'[:26]:<26} | Agent |    └ max_concurrent={self.max_workers}")

        # Warn about race condition: chat_many inside UseConversation without a pre-set conversation_id
        if len(request_list) > 1:
            conv_ctx = ConversationScope.get_current()
            if conv_ctx and not conv_ctx.has_conversation_id():
                logger.warning(
                    f"{'Agent-Batch-Execution'[:26]:<26} | Agent | "
                    "⚠️ chat_many() called inside UseConversation without a pre-set conversation_id. "
                    "Concurrent requests will race to capture the server-assigned ID, "
                    "likely starting independent conversations. "
                    "Consider using UseConversation.with_generated_id() or passing an explicit conversation_id."
                )

        # Use thread-pool for parallel calls to `_do_chat`
        # The ConversationScope::propagate method captures the active UseConversation context (if any)
        # and installs it in each worker thread — only conversation state is propagated.
        future_to_index = {
            self.executor.submit(
                ConversationScope.propagate(self._do_chat), # propagate conversation context
                request=req,                                # arg-1
                result_handler=result_handler,              # arg-2
            ): idx
            for idx, req in enumerate(request_list)
        }

        # Block and wait for all responses to be finished
        responses_map: dict[int, ChatResponse] = {}

        for future in as_completed(future_to_index):
            idx = future_to_index[future]
            correlated_request = request_list[idx]
            try:
                responses_map[idx] = future.result()
            except Exception as e:
                logger.error(
                    f"{correlated_request.id[:26]:<26} | Agent | ❌ Chat failed in batch(seq={idx}). {e}",
                    exc_info=logger.isEnabledFor(logging.DEBUG)
                )
                responses_map[idx] = ChatResponse(
                    request=correlated_request,
                    status=ChatStatus.ERROR,
                    error=str(e),
                )

        # Rebuild responses list in the same order of requests list
        responses = [
            responses_map[i] for i in range(len(request_list))
        ]

        # Race-condition check: ensure both lists have the same length
        assert len(responses) == len(request_list), (
            f"🌀 Sanity check | Unexpected mismatch: responses(size={len(responses)}) is different from requests(size={len(request_list)})."
        )
        # Race-condition check: ensure each response points to its respective request
        assert all(resp.request is req for req, resp in zip(request_list, responses, strict=True)), (
            "🌀 Sanity check | Unexpected mismatch: some responses do not reference their corresponding requests."
        )

        logger.info(
            f"{'Agent-Batch-Execution'[:26]:<26} | Agent | 🛜 Batch execution finished."
        )
        logger.info(f"{'Agent-Batch-Execution'[:26]:<26} | Agent |    ├ total of responses = {len(responses)}")

        from collections import Counter
        totals_per_status = Counter(r.status for r in responses)
        items = totals_per_status.items()
        for idx, (status, total) in enumerate(items):
            icon = "└" if idx == (len(items) - 1) else "├"
            logger.info(f"{'Agent-Batch-Execution'[:26]:<26} | Agent |    {icon} total of responses with status {status:<7} = {total}")

        return responses

    def chat_stream(
        self,
        request: ChatRequest,
        result_handler: ChatResultHandler | None = None,
    ) -> ChatResponseStream:
        """
        Send a message to the Agent and return a streaming response.

        Returns a ``ChatResponseStream`` context manager that yields SSE events
        as they arrive from the server. Must be used with ``with``::

            with agent.chat_stream(ChatRequest(user_prompt="Hello")) as stream:
                for event in stream:
                    if event.is_delta:
                        print(event.text, end="", flush=True)
                print(f"\\nTokens: {stream.response.tokens.total}")

        If retry is configured, retries only the initial connection (before
        the stream begins). Mid-stream errors are not retried.

        ``chat_many`` + stream is **not supported** — streaming is real-time
        by nature and batch execution defeats its purpose.

        Args:
            request: The request containing the user prompt and options.
            result_handler: Optional handler to process the final accumulated text.
                If None, uses RawResultHandler (returns accumulated text as-is).
                The handler is applied once after the stream is fully consumed,
                over the complete accumulated text — not per chunk.

        Returns:
            A ChatResponseStream context manager for iterating SSE events.

        Raises:
            requests.HTTPError: If the initial HTTP request fails (after retries).
            RuntimeError: If the HTTP client does not support streaming.

        Example:
            >>> with agent.chat_stream(ChatRequest(user_prompt="Hello")) as stream:
            ...     for text in stream.text_stream:
            ...         print(text, end="", flush=True)
            >>>
            >>> # With JSON result handler
            >>> from stkai.agents import JSON_RESULT_HANDLER
            >>> with agent.chat_stream(request, result_handler=JSON_RESULT_HANDLER) as stream:
            ...     response = stream.get_final_response()
            ...     print(response.result)  # Parsed dict
        """
        logger.info(f"{request.id[:26]:<26} | Agent | 🛜 Starting streaming chat with agent '{self.agent_id}'.")
        logger.info(f"{request.id[:26]:<26} | Agent |    ├ base_url={self.base_url}")
        logger.info(f"{request.id[:26]:<26} | Agent |    └ agent_id='{self.agent_id}'")

        # Assertion for type narrowing (mypy)
        assert self.options.request_timeout is not None, \
            "🌀 Sanity check | request_timeout must be set after with_defaults_from()"
        assert self.options.retry_max_retries is not None, \
            "🌀 Sanity check | retry_max_retries must be set after with_defaults_from()"
        assert self.options.retry_initial_delay is not None, \
            "🌀 Sanity check | retry_initial_delay must be set after with_defaults_from()"

        # Build payload with streaming=True
        payload = request.to_api_payload()
        payload["streaming"] = True

        # Apply UseConversation context (if active and request has no explicit conversation_id)
        conv_ctx = ConversationScope.get_current()
        if conv_ctx is not None and not request.conversation_id:
            payload["use_conversation"] = True
            if conv_ctx.conversation_id:
                payload["conversation_id"] = conv_ctx.conversation_id

        url = f"{self.base_url}/v1/agent/{self.agent_id}/chat"

        # Retry only the initial connection
        for attempt in Retrying(
            max_retries=self.options.retry_max_retries,
            initial_delay=self.options.retry_initial_delay,
            logger_prefix=f"{request.id[:26]:<26} | Agent",
        ):
            with attempt:
                logger.info(
                    f"{request.id[:26]:<26} | Agent | "
                    f"Opening stream to agent '{self.agent_id}' (attempt {attempt.attempt_number}/{attempt.max_attempts})..."
                )

                http_response = self.http_client.post_stream(
                    url=url,
                    data=payload,
                    timeout=self.options.request_timeout,
                )
                assert isinstance(http_response, requests.Response), \
                    f"🌀 Sanity check | Object returned by `post_stream` is not an instance of `requests.Response`. ({http_response.__class__})"

                http_response.raise_for_status()

                logger.info(f"{request.id[:26]:<26} | Agent | 🛜 Stream opened successfully.")

                stream = ChatResponseStream(
                    request=request,
                    http_response=http_response,
                    result_handler=result_handler,
                )

                # Hook: after stream is fully consumed, update UseConversation
                if conv_ctx is not None:
                    stream._build_response = self._make_conversation_tracking_hook(  # type: ignore[assignment,method-assign]
                        stream=stream,
                        conv_ctx=conv_ctx,
                    )

                return stream

        # Should never reach here - Retrying raises MaxRetriesExceededError
        raise RuntimeError(
            "Unexpected error while opening stream: "
            "reached end of `chat_stream` method without returning a ChatResponseStream."
        )

    @staticmethod
    def _make_conversation_tracking_hook(
        stream: ChatResponseStream,
        conv_ctx: ConversationContext,
    ) -> object:
        """Create a _build_response hook that tracks conversation_id in UseConversation."""
        original_build = stream._build_response

        def _build_with_conversation_tracking() -> None:
            original_build()
            if stream._response is not None:
                conv_id = stream._response.conversation_id
                if conv_id:
                    conv_ctx.update_if_absent(conversation_id=conv_id)

        return _build_with_conversation_tracking

    def _do_chat(
        self,
        request: ChatRequest,
        result_handler: ChatResultHandler | None = None,
    ) -> ChatResponse:
        """
        Internal method that executes the full chat workflow: retry, HTTP request,
        response parsing, and error handling.

        Always returns a ChatResponse (never raises exceptions).
        Called by both chat() and chat_many().

        Args:
            request: The request containing the user prompt and options.
            result_handler: Optional handler to process the response message.

        Returns:
            ChatResponse with the Agent's reply or error information.
        """
        assert request, "🌀 Sanity check | Chat-Request can not be None."
        assert request.id, "🌀 Sanity check | Chat-Request ID can not be None."

        # Assertion for type narrowing (mypy)
        assert self.options.request_timeout is not None, \
            "🌀 Sanity check | request_timeout must be set after with_defaults_from()"
        assert self.options.retry_max_retries is not None, \
            "🌀 Sanity check | retry_max_retries must be set after with_defaults_from()"
        assert self.options.retry_initial_delay is not None, \
            "🌀 Sanity check | retry_initial_delay must be set after with_defaults_from()"

        try:
            for attempt in Retrying(
                max_retries=self.options.retry_max_retries,
                initial_delay=self.options.retry_initial_delay,
                logger_prefix=f"{request.id[:26]:<26} | Agent",
            ):
                with attempt:
                    logger.info(
                        f"{request.id[:26]:<26} | Agent | "
                        f"Sending message to agent '{self.agent_id}' (attempt {attempt.attempt_number}/{attempt.max_attempts})..."
                    )

                    # HTTP request
                    payload = request.to_api_payload()

                    # Apply UseConversation context (if active and request has no explicit conversation_id)
                    conv_ctx = ConversationScope.get_current()
                    if conv_ctx is not None and not request.conversation_id:
                        payload["use_conversation"] = True
                        if conv_ctx.conversation_id:
                            payload["conversation_id"] = conv_ctx.conversation_id

                    sent_conv_id = payload.get("conversation_id")
                    if sent_conv_id:
                        logger.debug(
                            f"{request.id[:26]:<26} | Agent | Request sent with conversation_id='{sent_conv_id}'"
                        )

                    url = f"{self.base_url}/v1/agent/{self.agent_id}/chat"
                    http_response = self.http_client.post(
                        url=url,
                        data=payload,
                        timeout=self.options.request_timeout,
                    )
                    assert isinstance(http_response, requests.Response), \
                        f"🌀 Sanity check | Object returned by `post` method is not an instance of `requests.Response`. ({http_response.__class__})"

                    http_response.raise_for_status()
                    response_data = http_response.json()
                    raw_message = response_data.get("message")

                    # Process result through handler
                    if not result_handler:
                        from stkai.agents._handlers import DEFAULT_RESULT_HANDLER
                        result_handler = DEFAULT_RESULT_HANDLER

                    try:
                        from stkai.agents._handlers import ChatResultContext
                        context = ChatResultContext(request=request, raw_result=raw_message)
                        processed_result = result_handler.handle_result(context)
                    except Exception as e:
                        handler_name = type(result_handler).__name__
                        raise ChatResultHandlerError(
                            f"{request.id} | Agent | Result handler '{handler_name}' failed: {e}",
                            cause=e, result_handler=result_handler,
                        ) from e

                    response = ChatResponse(
                        request=request,
                        status=ChatStatus.SUCCESS,
                        result=processed_result,
                        raw_response=response_data,
                    )

                    # Auto-track conversation_id in UseConversation context
                    if conv_ctx is not None and response.conversation_id:
                        conv_ctx.update_if_absent(conversation_id=response.conversation_id)

                    logger.info(
                        f"{request.id[:26]:<26} | Agent | "
                        f"✅ Response received successfully (tokens: {response.tokens.total if response.tokens else 'N/A'})"
                    )
                    if response.conversation_id:
                        logger.debug(
                            f"{request.id[:26]:<26} | Agent | Response received with conversation_id='{response.conversation_id}'"
                        )

                    assert response.request is request, \
                        "🌀 Sanity check | Unexpected mismatch: response does not reference its corresponding request."
                    return response

            # Should never reach here - Retrying raises MaxRetriesExceededError
            raise RuntimeError(
                "Unexpected error while chatting the agent: "
                "reached end of `_do_chat` method without returning a response."
            )

        except Exception as e:
            error_status = ChatStatus.from_exception(e)
            error_msg = f"Chat message failed: {e}"
            if isinstance(e, requests.HTTPError) and e.response is not None:
                error_msg = f"Chat message failed due to an HTTP error {e.response.status_code}: {e.response.text}"
            logger.error(
                f"{request.id[:26]:<26} | Agent | ❌ {error_msg}",
                exc_info=logger.isEnabledFor(logging.DEBUG)
            )
            return ChatResponse(
                request=request,
                status=error_status,
                error=error_msg,
            )
