"""
Boring-OS SDK client.

Main client class for interacting with the Boring-OS API.
Provides both synchronous and asynchronous methods.
"""

from __future__ import annotations

import uuid
from typing import TYPE_CHECKING, Any

import httpx

if TYPE_CHECKING:
    from boring_os.agents import AgentsClient
    from boring_os.flows import FlowsClient
    from boring_os.models_client import ModelsClient
    from boring_os.tracing import OpikTracer

from boring_os.exceptions import (
    AgentNotFoundError,
    AuthenticationError,
    AuthorizationError,
    BoringOSError,
    ConnectionError,
    FlowNotFoundError,
    RateLimitError,
    ToolDisabledError,
    ToolExecutionError,
    ToolNotFoundError,
    ToolTimeoutError,
    ValidationError,
)
from boring_os.models import (
    AsyncToolResponse,
    ResultStatus,
    SessionInfo,
    Tool,
    ToolInfo,
    ToolListResponse,
    ToolResult,
    ToolStatus,
    ToolType,
)


# Error code to exception mapping
ERROR_CODE_MAP = {
    "AUTHENTICATION_ERROR": AuthenticationError,
    "AUTHORIZATION_ERROR": AuthorizationError,
    "RATE_LIMIT_EXCEEDED": RateLimitError,
    "VALIDATION_ERROR": ValidationError,
    "TOOL_NOT_FOUND": ToolNotFoundError,
    "TOOL_DISABLED": ToolDisabledError,
    "TOOL_TIMEOUT": ToolTimeoutError,
    "TOOL_EXECUTION_ERROR": ToolExecutionError,
    "AGENT_NOT_FOUND": AgentNotFoundError,
    "FLOW_NOT_FOUND": FlowNotFoundError,
}


