"""Waterlight SDK client — OpenAI-compatible, zero external dependencies.

Drop-in replacement for the OpenAI Python SDK:
    # Before
    from openai import OpenAI
    client = OpenAI(api_key="sk-...")

    # After
    from waterlight import Waterlight
    client = Waterlight(api_key="wl-...")

All methods mirror the OpenAI SDK interface:
    client.chat.completions.create(...)
    client.embeddings.create(...)
    client.models.list()
"""

from __future__ import annotations

import json
import os
import time
from dataclasses import dataclass, field
from typing import Any, Iterator, Optional
from urllib import error as urlerror
from urllib import request

from .errors import APIError, WaterlightError, _map_status


# =============================================================================
# Response Types
# =============================================================================

@dataclass
class Message:
    """Chat message."""
    role: str
    content: str | None = None
    tool_calls: list[dict] | None = None

@dataclass
class Choice:
    """Single completion choice."""
    index: int
    message: Message
    finish_reason: str | None = None

@dataclass
class Usage:
    """Token usage."""
    prompt_tokens: int = 0
    completion_tokens: int = 0
    total_tokens: int = 0

@dataclass
class ChatCompletion:
    """Non-streaming chat completion response."""
    id: str = ""
    object: str = "chat.completion"
    created: int = 0
    model: str = ""
    choices: list[Choice] = field(default_factory=list)
    usage: Usage = field(default_factory=Usage)

@dataclass
class Delta:
    """Streaming delta."""
    content: str | None = None
    tool_calls: list[dict] | None = None

@dataclass
class StreamChoice:
    """Single streaming choice."""
    index: int = 0
    delta: Delta = field(default_factory=Delta)
    finish_reason: str | None = None

@dataclass
class ChatCompletionChunk:
    """Streaming chunk."""
    id: str = ""
    object: str = "chat.completion.chunk"
    created: int = 0
    model: str = ""
    choices: list[StreamChoice] = field(default_factory=list)

@dataclass
class Embedding:
    """Single embedding result."""
    object: str = "embedding"
    embedding: list[float] = field(default_factory=list)
    index: int = 0

@dataclass
class EmbeddingResponse:
    """Embeddings response."""
    object: str = "list"
    data: list[Embedding] = field(default_factory=list)
    model: str = ""
    usage: Usage = field(default_factory=Usage)

@dataclass
class Model:
    """Model info."""
    id: str = ""
    object: str = "model"
    created: int = 0
    owned_by: str = "waterlight"

@dataclass
class ModelList:
    """Models list response."""
    object: str = "list"
    data: list[Model] = field(default_factory=list)


# =============================================================================
# Resource Classes
# =============================================================================

class Completions:
    """Chat completions resource — mirrors openai.chat.completions."""

    def __init__(self, client: Waterlight):
        self._client = client

    def create(
        self,
        *,
        model: str = "mist-1-turbo",
        messages: list[dict[str, Any]],
        temperature: float | None = None,
        max_tokens: int | None = None,
        stream: bool = False,
        tools: list[dict] | None = None,
        tool_choice: str | dict | None = None,
        **kwargs: Any,
    ) -> ChatCompletion | Iterator[ChatCompletionChunk]:
        """Create a chat completion.

        Args:
            model: Model to use (mist-1, mist-1-turbo, mist-1-flash, etc.)
            messages: List of message dicts with 'role' and 'content'
            temperature: Sampling temperature (0-2)
            max_tokens: Maximum tokens to generate
            stream: If True, returns an iterator of chunks
            tools: OpenAI-format tool definitions
            tool_choice: Tool choice strategy

        Returns:
            ChatCompletion if stream=False, Iterator[ChatCompletionChunk] if stream=True
        """
        payload: dict[str, Any] = {
            "model": model,
            "messages": messages,
            "stream": stream,
        }
        if temperature is not None:
            payload["temperature"] = temperature
        if max_tokens is not None:
            payload["max_tokens"] = max_tokens
        if tools is not None:
            payload["tools"] = tools
        if tool_choice is not None:
            payload["tool_choice"] = tool_choice
        payload.update(kwargs)

        if stream:
            return self._stream(payload)
        return self._complete(payload)

    def _complete(self, payload: dict) -> ChatCompletion:
        body = self._client._post("/v1/chat/completions", payload)
        choices = []
        for c in body.get("choices", []):
            msg = c.get("message", {})
            choices.append(Choice(
                index=c.get("index", 0),
                message=Message(
                    role=msg.get("role", "assistant"),
                    content=msg.get("content"),
                    tool_calls=msg.get("tool_calls"),
                ),
                finish_reason=c.get("finish_reason"),
            ))
        usage_data = body.get("usage", {})
        return ChatCompletion(
            id=body.get("id", ""),
            created=body.get("created", 0),
            model=body.get("model", ""),
            choices=choices,
            usage=Usage(
                prompt_tokens=usage_data.get("prompt_tokens", 0),
                completion_tokens=usage_data.get("completion_tokens", 0),
                total_tokens=usage_data.get("total_tokens", 0),
            ),
        )

    def _stream(self, payload: dict) -> Iterator[ChatCompletionChunk]:
        for event_data in self._client._post_stream("/v1/chat/completions", payload):
            if event_data == "[DONE]":
                return
            try:
                data = json.loads(event_data)
            except json.JSONDecodeError:
                continue
            choices = []
            for c in data.get("choices", []):
                delta_data = c.get("delta", {})
                choices.append(StreamChoice(
                    index=c.get("index", 0),
                    delta=Delta(
                        content=delta_data.get("content"),
                        tool_calls=delta_data.get("tool_calls"),
                    ),
                    finish_reason=c.get("finish_reason"),
                ))
            yield ChatCompletionChunk(
                id=data.get("id", ""),
                created=data.get("created", 0),
                model=data.get("model", ""),
                choices=choices,
            )


