"""
FluxTokens API Client
"""

import json
from typing import Any, AsyncIterator, Dict, Iterator, List, Optional, Union

import httpx

from fluxtokens.errors import (
    AuthenticationError,
    BadRequestError,
    ConnectionError,
    FluxTokensError,
    InsufficientBalanceError,
    InternalServerError,
    RateLimitError,
    TimeoutError,
)
from fluxtokens.types import (
    ChatCompletionChunk,
    ChatCompletionRequest,
    ChatCompletionResponse,
    ChatMessage,
    ModelInfo,
    MODELS,
)


DEFAULT_BASE_URL = "https://api.fluxtokens.io"
DEFAULT_TIMEOUT = 30.0
DEFAULT_MAX_RETRIES = 2


class ChatCompletions:
    """Chat completions API."""

    def __init__(self, client: "FluxTokens") -> None:
        self._client = client

    def create(
        self,
        *,
        model: str,
        messages: List[Union[ChatMessage, Dict[str, Any]]],
        temperature: Optional[float] = None,
        max_tokens: Optional[int] = None,
        top_p: Optional[float] = None,
        top_k: Optional[int] = None,
        frequency_penalty: Optional[float] = None,
        presence_penalty: Optional[float] = None,
        stop: Optional[Union[str, List[str]]] = None,
        stream: bool = False,
        user: Optional[str] = None,
        tools: Optional[List[Dict[str, Any]]] = None,
        tool_choice: Optional[Union[str, Dict[str, Any]]] = None,
        **kwargs: Any,
    ) -> Union[ChatCompletionResponse, Iterator[ChatCompletionChunk]]:
        """Create a chat completion.

        Args:
            model: Model to use for completion.
            messages: List of messages in the conversation.
            temperature: Sampling temperature (0-2).
            max_tokens: Maximum tokens to generate.
            top_p: Nucleus sampling parameter (0-1).
            top_k: Top-k sampling parameter.
            frequency_penalty: Frequency penalty (-2 to 2).
            presence_penalty: Presence penalty (-2 to 2).
            stop: Stop sequences.
            stream: Whether to stream the response.
            user: User identifier for abuse detection.
            tools: Tools available to the model.
            tool_choice: How the model should use tools.

        Returns:
            ChatCompletionResponse or Iterator of ChatCompletionChunk if streaming.
        """
        # Convert dict messages to ChatMessage
        processed_messages = [
            msg if isinstance(msg, ChatMessage) else ChatMessage(**msg)
            for msg in messages
        ]

        request = ChatCompletionRequest(
            model=model,
            messages=processed_messages,
            temperature=temperature,
            max_tokens=max_tokens,
            top_p=top_p,
            top_k=top_k,
            frequency_penalty=frequency_penalty,
            presence_penalty=presence_penalty,
            stop=stop,
            stream=stream,
            user=user,
            tools=tools,
            tool_choice=tool_choice,
        )

        if stream:
            return self._stream(request)
        return self._create(request)

    def _create(self, request: ChatCompletionRequest) -> ChatCompletionResponse:
        """Create a non-streaming chat completion."""
        response = self._client._request(
            "POST",
            "/v1/chat/completions",
            json=request.model_dump(exclude_none=True),
        )
        return ChatCompletionResponse(**response)

    def _stream(self, request: ChatCompletionRequest) -> Iterator[ChatCompletionChunk]:
        """Create a streaming chat completion."""
        request.stream = True
        
        with self._client._client.stream(
            "POST",
            f"{self._client._base_url}/v1/chat/completions",
            json=request.model_dump(exclude_none=True),
            headers=self._client._headers,
            timeout=self._client._timeout,
        ) as response:
            if response.status_code != 200:
                self._client._handle_error(response)

            for line in response.iter_lines():
                if not line or line == "data: [DONE]":
                    continue
                if line.startswith("data: "):
                    try:
                        data = json.loads(line[6:])
                        yield ChatCompletionChunk(**data)
                    except json.JSONDecodeError:
                        continue


class Chat:
    """Chat API namespace."""

    def __init__(self, client: "FluxTokens") -> None:
        self.completions = ChatCompletions(client)


class Models:
    """Models API."""

    def list(self) -> List[ModelInfo]:
        """List available models."""
        return list(MODELS.values())

    def get(self, model_id: str) -> Optional[ModelInfo]:
        """Get model info by ID."""
        return MODELS.get(model_id)