class ToolsClient:
    """
    Client for tool-related operations.
    
    Accessed via `client.tools` on the main BoringOS client.
    """
    
    def __init__(self, client: "BoringOS"):
        self._client = client
    
    def call(
        self,
        tool_name: str,
        timeout: int = 30,
        trace_id: str | None = None,
        **params: Any,
    ) -> dict[str, Any]:
        """
        Execute a tool synchronously (recommended for agents).
        
        Uses server-side waiting for minimal latency.
        
        Args:
            tool_name: Name of the tool to execute
            timeout: Max seconds to wait (1-60)
            trace_id: Optional trace ID for observability
            **params: Tool parameters as keyword arguments
        
        Returns:
            Tool execution result data
        
        Raises:
            ToolNotFoundError: If tool doesn't exist
            ToolTimeoutError: If execution times out
            ToolExecutionError: If execution fails
            ValidationError: If parameters are invalid
        
        Example:
            result = client.tools.call(
                "get_customer_profile",
                customer_id="cust_123"
            )
        """
        result = self._call_internal(tool_name, params, timeout, trace_id)
        return result.data
    
    def call_with_result(
        self,
        tool_name: str,
        timeout: int = 30,
        trace_id: str | None = None,
        **params: Any,
    ) -> ToolResult:
        """
        Execute a tool and return the full result object.
        
        Use this when you need access to metadata or deprecation warnings.
        
        Args:
            tool_name: Name of the tool to execute
            timeout: Max seconds to wait (1-60)
            trace_id: Optional trace ID for observability
            **params: Tool parameters as keyword arguments
        
        Returns:
            Full ToolResult object with metadata
        """
        return self._call_internal(tool_name, params, timeout, trace_id)
    
    def _call_internal(
        self,
        tool_name: str,
        params: dict[str, Any],
        timeout: int,
        trace_id: str | None,
    ) -> ToolResult:
        """Internal method to execute tool call."""
        # Generate trace_id if not provided
        if trace_id is None:
            trace_id = str(uuid.uuid4())
        
        # Use tracing if enabled
        if self._client._tracer:
            return self._traced_call(tool_name, params, timeout, trace_id)
        
        return self._execute_call(tool_name, params, timeout, trace_id)
    
    def _execute_call(
        self,
        tool_name: str,
        params: dict[str, Any],
        timeout: int,
        trace_id: str,
    ) -> ToolResult:
        """Execute the actual tool call without tracing."""
        request_body = {
            "tool_name": tool_name,
            "params": params,
            "options": {
                "timeout_seconds": timeout,
                "trace_id": trace_id,
            },
        }
        
        response = self._client._request(
            "POST",
            "/tools/execute",
            json=request_body,
            params={"timeout": timeout},
            timeout=timeout + 5,  # Add buffer for network
        )
        
        result = ToolResult.model_validate(response)
        
        # Raise exception if execution failed
        if result.status == ResultStatus.ERROR and result.error:
            self._raise_for_error(tool_name, result.error.code, result.error.message)
        
        return result
    
    def _traced_call(
        self,
        tool_name: str,
        params: dict[str, Any],
        timeout: int,
        trace_id: str,
    ) -> ToolResult:
        """Execute tool call with Opik tracing."""
        tracer = self._client._tracer
        assert tracer is not None  # Guaranteed by caller
        
        with tracer.trace_tool(tool_name, params, trace_id, timeout) as span:
            result = self._execute_call(tool_name, params, timeout, trace_id)
            span.set_result(result)
            return result
    
    async def call_async(
        self,
        tool_name: str,
        timeout: int = 30,
        trace_id: str | None = None,
        **params: Any,
    ) -> dict[str, Any]:
        """
        Execute a tool asynchronously.
        
        Async version of `call()` for use in async contexts.
        
        Args:
            tool_name: Name of the tool to execute
            timeout: Max seconds to wait (1-60)
            trace_id: Optional trace ID for observability
            **params: Tool parameters as keyword arguments
        
        Returns:
            Tool execution result data
        """
        result = await self._call_async_internal(tool_name, params, timeout, trace_id)
        return result.data
    
    async def _call_async_internal(
        self,
        tool_name: str,
        params: dict[str, Any],
        timeout: int,
        trace_id: str | None,
    ) -> ToolResult:
        """Internal async method to execute tool call."""
        if trace_id is None:
            trace_id = str(uuid.uuid4())
        
        # Use tracing if enabled
        if self._client._tracer:
            return await self._traced_call_async(tool_name, params, timeout, trace_id)
        
        return await self._execute_call_async(tool_name, params, timeout, trace_id)
    
    async def _execute_call_async(
        self,
        tool_name: str,
        params: dict[str, Any],
        timeout: int,
        trace_id: str,
    ) -> ToolResult:
        """Execute the actual async tool call without tracing."""
        request_body = {
            "tool_name": tool_name,
            "params": params,
            "options": {
                "timeout_seconds": timeout,
                "trace_id": trace_id,
            },
        }
        
        response = await self._client._request_async(
            "POST",
            "/tools/execute",
            json=request_body,
            params={"timeout": timeout},
            timeout=timeout + 5,
        )
        
        result = ToolResult.model_validate(response)
        
        if result.status == ResultStatus.ERROR and result.error:
            self._raise_for_error(tool_name, result.error.code, result.error.message)
        
        return result
    
    async def _traced_call_async(
        self,
        tool_name: str,
        params: dict[str, Any],
        timeout: int,
        trace_id: str,
    ) -> ToolResult:
        """Execute async tool call with Opik tracing."""
        tracer = self._client._tracer
        assert tracer is not None  # Guaranteed by caller
        
        # Note: Opik's context manager works synchronously,
        # but the actual call inside is async
        with tracer.trace_tool(tool_name, params, trace_id, timeout) as span:
            result = await self._execute_call_async(tool_name, params, timeout, trace_id)
            span.set_result(result)
            return result
    
    def submit(
        self,
        tool_name: str,
        trace_id: str | None = None,
        **params: Any,
    ) -> AsyncToolResponse:
        """
        Submit a tool for async execution (fire-and-forget).
        
        Use this for long-running operations or when you don't need
        to wait for the result immediately.
        
        Args:
            tool_name: Name of the tool to execute
            trace_id: Optional trace ID for observability
            **params: Tool parameters
        
        Returns:
            AsyncToolResponse with request_id for polling
        """
        if trace_id is None:
            trace_id = str(uuid.uuid4())
        
        request_body = {
            "tool_name": tool_name,
            "params": params,
            "options": {"trace_id": trace_id},
        }
        
        response = self._client._request("POST", "/tools/call", json=request_body)
        return AsyncToolResponse.model_validate(response)
    
    def get_result(
        self,
        request_id: str,
        wait: bool = False,
        timeout: int = 30,
    ) -> ToolResult:
        """
        Get the result of an async tool execution.
        
        Args:
            request_id: Request ID from submit()
            wait: If True, use long polling
            timeout: Max seconds to wait (if wait=True)
        
        Returns:
            ToolResult (may be pending if wait=False)
        """
        params = {"wait": str(wait).lower()}
        if wait:
            params["timeout"] = str(timeout)
        
        response = self._client._request(
            "GET",
            f"/tools/result/{request_id}",
            params=params,
            timeout=timeout + 5 if wait else 10,
        )
        return ToolResult.model_validate(response)
    
    def list(
        self,
        status: ToolStatus | None = None,
        type: ToolType | None = None,
        limit: int = 100,
        offset: int = 0,
    ) -> ToolListResponse:
        """
        List available tools.
        
        Args:
            status: Filter by status (active, deprecated, disabled)
            type: Filter by type (api_backed, pure_function, composite)
            limit: Max tools to return (1-1000)
            offset: Offset for pagination
        
        Returns:
            ToolListResponse with tools and total count
        """
        params: dict[str, Any] = {"limit": limit, "offset": offset}
        if status:
            params["status"] = status.value
        if type:
            params["type"] = type.value
        
        response = self._client._request("GET", "/tools", params=params)
        return ToolListResponse.model_validate(response)
    
    def get(self, tool_name: str, version: str | None = None) -> Tool:
        """
        Get details for a specific tool.
        
        Args:
            tool_name: Name of the tool
            version: Optional specific version
        
        Returns:
            Full Tool definition with schema and instructions
        """
        params = {}
        if version:
            params["version"] = version
        
        response = self._client._request("GET", f"/tools/{tool_name}", params=params)
        return Tool.model_validate(response)
    
    def _raise_for_error(self, tool_name: str, code: str, message: str) -> None:
        """Raise appropriate exception for error code."""
        exc_class = ERROR_CODE_MAP.get(code, ToolExecutionError)
        
        # Handle special cases that need tool_name
        if exc_class in (ToolNotFoundError, ToolDisabledError, ToolTimeoutError, ToolExecutionError):
            raise exc_class(tool_name=tool_name, message=message)
        
        raise exc_class(message=message)


