"""Main client for interacting with the Lumo API."""

import json
from typing import Optional, List, Iterator, Union

import requests

from lumo_sdk.models import (
    RunTaskRequest,
    RunTaskResponse,
    Message,
    MCPServer,
    StreamTokenEvent,
    StreamStepEvent,
    StreamDoneEvent,
)
from lumo_sdk.exceptions import (
    LumoAPIError,
    LumoAuthenticationError,
    LumoRateLimitError,
    LumoServerError,
)
import os


class LumoClient:
    """Client for interacting with the Lumo API."""

    BASE_URL = "https://api.starlight-search.com"

    def __init__(self, api_key: Optional[str] = None, base_url: Optional[str] = None):
        """
        Initialize the Lumo client.

        Args:
            api_key: Optional Lumo API key (can also be provided via LUMO_API_KEY environment variable)
            base_url: Optional custom base URL (defaults to official API URL)
        """
        if api_key:
            self.api_key = api_key
        else:
            self.api_key = os.getenv("LUMO_API_KEY")
        if not self.api_key:
            raise ValueError(
                "API key is required. Provide it as a constructor parameter or set LUMO_API_KEY environment variable."
            )
        self.base_url = base_url or self.BASE_URL
        self.session = requests.Session()
        headers = {
            "Content-Type": "application/json",
        }
        if self.api_key:
            headers["Authorization"] = f"Bearer {self.api_key}"
        self.session.headers.update(headers)

    def run_task(
        self,
        task: str,
        model: str,
        base_url: str,
        tools: Optional[List[str]] = None,
        max_steps: Optional[int] = None,
        agent_type: Optional[str] = None,
        history: Optional[List[Message]] = None,
        mcp_servers: Optional[List[MCPServer]] = None,
    ) -> RunTaskResponse:
        """
        Execute a task using the specified model and tools.

        Args:
            task: The task to execute
            model: Model ID (e.g., 'gpt-4', 'qwen2.5', 'gemini-2.0-flash')
            base_url: Base URL for the upstream API
            tools: Optional array of tool names to make available
            max_steps: Optional maximum number of steps to take
            agent_type: Optional type of agent to use ('function-calling' or 'mcp')
            history: Optional array of Message objects containing prior conversation context

        Returns:
            RunTaskResponse containing the final answer and steps

        Raises:
            LumoAuthenticationError: If authentication fails (401)
            LumoRateLimitError: If rate limit is exceeded (429)
            LumoServerError: If server error occurs (500)
            LumoAPIError: For other API errors
        """
        request_data = RunTaskRequest(
            task=task,
            model=model,
            base_url=base_url,
            tools=tools,
            max_steps=max_steps,
            agent_type=agent_type,
            history=history,
            mcp_servers=mcp_servers,
        )

        try:
            response = self.session.post(
                f"{self.base_url}/run",
                json=request_data.model_dump(exclude_none=True),
            )
            self._handle_response(response)
            return RunTaskResponse(**response.json())
        except requests.exceptions.RequestException as e:
            raise LumoAPIError(f"Request failed: {str(e)}") from e

    def _handle_response(self, response: requests.Response) -> None:
        """
        Handle API response and raise appropriate exceptions.

        Args:
            response: The HTTP response object

        Raises:
            LumoAuthenticationError: For 401 errors
            LumoRateLimitError: For 429 errors
            LumoServerError: For 500 errors
            LumoAPIError: For other error status codes
        """
        if response.status_code == 200:
            return

        try:
            error_data = response.json()
        except ValueError:
            error_data = {"message": response.text or "Unknown error"}

        error_message = error_data.get("message", error_data.get("error", "Unknown error"))
        status_code = response.status_code

        if status_code == 400:
            raise LumoAPIError(
                f"Bad Request: {error_message}", status_code=status_code, response=error_data
            )
        elif status_code == 401:
            raise LumoAuthenticationError(
                f"Authentication failed: {error_message}",
                status_code=status_code,
                response=error_data,
            )
        elif status_code == 429:
            raise LumoRateLimitError(
                f"Rate limit exceeded: {error_message}",
                status_code=status_code,
                response=error_data,
            )
        elif status_code >= 500:
            raise LumoServerError(
                f"Server error: {error_message}", status_code=status_code, response=error_data
            )
        else:
            raise LumoAPIError(
                f"API error ({status_code}): {error_message}",
                status_code=status_code,
                response=error_data,
            )

    def stream_task(
        self,
        task: str,
        model: str,
        base_url: str,
        tools: Optional[List[str]] = None,
        max_steps: Optional[int] = None,
        agent_type: Optional[str] = None,
        history: Optional[List[Message]] = None,
        mcp_servers: Optional[List[MCPServer]] = None,
    ) -> Iterator[Union[StreamTokenEvent, StreamStepEvent, StreamDoneEvent]]:
        """
        Stream task execution using the specified model and tools.
        
        This method yields events as they arrive from the API:
        - StreamTokenEvent: Individual tokens from the LLM response
        - StreamStepEvent: Step information including tool calls and output
        - StreamDoneEvent: Signals completion of the task

        Args:
            task: The task to execute
            model: Model ID (e.g., 'gpt-4', 'qwen2.5', 'gemini-2.0-flash')
            base_url: Base URL for the upstream API
            tools: Optional array of tool names to make available
            max_steps: Optional maximum number of steps to take
            agent_type: Optional type of agent to use ('function-calling' or 'mcp')
            history: Optional array of Message objects containing prior conversation context
            mcp_servers: Optional list of MCP server configurations (required when agent_type is 'mcp')

        Yields:
            Union[StreamTokenEvent, StreamStepEvent, StreamDoneEvent]: Stream events

        Raises:
            LumoAuthenticationError: If authentication fails (401)
            LumoRateLimitError: If rate limit is exceeded (429)
            LumoServerError: If server error occurs (500)
            LumoAPIError: For other API errors

        Example:
            ```python
            client = LumoClient(api_key="your-api-key")
            for event in client.stream_task(
                task="What is the weather?",
                model="gpt-4.1-mini",
                base_url="https://api.openai.com/v1/chat/completions",
                tools=["ExaSearchTool"]
            ):
                if isinstance(event, StreamTokenEvent):
                    print(event.content, end="", flush=True)
                elif isinstance(event, StreamStepEvent):
                    print(f"\\nStep: {event.step}")
                elif isinstance(event, StreamDoneEvent):
                    print("\\nDone!")
            ```
        """
        request_data = RunTaskRequest(
            task=task,
            model=model,
            base_url=base_url,
            tools=tools,
            max_steps=max_steps,
            agent_type=agent_type,
            history=history,
            mcp_servers=mcp_servers,
        )

        # Prepare headers for streaming
        headers = {
            "Content-Type": "application/json",
            "Accept": "text/event-stream",
        }
        if self.api_key:
            headers["Authorization"] = f"Bearer {self.api_key}"

        try:
            response = self.session.post(
                f"{self.base_url}/stream",
                json=request_data.model_dump(exclude_none=True),
                headers=headers,
                stream=True,
            )

            # Check for errors before streaming
            if response.status_code != 200:
                self._handle_response(response)
                return

            # Ensure response encoding is UTF-8 to handle special characters correctly
            response.encoding = 'utf-8'

            # Stream the response
            for line in response.iter_lines(decode_unicode=True):
                if not line:
                    continue

                # Parse Server-Sent Events format
                if line.startswith("data: "):
                    data = line[6:]  # Remove "data: " prefix

                    try:
                        event_data = json.loads(data)
                        event_type = event_data.get("type")

                        if event_type == "token":
                            yield StreamTokenEvent(**event_data)
                        elif event_type == "step":
                            yield StreamStepEvent(**event_data)
                        elif event_type == "done":
                            yield StreamDoneEvent(**event_data)
                        else:
                            # Unknown event type, skip or log
                            continue

                    except json.JSONDecodeError:
                        # Skip malformed JSON
                        continue
                    except Exception as e:
                        # Skip events that can't be parsed
                        continue

        except requests.exceptions.RequestException as e:
            raise LumoAPIError(f"Stream request failed: {str(e)}") from e