class Chat:
    """Chat namespace — mirrors openai.chat."""

    def __init__(self, client: Waterlight):
        self.completions = Completions(client)


class Embeddings:
    """Embeddings resource — mirrors openai.embeddings."""

    def __init__(self, client: Waterlight):
        self._client = client

    def create(
        self,
        *,
        input: str | list[str],
        model: str | None = None,
        encoding_format: str = "float",
    ) -> EmbeddingResponse:
        """Create embeddings for the given input."""
        payload: dict[str, Any] = {"input": input, "encoding_format": encoding_format}
        if model is not None:
            payload["model"] = model
        body = self._client._post("/v1/embeddings", payload)
        data = []
        for item in body.get("data", []):
            data.append(Embedding(
                embedding=item.get("embedding", []),
                index=item.get("index", 0),
            ))
        usage_data = body.get("usage", {})
        return EmbeddingResponse(
            data=data,
            model=body.get("model", ""),
            usage=Usage(
                prompt_tokens=usage_data.get("prompt_tokens", 0),
                completion_tokens=usage_data.get("completion_tokens", 0),
                total_tokens=usage_data.get("total_tokens", 0),
            ),
        )


class Models:
    """Models resource — mirrors openai.models."""

    def __init__(self, client: Waterlight):
        self._client = client

    def list(self) -> ModelList:
        """List available models."""
        body = self._client._get("/v1/models")
        models = []
        for item in body.get("data", []):
            models.append(Model(
                id=item.get("id", ""),
                created=item.get("created", 0),
                owned_by=item.get("owned_by", "waterlight"),
            ))
        return ModelList(data=models)


class Billing:
    """Billing resource — Waterlight-specific (not in OpenAI SDK)."""

    def __init__(self, client: Waterlight):
        self._client = client

    def get(self) -> dict:
        """Get billing info for the authenticated key.

        Returns a dict with plan, billing_mode, spent_usd, balance_usd (PAYG),
        remaining_usd (subscription), daily_limit, daily_used, etc.
        """
        return self._client._get("/v1/billing")


# =============================================================================
# Main Client
# =============================================================================