class Session:
    """
    Session context for grouping related tool calls.
    
    Created via `client.create_session()`. All tool calls made through
    the session automatically include the session_id in their options,
    enabling session-level trace grouping in the admin console.
    
    Can be used as a context manager:
    
        with client.create_session() as session:
            session.tools.call("get_billing_details", customer_id="123")
            session.tools.call("calculate_proration", amount=100, days_used=15, total_days=30)
    """
    
    def __init__(self, client: "BoringOS", session_id: str | None = None, metadata: dict[str, Any] | None = None):
        """
        Initialize a session.
        
        Args:
            client: The BoringOS client instance
            session_id: Optional session ID (auto-generated if not provided)
            metadata: Optional metadata for the session
        """
        self._client = client
        self.session_id = session_id or f"ses_{uuid.uuid4().hex[:16]}"
        self.metadata = metadata or {}
        self._tools = SessionToolsClient(self)
    
    @property
    def tools(self) -> "SessionToolsClient":
        """Access session-scoped tool operations."""
        return self._tools
    
    def __enter__(self) -> "Session":
        """Context manager entry."""
        return self
    
    def __exit__(self, *args: Any) -> None:
        """Context manager exit."""
        pass  # Session is stateless from Boring-OS perspective
    
    async def __aenter__(self) -> "Session":
        """Async context manager entry."""
        return self
    
    async def __aexit__(self, *args: Any) -> None:
        """Async context manager exit."""
        pass