class FluxTokens:
    """FluxTokens API client.

    Example:
        >>> from fluxtokens import FluxTokens
        >>>
        >>> client = FluxTokens(api_key="sk-flux-...")
        >>>
        >>> response = client.chat.completions.create(
        ...     model="gpt-4.1-mini",
        ...     messages=[{"role": "user", "content": "Hello!"}]
        ... )
        >>>
        >>> print(response.choices[0].message.content)
    """

    def __init__(
        self,
        api_key: str,
        *,
        base_url: str = DEFAULT_BASE_URL,
        timeout: float = DEFAULT_TIMEOUT,
        max_retries: int = DEFAULT_MAX_RETRIES,
    ) -> None:
        """Initialize the FluxTokens client.

        Args:
            api_key: Your FluxTokens API key (starts with sk-flux-).
            base_url: Base URL for the API.
            timeout: Request timeout in seconds.
            max_retries: Maximum number of retries.
        """
        if not api_key:
            raise ValueError("API key is required. Get one at https://fluxtokens.io")

        self._api_key = api_key
        self._base_url = base_url.rstrip("/")
        self._timeout = timeout
        self._max_retries = max_retries
        self._client = httpx.Client(timeout=timeout)

        # API namespaces
        self.chat = Chat(self)
        self.models = Models()

    @property
    def _headers(self) -> Dict[str, str]:
        return {
            "Authorization": f"Bearer {self._api_key}",
            "Content-Type": "application/json",
        }

    def _request(
        self,
        method: str,
        path: str,
        *,
        json: Optional[Dict[str, Any]] = None,
        retries: int = 0,
    ) -> Dict[str, Any]:
        """Make an HTTP request."""
        try:
            response = self._client.request(
                method,
                f"{self._base_url}{path}",
                json=json,
                headers=self._headers,
            )

            if response.status_code != 200:
                self._handle_error(response)

            return response.json()

        except httpx.TimeoutException:
            raise TimeoutError()
        except httpx.ConnectError as e:
            raise ConnectionError(str(e))
        except (RateLimitError, InternalServerError) as e:
            if retries < self._max_retries:
                import time
                time.sleep(min(2 ** retries, 10))
                return self._request(method, path, json=json, retries=retries + 1)
            raise

    def _handle_error(self, response: httpx.Response) -> None:
        """Handle error response."""
        try:
            data = response.json()
            message = data.get("error", {}).get("message", response.text)
            code = data.get("error", {}).get("code")
            param = data.get("error", {}).get("param")
        except Exception:
            message = response.text or "Unknown error"
            code = None
            param = None

        status = response.status_code

        if status == 401:
            raise AuthenticationError(message)
        elif status == 402:
            raise InsufficientBalanceError(message)
        elif status == 429:
            retry_after = response.headers.get("Retry-After")
            raise RateLimitError(
                message,
                retry_after=int(retry_after) if retry_after else None,
            )
        elif status == 400:
            raise BadRequestError(message, code=code, param=param)
        elif status >= 500:
            raise InternalServerError(message)
        else:
            raise FluxTokensError(message, status=status)

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

    def __enter__(self) -> "FluxTokens":
        return self

    def __exit__(self, *args: Any) -> None:
        self.close()


# Async Client
class AsyncChatCompletions:
    """Async chat completions API."""

    def __init__(self, client: "AsyncFluxTokens") -> None:
        self._client = client

    async def create(
        self,
        *,
        model: str,
        messages: List[Union[ChatMessage, Dict[str, Any]]],
        temperature: Optional[float] = None,
        max_tokens: Optional[int] = None,
        top_p: Optional[float] = None,
        top_k: Optional[int] = None,
        frequency_penalty: Optional[float] = None,
        presence_penalty: Optional[float] = None,
        stop: Optional[Union[str, List[str]]] = None,
        stream: bool = False,
        user: Optional[str] = None,
        tools: Optional[List[Dict[str, Any]]] = None,
        tool_choice: Optional[Union[str, Dict[str, Any]]] = None,
        **kwargs: Any,
    ) -> Union[ChatCompletionResponse, AsyncIterator[ChatCompletionChunk]]:
        """Create a chat completion asynchronously."""
        processed_messages = [
            msg if isinstance(msg, ChatMessage) else ChatMessage(**msg)
            for msg in messages
        ]

        request = ChatCompletionRequest(
            model=model,
            messages=processed_messages,
            temperature=temperature,
            max_tokens=max_tokens,
            top_p=top_p,
            top_k=top_k,
            frequency_penalty=frequency_penalty,
            presence_penalty=presence_penalty,
            stop=stop,
            stream=stream,
            user=user,
            tools=tools,
            tool_choice=tool_choice,
        )

        if stream:
            return self._stream(request)
        return await self._create(request)

    async def _create(self, request: ChatCompletionRequest) -> ChatCompletionResponse:
        """Create a non-streaming chat completion."""
        response = await self._client._request(
            "POST",
            "/v1/chat/completions",
            json=request.model_dump(exclude_none=True),
        )
        return ChatCompletionResponse(**response)

    async def _stream(
        self, request: ChatCompletionRequest
    ) -> AsyncIterator[ChatCompletionChunk]:
        """Create a streaming chat completion."""
        request.stream = True

        async with self._client._client.stream(
            "POST",
            f"{self._client._base_url}/v1/chat/completions",
            json=request.model_dump(exclude_none=True),
            headers=self._client._headers,
            timeout=self._client._timeout,
        ) as response:
            if response.status_code != 200:
                await self._client._handle_error_async(response)

            async for line in response.aiter_lines():
                if not line or line == "data: [DONE]":
                    continue
                if line.startswith("data: "):
                    try:
                        data = json.loads(line[6:])
                        yield ChatCompletionChunk(**data)
                    except json.JSONDecodeError:
                        continue


