"""
API client for fetching mocked tool responses from Snowglobe servers.
"""

from __future__ import annotations

import asyncio
import json
import logging
import os
import threading
import time
from typing import Any, Dict, List, Optional

import httpx
from snowglobe.http_retry import (
    DEFAULT_HTTP_RETRY_ATTEMPTS,
    DEFAULT_HTTP_RETRY_BASE_SECONDS,
    DEFAULT_HTTP_RETRY_MAX_SECONDS,
    HttpRetryPolicy,
    async_request_with_retries,
    build_retry_policy,
    sync_request_with_retries,
)

from snowglobe.client.src.config import build_auth_headers, config

from .models import (
    ToolCall,
    ToolExecutionCreated,
    ToolExecutionRecord,
    ToolExecutionsBatchRequest,
    ToolExecutionsBatchResponse,
    ToolSpec,
)
from .registry import get_registered_tool_specs

LOGGER = logging.getLogger(__name__)

DEFAULT_POLL_INTERVAL_SEC = 1.0
DEFAULT_POLL_TIMEOUT_SEC = 300.0


class ToolExecutionValidationError(RuntimeError):
    """Raised when a tool execution fails validation."""


def _get_poll_timeout_sec() -> float:
    value = os.getenv("SNOWGLOBE_TOOL_POLL_TIMEOUT_SECONDS")
    if not value:
        return DEFAULT_POLL_TIMEOUT_SEC
    try:
        return float(value)
    except ValueError:
        LOGGER.warning(
            "Invalid SNOWGLOBE_TOOL_POLL_TIMEOUT_SECONDS=%r, using %s",
            value,
            DEFAULT_POLL_TIMEOUT_SEC,
        )
        return DEFAULT_POLL_TIMEOUT_SEC


POLL_TIMEOUT_SEC = _get_poll_timeout_sec()


def _get_poll_interval_sec() -> float:
    value = os.getenv("SNOWGLOBE_TOOL_POLL_INTERVAL_SECONDS")
    if not value:
        return DEFAULT_POLL_INTERVAL_SEC
    try:
        return max(0.1, float(value))
    except ValueError:
        LOGGER.warning(
            "Invalid SNOWGLOBE_TOOL_POLL_INTERVAL_SECONDS=%r, using %s",
            value,
            DEFAULT_POLL_INTERVAL_SEC,
        )
        return DEFAULT_POLL_INTERVAL_SEC


POLL_INTERVAL_SEC = _get_poll_interval_sec()
HTTP_RETRY_POLICY: HttpRetryPolicy = build_retry_policy(
    logger=LOGGER,
    attempts_env_names=("SNOWGLOBE_HTTP_RETRY_ATTEMPTS",),
    base_seconds_env_names=("SNOWGLOBE_HTTP_RETRY_BASE_SECONDS",),
    max_seconds_env_names=("SNOWGLOBE_HTTP_RETRY_MAX_SECONDS",),
    attempts_default=DEFAULT_HTTP_RETRY_ATTEMPTS,
    base_seconds_default=DEFAULT_HTTP_RETRY_BASE_SECONDS,
    max_seconds_default=DEFAULT_HTTP_RETRY_MAX_SECONDS,
)


async def _async_request_with_retries(
    client: httpx.AsyncClient,
    method: str,
    url: str,
    **kwargs,
) -> httpx.Response:
    return await async_request_with_retries(
        client,
        method,
        url,
        policy=HTTP_RETRY_POLICY,
        logger=LOGGER,
        **kwargs,
    )


def _sync_request_with_retries(
    client: httpx.Client,
    method: str,
    url: str,
    **kwargs,
) -> httpx.Response:
    return sync_request_with_retries(
        client,
        method,
        url,
        policy=HTTP_RETRY_POLICY,
        logger=LOGGER,
        **kwargs,
    )