class Waterlight:
    """Waterlight API client — OpenAI-compatible interface.

    Usage:
        client = Waterlight(api_key="wl-...")

        # Chat completion
        response = client.chat.completions.create(
            model="mist-1-turbo",
            messages=[{"role": "user", "content": "Hello"}],
        )
        print(response.choices[0].message.content)

        # Streaming
        stream = client.chat.completions.create(
            model="mist-1-turbo",
            messages=[{"role": "user", "content": "Tell me a story"}],
            stream=True,
        )
        for chunk in stream:
            if chunk.choices[0].delta.content:
                print(chunk.choices[0].delta.content, end="", flush=True)

        # Embeddings
        emb = client.embeddings.create(input="Hello world")
        print(len(emb.data[0].embedding))
    """

    DEFAULT_BASE_URL = "https://api.waterlight.io"

    def __init__(
        self,
        api_key: str | None = None,
        base_url: str | None = None,
        timeout: float = 120.0,
        max_retries: int = 2,
    ):
        self.api_key = api_key or os.getenv("WATERLIGHT_API_KEY")
        if not self.api_key:
            raise ValueError(
                "API key required. Pass api_key= or set WATERLIGHT_API_KEY env var. "
                "Get your key at https://waterlight.io"
            )
        resolved_url = (base_url or os.getenv("WATERLIGHT_BASE_URL") or self.DEFAULT_BASE_URL).rstrip("/")
        if not resolved_url.startswith(("https://", "http://localhost", "http://127.0.0.1")):
            raise ValueError(
                f"base_url must use HTTPS (got: {resolved_url[:40]}...). "
                "HTTP is only allowed for localhost development."
            )
        self.base_url = resolved_url
        self._timeout = timeout
        self._max_retries = max_retries

        self.chat = Chat(self)
        self.embeddings = Embeddings(self)
        self.models = Models(self)
        self.billing = Billing(self)

    def __repr__(self) -> str:
        masked = f"{self.api_key[:6]}...{self.api_key[-4:]}" if len(self.api_key) > 10 else "***"
        return f"Waterlight(api_key='{masked}', base_url='{self.base_url}')"

    def _headers(self) -> dict[str, str]:
        return {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json",
            "User-Agent": "waterlight-python/0.2.1",
        }

    def _post(self, path: str, payload: dict) -> dict:
        """POST JSON and return parsed response body."""
        data = json.dumps(payload).encode("utf-8")
        req = request.Request(
            url=f"{self.base_url}{path}",
            data=data,
            headers=self._headers(),
            method="POST",
        )
        return self._execute(req)

    def _get(self, path: str) -> dict:
        """GET and return parsed response body."""
        req = request.Request(
            url=f"{self.base_url}{path}",
            headers=self._headers(),
            method="GET",
        )
        return self._execute(req)

    def _post_stream(self, path: str, payload: dict) -> Iterator[str]:
        """POST and yield SSE event data strings."""
        data = json.dumps(payload).encode("utf-8")
        headers = self._headers()
        headers["Accept"] = "text/event-stream"
        req = request.Request(
            url=f"{self.base_url}{path}",
            data=data,
            headers=headers,
            method="POST",
        )
        try:
            with request.urlopen(req, timeout=self._timeout) as resp:
                for raw_line in resp:
                    line = raw_line.decode("utf-8").rstrip("\r\n")
                    if line.startswith("data: "):
                        yield line[6:]
        except urlerror.HTTPError as exc:
            self._handle_http_error(exc)

    def _execute(self, req: request.Request) -> dict:
        """Execute request with retry on transient failures."""
        attempt = 0
        while True:
            try:
                with request.urlopen(req, timeout=self._timeout) as resp:
                    return json.loads(resp.read().decode("utf-8") or "{}")
            except urlerror.HTTPError as exc:
                status = exc.code
                if status in (429, 500, 502, 503, 504) and attempt < self._max_retries:
                    retry_after = exc.headers.get("Retry-After") if exc.headers else None
                    delay = float(retry_after) if retry_after else 0.5 * (2 ** attempt)
                    time.sleep(delay)
                    attempt += 1
                    continue
                self._handle_http_error(exc)
            except urlerror.URLError as exc:
                raise APIError(f"Network error: {exc.reason}", status_code=503)

    def _handle_http_error(self, exc: urlerror.HTTPError) -> None:
        """Parse error body and raise appropriate exception."""
        body = ""
        try:
            body = exc.read().decode("utf-8") if exc.fp else ""
        except Exception:
            pass
        message = "Request failed"
        request_id = None
        retry_after = None
        if body:
            try:
                data = json.loads(body)
                if isinstance(data, dict):
                    err = data.get("error", {})
                    if isinstance(err, dict):
                        message = err.get("message", message)
                    elif isinstance(err, str):
                        message = err
                    elif "message" in data:
                        message = data["message"]
            except json.JSONDecodeError:
                message = body
        if exc.headers:
            request_id = exc.headers.get("X-Request-Id")
            ra = exc.headers.get("Retry-After")
            if ra:
                try:
                    retry_after = float(ra)
                except ValueError:
                    pass
        raise _map_status(exc.code, message, request_id=request_id, retry_after=retry_after)
