"""Typed Python SDK for the meshnet OpenAI-compatible gateway."""

from __future__ import annotations

from dataclasses import dataclass
import json
import urllib.parse
import urllib.request
from typing import Any

import openai
from openai.types.chat import ChatCompletion


@dataclass(frozen=True)
class WalletBalance:
    """API-key balance in client-facing currencies."""

    sol: float
    usdc: float


@dataclass(frozen=True)
class TopUpQuote:
    """Payment details for funding an API key account."""

    payment_address: str
    qr: str
    amount_sol: float
    amount_usdc: float


@dataclass(frozen=True)
class CostEstimate:
    """Preflight cost estimate for a request."""

    model: str
    tokens: int
    lamports: int
    sol: float


@dataclass(frozen=True)
class RedundantRequestResult:
    """Result of sending a prompt through multiple independent routes."""

    redundancy: int
    majority_response: ChatCompletion
    responses: list[ChatCompletion]


class Client:
    """meshnet SDK client.

    ``chat`` is the official OpenAI SDK chat resource, configured against the
    meshnet gateway. Network-specific controls live on ``wallet``, ``models``,
    ``estimate_cost``, and ``request``.
    """

    def __init__(self, *, api_key: str, base_url: str = "https://api.meshnet.ai") -> None:
        self.api_key = api_key
        self.base_url = base_url.rstrip("/")
        self._openai = openai.OpenAI(
            api_key=api_key,
            base_url=_openai_base_url(self.base_url),
        )
        self.chat = self._openai.chat
        self.wallet = _WalletClient(self)
        self.models = _ModelsClient(self)

    def estimate_cost(self, *, model: str, tokens: int) -> CostEstimate:
        query = urllib.parse.urlencode({"model": model, "tokens": tokens})
        data = self._get_json(f"/v1/estimate_cost?{query}")
        return CostEstimate(
            model=str(data["model"]),
            tokens=int(data["tokens"]),
            lamports=int(data["lamports"]),
            sol=float(data["sol"]),
        )

    def request(
        self,
        *,
        model: str,
        messages: list[dict[str, Any]],
        redundancy: int = 1,
        **params: Any,
    ) -> RedundantRequestResult:
        payload: dict[str, Any] = {
            "model": model,
            "messages": messages,
            "redundancy": redundancy,
            **params,
        }
        data = self._post_json("/v1/request", payload)
        return RedundantRequestResult(
            redundancy=int(data["redundancy"]),
            majority_response=_chat_completion(data["majority_response"]),
            responses=[_chat_completion(item) for item in data["responses"]],
        )

    def _get_json(self, path: str) -> dict[str, Any]:
        req = urllib.request.Request(
            f"{self.base_url}{path}",
            headers=self._headers(),
            method="GET",
        )
        with urllib.request.urlopen(req, timeout=10) as resp:
            return json.loads(resp.read())

    def _post_json(self, path: str, payload: dict[str, Any]) -> dict[str, Any]:
        data = json.dumps(payload).encode()
        req = urllib.request.Request(
            f"{self.base_url}{path}",
            data=data,
            headers={**self._headers(), "Content-Type": "application/json"},
            method="POST",
        )
        with urllib.request.urlopen(req, timeout=30) as resp:
            return json.loads(resp.read())

    def _headers(self) -> dict[str, str]:
        return {"Authorization": f"Bearer {self.api_key}"}


class _WalletClient:
    def __init__(self, client: Client) -> None:
        self._client = client

    def balance(self) -> WalletBalance:
        data = self._client._get_json("/v1/wallet/balance")
        return WalletBalance(sol=float(data["sol"]), usdc=float(data["usdc"]))

    def top_up(self, *, amount_sol: float = 0.0, amount_usdc: float = 0.0) -> TopUpQuote:
        data = self._client._post_json(
            "/v1/wallet/top_up",
            {"amount_sol": amount_sol, "amount_usdc": amount_usdc},
        )
        return TopUpQuote(
            payment_address=str(data["payment_address"]),
            qr=str(data["qr"]),
            amount_sol=float(data["amount_sol"]),
            amount_usdc=float(data["amount_usdc"]),
        )


class _ModelsClient:
    def __init__(self, client: Client) -> None:
        self._client = client

    def available(self) -> list[dict[str, float | str]]:
        data = self._client._get_json("/v1/models/available")
        return [
            {
                "id": str(model["id"]),
                "shard_coverage_percentage": float(model["shard_coverage_percentage"]),
            }
            for model in data["data"]
        ]


def _openai_base_url(base_url: str) -> str:
    return base_url if base_url.endswith("/v1") else f"{base_url}/v1"


def _chat_completion(data: dict[str, Any]) -> ChatCompletion:
    return ChatCompletion.model_validate(data)