def _build_batch_request_payload(
    test_id: str,
    tool_calls: List[ToolCall],
    tool_specs: Optional[List[ToolSpec]],
) -> Dict[str, Any]:
    request = ToolExecutionsBatchRequest(
        test_id=test_id,
        tool_calls=tool_calls,
        tool_specs=tool_specs,
    )
    return request.model_dump(exclude_none=True)


def _parse_batch_response(response: httpx.Response) -> List[ToolExecutionCreated]:
    response_data = response.json()
    batch_response = ToolExecutionsBatchResponse(**response_data)
    return batch_response.tool_executions


def _response_snippet(response: httpx.Response, limit: int = 500) -> str:
    body = (response.text or "").strip()
    if not body:
        return "<empty>"
    if len(body) > limit:
        return f"{body[:limit]}...<truncated>"
    return body


def _log_tool_execution_batch_http_error(
    *,
    exc: httpx.HTTPStatusError,
    test_id: str,
    included_tool_specs: bool,
) -> None:
    status = exc.response.status_code
    snippet = _response_snippet(exc.response)
    LOGGER.error(
        (
            "Error creating tool execution batch for test %s: status=%s, "
            "included_tool_specs=%s, response_body=%s"
        ),
        test_id,
        status,
        included_tool_specs,
        snippet,
    )


def _summarize_tool_executions(rows: List[ToolExecutionRecord]) -> str:
    payload = [
        {
            "id": row.id,
            "tool_name": row.tool_name,
            "validation_status": row.validation_status,
            "validation_errors": row.validation_errors,
            "tool_response_present": row.tool_response is not None,
            "updated_at": row.updated_at,
        }
        for row in rows
    ]
    return json.dumps(payload, ensure_ascii=True)


def _find_validation_failures(
    rows: List[ToolExecutionRecord],
) -> List[ToolExecutionRecord]:
    failures = []
    for row in rows:
        status = (row.validation_status or "").lower()
        if row.validation_errors or status in {"failed", "error", "invalid"}:
            failures.append(row)
    return failures


class _AsyncToolBatcher:
    def __init__(self, test_id: str, application_id: str) -> None:
        self._test_id = test_id
        self._application_id = application_id
        self._pending: List[tuple[str, Dict[str, Any], asyncio.Future]] = []
        self._lock = asyncio.Lock()
        self._flush_task: Optional[asyncio.Task] = None

    async def enqueue(self, tool_name: str, arguments: Dict[str, Any]) -> Optional[str]:
        loop = asyncio.get_running_loop()
        future: asyncio.Future = loop.create_future()
        async with self._lock:
            self._pending.append((tool_name, arguments, future))
            if self._flush_task is None:
                self._flush_task = loop.create_task(self._flush())

        return await future

    async def _drain_pending(
        self,
    ) -> List[tuple[str, Dict[str, Any], asyncio.Future]]:
        # Yield once so tool calls in the same event-loop tick can share one batch.
        await asyncio.sleep(0)
        async with self._lock:
            pending = self._pending
            self._pending = []
            self._flush_task = None
        return pending

    async def _fetch_batch_responses(
        self,
        pending: List[tuple[str, Dict[str, Any], asyncio.Future]],
    ) -> List[Optional[str]]:
        async with httpx.AsyncClient() as client:
            tool_calls = [
                ToolCall(tool_name=tool_name, arguments=arguments)
                for tool_name, arguments, _ in pending
            ]
            created = await create_tool_executions_batch_async(
                self._test_id,
                self._application_id,
                tool_calls,
                client=client,
            )
            if created is None:
                return [None] * len(pending)

            if len(created) != len(pending):
                raise RuntimeError(
                    "Tool execution batch size mismatch while creating executions."
                )

            ids = [item.id for item in created]
            response_map = await poll_tool_executions_async(ids, client=client)
            responses: List[Optional[str]] = []
            for execution_id in ids:
                response = response_map.get(execution_id)
                if response is None:
                    raise RuntimeError(
                        f"Missing tool response for execution {execution_id}"
                    )
                responses.append(response)
            return responses

    def _resolve_pending_results(
        self,
        pending: List[tuple[str, Dict[str, Any], asyncio.Future]],
        responses: List[Optional[str]],
    ) -> None:
        for (_, _, future), response in zip(pending, responses):
            if not future.done():
                future.set_result(response)

    def _resolve_pending_exception(
        self,
        pending: List[tuple[str, Dict[str, Any], asyncio.Future]],
        exc: Exception,
    ) -> None:
        for _, _, future in pending:
            if not future.done():
                future.set_exception(exc)

    async def _flush(self) -> None:
        pending = await self._drain_pending()
        if not pending:
            return

        try:
            responses = await self._fetch_batch_responses(pending)
            self._resolve_pending_results(pending, responses)
        except Exception as exc:
            self._resolve_pending_exception(pending, exc)