class SessionToolsClient:
    """
    Tool operations scoped to a session.
    
    Automatically injects session_id into all tool call options.
    """
    
    def __init__(self, session: Session):
        self._session = session
        self._client = session._client
    
    def call(
        self,
        tool_name: str,
        timeout: int = 30,
        trace_id: str | None = None,
        **params: Any,
    ) -> dict[str, Any]:
        """
        Execute a tool within this session (synchronous).
        
        The session_id is automatically included in the request options.
        """
        if trace_id is None:
            trace_id = str(uuid.uuid4())
        
        # Build request with session_id in options
        request_body = {
            "tool_name": tool_name,
            "params": params,
            "options": {
                "timeout_seconds": timeout,
                "trace_id": trace_id,
                "session_id": self._session.session_id,
            },
        }
        
        response = self._client._request(
            "POST",
            "/tools/execute",
            json=request_body,
            params={"timeout": timeout},
            timeout=timeout + 5,
        )
        
        result = ToolResult.model_validate(response)
        
        if result.status == ResultStatus.ERROR and result.error:
            self._client.tools._raise_for_error(tool_name, result.error.code, result.error.message)
        
        return result.data
    
    async def call_async(
        self,
        tool_name: str,
        timeout: int = 30,
        trace_id: str | None = None,
        **params: Any,
    ) -> dict[str, Any]:
        """
        Execute a tool within this session (asynchronous).
        
        The session_id is automatically included in the request options.
        """
        if trace_id is None:
            trace_id = str(uuid.uuid4())
        
        request_body = {
            "tool_name": tool_name,
            "params": params,
            "options": {
                "timeout_seconds": timeout,
                "trace_id": trace_id,
                "session_id": self._session.session_id,
            },
        }
        
        response = await self._client._request_async(
            "POST",
            "/tools/execute",
            json=request_body,
            params={"timeout": timeout},
            timeout=timeout + 5,
        )
        
        result = ToolResult.model_validate(response)
        
        if result.status == ResultStatus.ERROR and result.error:
            self._client.tools._raise_for_error(tool_name, result.error.code, result.error.message)
        
        return result.data


