"""Low-level HTTP client wrapper for sync and async usage."""

from __future__ import annotations

import asyncio
import json
import time
from collections.abc import Generator, AsyncGenerator
from typing import Any

import httpx

from melogenai.auth import AuthManager
from melogenai.config import Config
from melogenai.exceptions import (
    AuthenticationError,
    InsufficientCreditsError,
    MelogenError,
    RateLimitError,
    TaskFailedError,
    TaskTimeoutError,
    ValidationError,
)
from melogenai.types import COMPLETED_STATUSES, FAILED_STATUSES


def _raise_for_status(resp: httpx.Response) -> None:
    """Map HTTP status codes to SDK exceptions."""
    if resp.is_success:
        return
    status = resp.status_code
    try:
        body = resp.json()
        detail = body.get("detail") or body.get("error") or body.get("message") or resp.text
    except Exception:
        detail = resp.text

    if status == 401:
        raise AuthenticationError(str(detail))
    if status in (402, 403):
        raise InsufficientCreditsError(str(detail))
    if status == 422:
        raise ValidationError(str(detail))
    if status == 429:
        raise RateLimitError(str(detail))
    raise MelogenError(f"HTTP {status}: {detail}", status_code=status)


class HttpClient:
    """Synchronous HTTP client."""

    def __init__(self, config: Config, auth: AuthManager):
        self._config = config
        self._auth = auth

    def _headers(self, token: str) -> dict[str, str]:
        return {
            "Authorization": f"Bearer {token}",
            "Content-Type": "application/json",
        }

    def _get_token(self) -> str:
        token_data = self._auth.get_token(self._config.api_key)
        return token_data["token"]

    def request(
        self,
        method: str,
        path: str,
        *,
        json_body: dict[str, Any] | None = None,
        timeout: float | None = None,
    ) -> dict[str, Any]:
        token = self._get_token()
        with httpx.Client(timeout=timeout or self._config.timeout) as client:
            resp = client.request(
                method,
                f"{self._config.base_url}{path}",
                json=json_body,
                headers=self._headers(token),
            )
        if resp.status_code == 401:
            self._auth.invalidate(self._config.api_key)
        _raise_for_status(resp)
        return resp.json()

    def post(self, path: str, payload: dict[str, Any], timeout: float | None = None) -> dict[str, Any]:
        return self.request("POST", path, json_body=payload, timeout=timeout)

    def get(self, path: str, timeout: float | None = None) -> dict[str, Any]:
        return self.request("GET", path, timeout=timeout)

    def poll_task(
        self,
        path: str,
        task_id: str,
        interval: float | None = None,
        timeout: float | None = None,
    ) -> dict[str, Any]:
        interval = interval or self._config.poll_interval
        timeout = timeout or self._config.poll_timeout
        deadline = time.monotonic() + timeout

        while time.monotonic() < deadline:
            result = self.get(path)
            task = result.get("task", result)
            status = task.get("status", "")

            if status in COMPLETED_STATUSES:
                return result
            if status in FAILED_STATUSES:
                error_msg = task.get("error") or task.get("message") or "Task failed"
                raise TaskFailedError(task_id, error_msg)

            time.sleep(interval)

        raise TaskTimeoutError(task_id, timeout)

    def stream_sse(
        self,
        path: str,
        payload: dict[str, Any],
        timeout: float | None = None,
    ) -> Generator[str, None, None]:
        """POST to SSE endpoint and yield content chunks."""
        token = self._get_token()
        t = timeout or self._config.poll_timeout
        with httpx.Client(timeout=httpx.Timeout(t, connect=10)) as client:
            with client.stream(
                "POST",
                f"{self._config.base_url}{path}",
                json=payload,
                headers=self._headers(token),
            ) as resp:
                if resp.status_code == 401:
                    self._auth.invalidate(self._config.api_key)
                _raise_for_status(resp)

                for line in resp.iter_lines():
                    if not line.startswith("data: "):
                        continue
                    data_str = line[6:].strip()
                    if data_str == "[DONE]":
                        return

                    try:
                        event = json.loads(data_str)
                    except json.JSONDecodeError:
                        continue

                    chunk = _extract_sse_content(event)
                    if chunk is not None:
                        yield chunk

    def consume_sse(
        self,
        path: str,
        payload: dict[str, Any],
        timeout: float | None = None,
    ) -> str:
        """POST to SSE endpoint and return the full aggregated content."""
        return "".join(self.stream_sse(path, payload, timeout))