_batchers = threading.local()


def _get_async_batcher(test_id: str, application_id: str) -> _AsyncToolBatcher:
    if not hasattr(_batchers, "async_by_test"):
        _batchers.async_by_test = {}
    batcher = _batchers.async_by_test.get(test_id)
    if batcher is None:
        batcher = _AsyncToolBatcher(test_id, application_id)
        _batchers.async_by_test[test_id] = batcher
    return batcher


async def create_tool_executions_batch_async(
    test_id: str,
    application_id: str,
    tool_calls: List[ToolCall],
    client: Optional[httpx.AsyncClient] = None,
) -> Optional[List[ToolExecutionCreated]]:
    raw_specs = get_registered_tool_specs(application_id)
    tool_specs = [ToolSpec(**s) for s in raw_specs] if raw_specs else None
    url = f"{config.CONTROL_PLANE_URL}/api/tool-executions"
    payload = _build_batch_request_payload(test_id, tool_calls, tool_specs)
    LOGGER.debug("POST %s payload=%s", url, payload)

    async def _post(payload_to_send: Dict[str, Any]) -> httpx.Response:
        if client is None:
            async with httpx.AsyncClient() as async_client:
                return await _async_request_with_retries(
                    async_client,
                    "POST",
                    url,
                    json=payload_to_send,
                    headers=build_auth_headers(),
                    timeout=180,
                )
        return await _async_request_with_retries(
            client,
            "POST",
            url,
            json=payload_to_send,
            headers=build_auth_headers(),
            timeout=180,
        )

    try:
        response = await _post(payload)
        response.raise_for_status()
        return _parse_batch_response(response)
    except httpx.HTTPStatusError as exc:
        if exc.response.status_code == 404:
            snippet = _response_snippet(exc.response)
            LOGGER.warning(
                "Tool execution batch returned 404 for test %s: %s",
                test_id,
                snippet,
            )
            return None
        if tool_specs and exc.response.status_code >= 500:
            LOGGER.warning(
                (
                    "Received %s creating tool execution batch for test %s with "
                    "tool specs. Retrying once without tool_specs."
                ),
                exc.response.status_code,
                test_id,
            )
            retry_payload = _build_batch_request_payload(test_id, tool_calls, None)
            try:
                retry_response = await _post(retry_payload)
                retry_response.raise_for_status()
                LOGGER.info(
                    (
                        "Recovered tool execution batch creation for test %s by "
                        "retrying without tool_specs."
                    ),
                    test_id,
                )
                return _parse_batch_response(retry_response)
            except httpx.HTTPStatusError as retry_exc:
                if retry_exc.response.status_code == 404:
                    snippet = _response_snippet(retry_exc.response)
                    LOGGER.warning(
                        "Tool execution batch returned 404 for test %s: %s",
                        test_id,
                        snippet,
                    )
                    return None
                _log_tool_execution_batch_http_error(
                    exc=retry_exc,
                    test_id=test_id,
                    included_tool_specs=False,
                )
                raise

        _log_tool_execution_batch_http_error(
            exc=exc,
            test_id=test_id,
            included_tool_specs=bool(tool_specs),
        )
        raise