class BoringOS:
    """
    Main client for the Boring-OS API.
    
    Provides access to tool operations via the `tools` property.
    Supports optional Opik tracing for observability.
    
    Example:
        from boring_os import BoringOS
        
        client = BoringOS(
            base_url="http://localhost:8080",
            api_key="bos_dev_key_for_local_testing"
        )
        
        # Call a tool
        result = client.tools.call("get_customer_profile", customer_id="123")
        
        # List tools
        tools = client.tools.list()
        
        # With tracing enabled
        client = BoringOS(
            api_key="...",
            enable_tracing=True,
            opik_project="my-agent"
        )
    """
    
    def __init__(
        self,
        api_key: str,
        base_url: str = "http://localhost:8080",
        timeout: float = 30.0,
        max_retries: int = 3,
        # Tracing options
        enable_tracing: bool = False,
        opik_project: str = "boring-os",
        opik_use_local: bool = True,
        opik_url: str | None = None,
        opik_api_key: str | None = None,
        opik_workspace: str | None = None,
    ):
        """
        Initialize the Boring-OS client.
        
        Args:
            api_key: API key for authentication
            base_url: Base URL of the Boring-OS gateway
            timeout: Default request timeout in seconds
            max_retries: Max retries for transient errors
            enable_tracing: Enable Opik tracing for all tool calls
            opik_project: Project name for Opik traces
            opik_use_local: If True, use local Opik server; False for Comet cloud
            opik_url: Custom Opik server URL (defaults to http://localhost:5173 for local)
            opik_api_key: Comet API key (required if opik_use_local=False)
            opik_workspace: Comet workspace (required if opik_use_local=False)
        """
        self._api_key = api_key
        self._base_url = base_url.rstrip("/")
        self._timeout = timeout
        self._max_retries = max_retries
        
        # Initialize tracing if enabled
        self._tracer: "OpikTracer | None" = None
        if enable_tracing:
            from boring_os.tracing import OpikTracer
            self._tracer = OpikTracer(
                project_name=opik_project,
                use_local=opik_use_local,
                url=opik_url,
                api_key=opik_api_key,
                workspace=opik_workspace,
            )
        
        # Initialize sub-clients
        self._tools = ToolsClient(self)

        # Lazy-imported sub-clients for agents, models, and flows
        from boring_os.agents import AgentsClient
        from boring_os.models_client import ModelsClient
        from boring_os.flows import FlowsClient

        self._agents = AgentsClient(self)
        self._models_client = ModelsClient(self)
        self._flows = FlowsClient(self)
        
        # HTTP clients (created lazily)
        self._sync_client: httpx.Client | None = None
        self._async_client: httpx.AsyncClient | None = None
    
    @property
    def tracer(self) -> "OpikTracer | None":
        """Access the Opik tracer (if tracing is enabled)."""
        return self._tracer
    
    @property
    def tools(self) -> ToolsClient:
        """Access tool operations."""
        return self._tools

    @property
    def agents(self) -> "AgentsClient":
        """Access agent operations (run, stream, sessions, CRUD)."""
        return self._agents

    @property
    def models(self) -> "ModelsClient":
        """Access direct model operations (complete, stream, embed)."""
        return self._models_client

    @property
    def flows(self) -> "FlowsClient":
        """Access multi-agent flow operations (run, stream, CRUD)."""
        return self._flows
    
    def create_session(
        self,
        session_id: str | None = None,
        metadata: dict[str, Any] | None = None,
    ) -> Session:
        """
        Create a new session for grouping related tool calls.
        
        All tool calls made through the session automatically include
        the session_id, enabling session-level trace grouping in the admin console.
        
        Args:
            session_id: Optional session ID (auto-generated if not provided)
            metadata: Optional metadata to associate with the session
        
        Returns:
            Session instance that can be used as a context manager
        
        Example:
            with client.create_session() as session:
                session.tools.call("get_billing_details", customer_id="123")
                session.tools.call("calculate_proration", amount=100, days_used=15, total_days=30)
        """
        return Session(self, session_id=session_id, metadata=metadata)
    
    def _get_sync_client(self) -> httpx.Client:
        """Get or create sync HTTP client."""
        if self._sync_client is None:
            self._sync_client = httpx.Client(
                base_url=self._base_url,
                headers=self._get_headers(),
                timeout=self._timeout,
            )
        return self._sync_client
    
    def _get_async_client(self) -> httpx.AsyncClient:
        """Get or create async HTTP client."""
        if self._async_client is None:
            self._async_client = httpx.AsyncClient(
                base_url=self._base_url,
                headers=self._get_headers(),
                timeout=self._timeout,
            )
        return self._async_client
    
    def _get_headers(self) -> dict[str, str]:
        """Get common request headers."""
        return {
            "X-API-Key": self._api_key,
            "Content-Type": "application/json",
            "Accept": "application/json",
        }
    
    def _request(
        self,
        method: str,
        path: str,
        timeout: float | None = None,
        **kwargs: Any,
    ) -> dict[str, Any]:
        """Make a synchronous HTTP request."""
        client = self._get_sync_client()
        
        try:
            response = client.request(
                method,
                path,
                timeout=timeout or self._timeout,
                **kwargs,
            )
        except httpx.ConnectError as e:
            raise ConnectionError(f"Failed to connect to {self._base_url}: {e}")
        except httpx.TimeoutException as e:
            raise ToolTimeoutError(
                tool_name="unknown",
                message=f"Request timed out: {e}",
            )
        
        return self._handle_response(response)
    
    async def _request_async(
        self,
        method: str,
        path: str,
        timeout: float | None = None,
        **kwargs: Any,
    ) -> dict[str, Any]:
        """Make an asynchronous HTTP request."""
        client = self._get_async_client()
        
        try:
            response = await client.request(
                method,
                path,
                timeout=timeout or self._timeout,
                **kwargs,
            )
        except httpx.ConnectError as e:
            raise ConnectionError(f"Failed to connect to {self._base_url}: {e}")
        except httpx.TimeoutException as e:
            raise ToolTimeoutError(
                tool_name="unknown",
                message=f"Request timed out: {e}",
            )
        
        return self._handle_response(response)
    
    def _handle_response(self, response: httpx.Response) -> dict[str, Any]:
        """Handle HTTP response and raise exceptions for errors."""
        # Parse JSON response
        try:
            data = response.json()
        except Exception:
            if response.status_code >= 400:
                raise BoringOSError(
                    f"Request failed with status {response.status_code}",
                    code="HTTP_ERROR",
                )
            raise BoringOSError("Invalid JSON response")
        
        # Handle HTTP-level errors
        if response.status_code == 401:
            raise AuthenticationError(data.get("detail", "Authentication failed"))
        
        if response.status_code == 403:
            raise AuthorizationError(data.get("detail", "Access denied"))
        
        if response.status_code == 429:
            retry_after = response.headers.get("Retry-After")
            raise RateLimitError(
                data.get("detail", "Rate limit exceeded"),
                retry_after=int(retry_after) if retry_after else None,
            )
        
        if response.status_code == 404:
            raise ToolNotFoundError(
                tool_name="unknown",
                message=data.get("detail", "Not found"),
            )
        
        if response.status_code == 422:
            raise ValidationError(
                data.get("detail", "Validation failed"),
                details=data.get("errors"),
            )
        
        if response.status_code >= 500:
            raise BoringOSError(
                f"Server error: {data.get('detail', 'Internal error')}",
                code="SERVER_ERROR",
                retryable=True,
            )
        
        # Success - return data
        return data
    
    async def _stream_request(
        self,
        method: str,
        path: str,
        **kwargs: Any,
    ) -> Any:
        """
        Make an SSE streaming request and yield parsed event data.

        Reads the response line-by-line, parsing SSE events in the
        standard format (data: {...}). Yields each parsed JSON dict.

        Args:
            method: HTTP method
            path: URL path
            **kwargs: Additional httpx request arguments

        Yields:
            Parsed JSON dicts from SSE data lines
        """
        import json as json_lib

        client = self._get_async_client()

        try:
            async with client.stream(
                method,
                path,
                timeout=kwargs.pop("timeout", 300),
                **kwargs,
            ) as response:
                # Check for errors before streaming
                if response.status_code >= 400:
                    body = await response.aread()
                    try:
                        data = json_lib.loads(body)
                    except Exception:
                        raise BoringOSError(
                            f"Streaming request failed with status {response.status_code}",
                            code="HTTP_ERROR",
                        )
                    self._handle_response_data(response.status_code, data)

                # Parse SSE lines
                async for line in response.aiter_lines():
                    line = line.strip()
                    if not line:
                        continue
                    if line.startswith("data: "):
                        payload = line[6:]
                        if payload == "[DONE]":
                            break
                        try:
                            yield json_lib.loads(payload)
                        except json_lib.JSONDecodeError:
                            continue
        except httpx.ConnectError as e:
            raise ConnectionError(f"Failed to connect to {self._base_url}: {e}")
        except httpx.TimeoutException as e:
            raise ToolTimeoutError(
                tool_name="unknown",
                message=f"Streaming request timed out: {e}",
            )

    def _handle_response_data(
        self, status_code: int, data: dict[str, Any]
    ) -> None:
        """Handle error response data from streaming requests."""
        if status_code == 401:
            raise AuthenticationError(data.get("detail", "Authentication failed"))
        if status_code == 403:
            raise AuthorizationError(data.get("detail", "Access denied"))
        if status_code == 404:
            raise ToolNotFoundError(
                tool_name="unknown",
                message=data.get("detail", "Not found"),
            )
        if status_code >= 500:
            raise BoringOSError(
                f"Server error: {data.get('detail', 'Internal error')}",
                code="SERVER_ERROR",
                retryable=True,
            )

    def close(self) -> None:
        """Close HTTP clients."""
        if self._sync_client:
            self._sync_client.close()
            self._sync_client = None
    
    async def aclose(self) -> None:
        """Close async HTTP clients."""
        if self._async_client:
            await self._async_client.aclose()
            self._async_client = None
    
    def __enter__(self) -> "BoringOS":
        """Context manager entry."""
        return self
    
    def __exit__(self, *args: Any) -> None:
        """Context manager exit."""
        self.close()
    
    async def __aenter__(self) -> "BoringOS":
        """Async context manager entry."""
        return self
    
    async def __aexit__(self, *args: Any) -> None:
        """Async context manager exit."""
        await self.aclose()