class AsyncChat:
    """Async chat API namespace."""

    def __init__(self, client: "AsyncFluxTokens") -> None:
        self.completions = AsyncChatCompletions(client)


class AsyncFluxTokens:
    """Async FluxTokens API client.

    Example:
        >>> from fluxtokens import AsyncFluxTokens
        >>> import asyncio
        >>>
        >>> async def main():
        ...     client = AsyncFluxTokens(api_key="sk-flux-...")
        ...     response = await client.chat.completions.create(
        ...         model="gpt-4.1-mini",
        ...         messages=[{"role": "user", "content": "Hello!"}]
        ...     )
        ...     print(response.choices[0].message.content)
        ...     await client.close()
        >>>
        >>> asyncio.run(main())
    """

    def __init__(
        self,
        api_key: str,
        *,
        base_url: str = DEFAULT_BASE_URL,
        timeout: float = DEFAULT_TIMEOUT,
        max_retries: int = DEFAULT_MAX_RETRIES,
    ) -> None:
        if not api_key:
            raise ValueError("API key is required. Get one at https://fluxtokens.io")

        self._api_key = api_key
        self._base_url = base_url.rstrip("/")
        self._timeout = timeout
        self._max_retries = max_retries
        self._client = httpx.AsyncClient(timeout=timeout)

        self.chat = AsyncChat(self)
        self.models = Models()

    @property
    def _headers(self) -> Dict[str, str]:
        return {
            "Authorization": f"Bearer {self._api_key}",
            "Content-Type": "application/json",
        }

    async def _request(
        self,
        method: str,
        path: str,
        *,
        json: Optional[Dict[str, Any]] = None,
        retries: int = 0,
    ) -> Dict[str, Any]:
        """Make an async HTTP request."""
        try:
            response = await self._client.request(
                method,
                f"{self._base_url}{path}",
                json=json,
                headers=self._headers,
            )

            if response.status_code != 200:
                await self._handle_error_async(response)

            return response.json()

        except httpx.TimeoutException:
            raise TimeoutError()
        except httpx.ConnectError as e:
            raise ConnectionError(str(e))
        except (RateLimitError, InternalServerError) as e:
            if retries < self._max_retries:
                import asyncio
                await asyncio.sleep(min(2 ** retries, 10))
                return await self._request(method, path, json=json, retries=retries + 1)
            raise

    async def _handle_error_async(self, response: httpx.Response) -> None:
        """Handle error response."""
        try:
            data = response.json()
            message = data.get("error", {}).get("message", response.text)
            code = data.get("error", {}).get("code")
            param = data.get("error", {}).get("param")
        except Exception:
            message = response.text or "Unknown error"
            code = None
            param = None

        status = response.status_code

        if status == 401:
            raise AuthenticationError(message)
        elif status == 402:
            raise InsufficientBalanceError(message)
        elif status == 429:
            retry_after = response.headers.get("Retry-After")
            raise RateLimitError(
                message,
                retry_after=int(retry_after) if retry_after else None,
            )
        elif status == 400:
            raise BadRequestError(message, code=code, param=param)
        elif status >= 500:
            raise InternalServerError(message)
        else:
            raise FluxTokensError(message, status=status)

    async def close(self) -> None:
        """Close the HTTP client."""
        await self._client.aclose()

    async def __aenter__(self) -> "AsyncFluxTokens":
        return self

    async def __aexit__(self, *args: Any) -> None:
        await self.close()