def create_tool_executions_batch_sync(
    test_id: str, tool_calls: List[ToolCall], application_id: str
) -> Optional[List[ToolExecutionCreated]]:
    raw_specs = get_registered_tool_specs(application_id)
    tool_specs = [ToolSpec(**s) for s in raw_specs] if raw_specs else None
    url = f"{config.CONTROL_PLANE_URL}/api/tool-executions"
    payload = _build_batch_request_payload(test_id, tool_calls, tool_specs)
    LOGGER.debug("POST %s payload=%s", url, payload)

    def _post(payload_to_send: Dict[str, Any]) -> httpx.Response:
        with httpx.Client() as client:
            return _sync_request_with_retries(
                client,
                "POST",
                url,
                json=payload_to_send,
                headers=build_auth_headers(),
                timeout=180,
            )

    try:
        response = _post(payload)
        response.raise_for_status()
        return _parse_batch_response(response)
    except httpx.HTTPStatusError as exc:
        if exc.response.status_code == 404:
            snippet = _response_snippet(exc.response)
            LOGGER.warning(
                "Tool execution batch returned 404 for test %s: %s",
                test_id,
                snippet,
            )
            return None
        if tool_specs and exc.response.status_code >= 500:
            LOGGER.warning(
                (
                    "Received %s creating tool execution batch for test %s with "
                    "tool specs. Retrying once without tool_specs."
                ),
                exc.response.status_code,
                test_id,
            )
            retry_payload = _build_batch_request_payload(test_id, tool_calls, None)
            try:
                retry_response = _post(retry_payload)
                retry_response.raise_for_status()
                LOGGER.info(
                    (
                        "Recovered tool execution batch creation for test %s by "
                        "retrying without tool_specs."
                    ),
                    test_id,
                )
                return _parse_batch_response(retry_response)
            except httpx.HTTPStatusError as retry_exc:
                if retry_exc.response.status_code == 404:
                    snippet = _response_snippet(retry_exc.response)
                    LOGGER.warning(
                        "Tool execution batch returned 404 for test %s: %s",
                        test_id,
                        snippet,
                    )
                    return None
                _log_tool_execution_batch_http_error(
                    exc=retry_exc,
                    test_id=test_id,
                    included_tool_specs=False,
                )
                raise

        _log_tool_execution_batch_http_error(
            exc=exc,
            test_id=test_id,
            included_tool_specs=bool(tool_specs),
        )
        raise


async def get_tool_executions_by_ids_async(
    ids: List[str],
    client: Optional[httpx.AsyncClient] = None,
) -> List[ToolExecutionRecord]:
    url = f"{config.CONTROL_PLANE_URL}/api/tool-executions"

    try:
        if client is None:
            async with httpx.AsyncClient() as async_client:
                response = await _async_request_with_retries(
                    async_client,
                    "GET",
                    url,
                    params={"ids": ",".join(ids)},
                    headers=build_auth_headers(),
                    timeout=180,
                )
        else:
            response = await _async_request_with_retries(
                client,
                "GET",
                url,
                params={"ids": ",".join(ids)},
                headers=build_auth_headers(),
                timeout=180,
            )
        response.raise_for_status()

        response_data = response.json()
        return [ToolExecutionRecord(**row) for row in response_data]
    except httpx.HTTPStatusError:
        raise


def get_tool_executions_by_ids_sync(ids: List[str]) -> List[ToolExecutionRecord]:
    url = f"{config.CONTROL_PLANE_URL}/api/tool-executions"

    try:
        with httpx.Client() as client:
            response = _sync_request_with_retries(
                client,
                "GET",
                url,
                params={"ids": ",".join(ids)},
                headers=build_auth_headers(),
                timeout=180,
            )
            response.raise_for_status()

            response_data = response.json()
            return [ToolExecutionRecord(**row) for row in response_data]
    except httpx.HTTPStatusError:
        raise