class AsyncHttpClient:
    """Asynchronous HTTP client."""

    def __init__(self, config: Config, auth: AuthManager):
        self._config = config
        self._auth = auth

    def _headers(self, token: str) -> dict[str, str]:
        return {
            "Authorization": f"Bearer {token}",
            "Content-Type": "application/json",
        }

    async def _get_token(self) -> str:
        token_data = await self._auth.async_get_token(self._config.api_key)
        return token_data["token"]

    async def request(
        self,
        method: str,
        path: str,
        *,
        json_body: dict[str, Any] | None = None,
        timeout: float | None = None,
    ) -> dict[str, Any]:
        token = await self._get_token()
        async with httpx.AsyncClient(timeout=timeout or self._config.timeout) as client:
            resp = await client.request(
                method,
                f"{self._config.base_url}{path}",
                json=json_body,
                headers=self._headers(token),
            )
        if resp.status_code == 401:
            self._auth.invalidate(self._config.api_key)
        _raise_for_status(resp)
        return resp.json()

    async def post(self, path: str, payload: dict[str, Any], timeout: float | None = None) -> dict[str, Any]:
        return await self.request("POST", path, json_body=payload, timeout=timeout)

    async def get(self, path: str, timeout: float | None = None) -> dict[str, Any]:
        return await self.request("GET", path, timeout=timeout)

    async def poll_task(
        self,
        path: str,
        task_id: str,
        interval: float | None = None,
        timeout: float | None = None,
    ) -> dict[str, Any]:
        interval = interval or self._config.poll_interval
        timeout = timeout or self._config.poll_timeout
        deadline = time.monotonic() + timeout

        while time.monotonic() < deadline:
            result = await self.get(path)
            task = result.get("task", result)
            status = task.get("status", "")

            if status in COMPLETED_STATUSES:
                return result
            if status in FAILED_STATUSES:
                error_msg = task.get("error") or task.get("message") or "Task failed"
                raise TaskFailedError(task_id, error_msg)

            await asyncio.sleep(interval)

        raise TaskTimeoutError(task_id, timeout)

    async def stream_sse(
        self,
        path: str,
        payload: dict[str, Any],
        timeout: float | None = None,
    ) -> AsyncGenerator[str, None]:
        """POST to SSE endpoint and yield content chunks."""
        token = await self._get_token()
        t = timeout or self._config.poll_timeout
        async with httpx.AsyncClient(timeout=httpx.Timeout(t, connect=10)) as client:
            async with client.stream(
                "POST",
                f"{self._config.base_url}{path}",
                json=payload,
                headers=self._headers(token),
            ) as resp:
                if resp.status_code == 401:
                    self._auth.invalidate(self._config.api_key)
                _raise_for_status(resp)

                async for line in resp.aiter_lines():
                    if not line.startswith("data: "):
                        continue
                    data_str = line[6:].strip()
                    if data_str == "[DONE]":
                        return

                    try:
                        event = json.loads(data_str)
                    except json.JSONDecodeError:
                        continue

                    chunk = _extract_sse_content(event)
                    if chunk is not None:
                        yield chunk

    async def consume_sse(
        self,
        path: str,
        payload: dict[str, Any],
        timeout: float | None = None,
    ) -> str:
        """POST to SSE endpoint and return the full aggregated content."""
        chunks: list[str] = []
        async for chunk in self.stream_sse(path, payload, timeout):
            chunks.append(chunk)
        return "".join(chunks)


def _extract_sse_content(event: dict[str, Any]) -> str | None:
    """Extract content from an SSE event."""
    if event.get("status") in FAILED_STATUSES:
        error_msg = event.get("error") or event.get("message") or "Stream error"
        raise MelogenError(f"SSE stream error: {error_msg}")

    if "content" in event:
        return event["content"]
    if "text" in event:
        return event["text"]
    if "result" in event:
        result = event["result"]
        return result if isinstance(result, str) else json.dumps(result, ensure_ascii=False)
    return None