async def poll_tool_executions_async(
    ids: List[str],
    client: Optional[httpx.AsyncClient] = None,
) -> Dict[str, str]:
    start_time = time.monotonic()
    last_rows: List[ToolExecutionRecord] = []
    while time.monotonic() - start_time < POLL_TIMEOUT_SEC:
        rows = await get_tool_executions_by_ids_async(ids, client=client)
        last_rows = rows
        failures = _find_validation_failures(rows)
        if failures:
            details = _summarize_tool_executions(failures)
            message = f"Tool execution validation failed: {details}"
            LOGGER.error(message)
            raise ToolExecutionValidationError(message)
        response_map = {
            row.id: row.tool_response for row in rows if row.tool_response is not None
        }
        if len(response_map) == len(ids):
            return response_map
        await asyncio.sleep(POLL_INTERVAL_SEC)

    last_summary = _summarize_tool_executions(last_rows) if last_rows else "[]"
    elapsed = time.monotonic() - start_time
    message = (
        "Tool execution timed out waiting for response. "
        f"execution_ids={ids}, elapsed_seconds={elapsed:.1f}, "
        f"poll_timeout_seconds={POLL_TIMEOUT_SEC}. "
        f"Last records: {last_summary}"
    )
    LOGGER.error(message)
    raise TimeoutError(message)


def poll_tool_executions_sync(ids: List[str]) -> Dict[str, str]:
    start_time = time.monotonic()
    last_rows: List[ToolExecutionRecord] = []
    while time.monotonic() - start_time < POLL_TIMEOUT_SEC:
        rows = get_tool_executions_by_ids_sync(ids)
        last_rows = rows
        failures = _find_validation_failures(rows)
        if failures:
            details = _summarize_tool_executions(failures)
            message = f"Tool execution validation failed: {details}"
            LOGGER.error(message)
            raise ToolExecutionValidationError(message)
        response_map = {
            row.id: row.tool_response for row in rows if row.tool_response is not None
        }
        if len(response_map) == len(ids):
            return response_map
        time.sleep(POLL_INTERVAL_SEC)

    last_summary = _summarize_tool_executions(last_rows) if last_rows else "[]"
    elapsed = time.monotonic() - start_time
    message = (
        "Tool execution timed out waiting for response. "
        f"execution_ids={ids}, elapsed_seconds={elapsed:.1f}, "
        f"poll_timeout_seconds={POLL_TIMEOUT_SEC}. "
        f"Last records: {last_summary}"
    )
    LOGGER.error(message)
    raise TimeoutError(message)


async def fetch_mocked_tool_response_async(
    tool_name: str,
    arguments: Dict[str, Any],
    test_id: str,
    application_id: str,
) -> Optional[str]:
    """
    Fetch a mocked tool response from Snowglobe servers (async).

    Args:
        tool_name: Name of the tool being called
        arguments: Tool arguments
        test_id: Snowglobe test ID

    Returns:
        Mocked tool response as JSON string

    Raises:
        httpx.HTTPStatusError: If the API request fails
    """
    batcher = _get_async_batcher(test_id, application_id)
    return await batcher.enqueue(tool_name, arguments)


def fetch_mocked_tool_response_sync(
    tool_name: str,
    arguments: Dict[str, Any],
    test_id: str,
    application_id: str,
) -> Optional[str]:
    """
    Fetch a mocked tool response from Snowglobe servers (sync).

    This is a synchronous wrapper around the async version.
    Uses httpx.Client for synchronous HTTP requests.

    Args:
        tool_name: Name of the tool being called
        arguments: Tool arguments
        test_id: Snowglobe test ID

    Returns:
        Mocked tool response as JSON string

    Raises:
        httpx.HTTPStatusError: If the API request fails
    """
    tool_calls = [ToolCall(tool_name=tool_name, arguments=arguments)]
    created = create_tool_executions_batch_sync(test_id, tool_calls, application_id)
    if created is None:
        return None

    if len(created) != 1:
        raise RuntimeError("Unexpected tool execution batch size.")

    response_map = poll_tool_executions_sync([created[0].id])
    return response_map.get(created[0].id)
