"""Raeto SDK for Python 3.8+. Standard library only.

    from raeto import Raeto
    raeto = Raeto(api_key="sk-raeto-...")
    print(raeto.ask("Say hi"))

Agents can use their own wallet (an eth_account LocalAccount) instead of a shared key:
    raeto = Raeto.agent(account, name="research-bot", limit_usd=5)
or pay per request from the wallet, no account needed:
    raeto = Raeto.pay_per_request(account, tx_hash)  # USDG sent to Raeto, or rCREDIT redeemed to the wallet

Docs: https://raeto.cc/developers
"""
import base64
import hashlib
import json
import secrets
import urllib.error
import urllib.parse
import urllib.request

__all__ = ["Raeto", "RaetoError", "create_pkce", "authorize_url", "exchange_code"]
__version__ = "1.0.0"
DEFAULT_BASE = "https://api.raeto.cc"


class RaetoError(Exception):
    def __init__(self, message, status=None, type=None, body=None):
        super().__init__(message)
        self.status, self.type, self.body = status, type, body


def _base(url):
    url = url.rstrip("/")
    return url[:-3] if url.endswith("/v1") else url


def _open(req, timeout):
    try:
        return urllib.request.urlopen(req, timeout=timeout)
    except urllib.error.HTTPError as e:
        raw = e.read().decode("utf-8", "replace")
        try:
            j = json.loads(raw)
        except ValueError:
            j = None
        err = j.get("error") if isinstance(j, dict) else None
        if isinstance(err, dict):
            raise RaetoError(err.get("message") or "Request failed", e.code, err.get("type"), j) from None
        msg = (j or {}).get("error_description") or err or raw or f"Request failed ({e.code})"
        raise RaetoError(msg, e.code, err if isinstance(err, str) else None, j or raw) from None


def _sign(account, message):
    try:
        from eth_account.messages import encode_defunct
    except ImportError:
        raise RaetoError("Wallet sign-in needs eth-account: pip install eth-account") from None
    sig = account.sign_message(encode_defunct(text=message)).signature.hex()
    return sig if sig.startswith("0x") else "0x" + sig


class Raeto:
    """api_key: a key from Developers, or an access token from Sign in with Raeto."""

    def __init__(self, api_key=None, base_url=DEFAULT_BASE, timeout=120, _payment=None):
        if not api_key and not _payment:
            raise RaetoError("api_key is required. Create one at https://raeto.cc/dashboard?tab=dev, or use Raeto.agent(account).")
        self.api_key, self.base, self.timeout, self.last = api_key, _base(base_url), timeout, None
        self._payment, self.payment_left_usd, self.address, self.fund = _payment, None, None, None

    def _req(self, path, body=None):
        data = json.dumps(body).encode() if body is not None else None
        auth = {"X-Payment-Tx": self._payment[0], "X-Payment-Sig": self._payment[1]} if self._payment else {"Authorization": f"Bearer {self.api_key}"}
        req = urllib.request.Request(self.base + path, data=data, method="POST" if body is not None else "GET",
                                     headers={**auth, "Content-Type": "application/json", "User-Agent": f"raeto-python/{__version__}"})
        r = _open(req, self.timeout)
        left = r.headers.get("x-raeto-payment-left")
        if left is not None:
            self.payment_left_usd = float(left)
        return r

    @classmethod
    def agent(cls, account, name=None, limit_usd=None, models=None, base_url=DEFAULT_BASE, timeout=120):
        """Sign in as an agent with its own wallet (eth_account LocalAccount). The agent gets its own account,
        balance, limits and history. The key can call models but cannot move the account's credits."""
        base = _base(base_url)
        def post(path, body):
            req = urllib.request.Request(base + path, data=json.dumps(body).encode(), method="POST",
                                         headers={"Content-Type": "application/json", "User-Agent": f"raeto-python/{__version__}"})
            with _open(req, timeout) as r:
                return json.loads(r.read())
        ch = post("/v1/agents/challenge", {"address": account.address})
        body = {"address": account.address, "signature": _sign(account, ch["message"])}
        if name: body["name"] = name
        if limit_usd is not None: body["limitUsd"] = limit_usd
        if models: body["models"] = list(models)
        r = post("/v1/agents/login", body)
        raeto = cls(r["apiKey"], base_url, timeout)
        raeto.address, raeto.fund = r["address"], r.get("fund")
        return raeto

    @classmethod
    def pay_per_request(cls, account, tx_hash, base_url=DEFAULT_BASE, timeout=120):
        """No account or key: pay from the wallet first (USDG to the address in Raeto.how_to_pay(), or rCREDIT
        redeemed into the wallet), then pass that transaction hash. Calls are paid from it until it is used up
        (see payment_left_usd); anything left stays in the wallet's own account."""
        tx = str(tx_hash or "").lower()
        if len(tx) != 66 or not tx.startswith("0x"):
            raise RaetoError("tx_hash must be the payment transaction hash.")
        raeto = cls(None, base_url, timeout, _payment=(tx, _sign(account, f"Pay raeto with transaction {tx}")))
        raeto.address = account.address.lower()
        return raeto

    @staticmethod
    def how_to_pay(base_url=DEFAULT_BASE, timeout=30):
        """Where and how to pay: treasury address, USDG and rCREDIT contracts, chain id."""
        with _open(urllib.request.Request(_base(base_url) + "/v1/agents", headers={"User-Agent": f"raeto-python/{__version__}"}), timeout) as r:
            return json.loads(r.read())

    def chat(self, **params):
        """Full OpenAI-style chat completion. Returns the response dict."""
        params["stream"] = False
        with self._req("/v1/chat/completions", params) as r:
            j = json.loads(r.read())
            cost = r.headers.get("x-raeto-cost")
            self.last = {"model": r.headers.get("x-raeto-model"), "tier": r.headers.get("x-raeto-tier"), "cost_usd": float(cost) if cost else None}
            return j

    def ask(self, prompt, model=None, system=None, **rest):
        """Ask one question, get the reply text."""
        messages = ([{"role": "system", "content": system}] if system else []) + [{"role": "user", "content": prompt}]
        if model:
            rest["model"] = model
        j = self.chat(messages=messages, **rest)
        c = j["choices"][0]["message"]["content"]
        return c if isinstance(c, str) else json.dumps(c)

    def stream(self, **params):
        """Yield text pieces as they arrive: for t in raeto.stream(messages=[...]): print(t, end="")"""
        params["stream"] = True
        with self._req("/v1/chat/completions", params) as r:
            for raw in r:
                line = raw.decode("utf-8", "replace").strip()
                if not line.startswith("data:"):
                    continue
                d = line[5:].strip()
                if d == "[DONE]":
                    return
                try:
                    t = json.loads(d)["choices"][0]["delta"].get("content")
                except (ValueError, KeyError, IndexError):
                    continue
                if t:
                    yield t

    def models(self):
        """All models, with prices."""
        with self._req("/v1/models") as r:
            return json.loads(r.read())["data"]

    def me(self):
        """Who this key belongs to and how much it can still spend."""
        with self._req("/v1/me") as r:
            return json.loads(r.read())

    def route(self, **request):
        """Routing only: ask Raeto which model fits this request, then call that model yourself with your own
        provider key. Pass the chat request you are about to send (messages=..., tools=...). Optional models=
        names your own model per tier, e.g. {"fast": "gpt-4o-mini", "mid": "gpt-4.1", "top": "o3"}, and floor=
        ("mid" or "top") sets a minimum. Returns {"model", "tier", "reasons", "fallbacks", "fee_usd", ...}.
        Costs a small flat fee per call."""
        with self._req("/v1/route", request) as r:
            return json.loads(r.read())


# ---------------------------------------------------------------- Sign in with Raeto (OAuth 2.0 + PKCE)
def create_pkce():
    """Returns (verifier, challenge). Keep the verifier; send the challenge in the sign-in link."""
    verifier = base64.urlsafe_b64encode(secrets.token_bytes(32)).rstrip(b"=").decode()
    challenge = base64.urlsafe_b64encode(hashlib.sha256(verifier.encode()).digest()).rstrip(b"=").decode()
    return verifier, challenge


def authorize_url(client_id, redirect_uri, state, code_challenge, site="https://raeto.cc"):
    """The link to send users to. `state` is any random string you check when they come back."""
    q = urllib.parse.urlencode({"response_type": "code", "client_id": client_id, "redirect_uri": redirect_uri, "state": state,
                                "code_challenge": code_challenge, "code_challenge_method": "S256"})
    return f"{site.rstrip('/')}/authorize?{q}"


def exchange_code(client_id, code, code_verifier, redirect_uri, client_secret=None, base_url=DEFAULT_BASE, timeout=30):
    """Swap the `code` from the redirect for an access token. Returns {access_token, token_type, user_id, limit_usd}."""
    form = {"grant_type": "authorization_code", "client_id": client_id, "code": code, "code_verifier": code_verifier, "redirect_uri": redirect_uri}
    if client_secret:
        form["client_secret"] = client_secret
    req = urllib.request.Request(_base(base_url) + "/oauth/token", data=urllib.parse.urlencode(form).encode(), method="POST",
                                 headers={"Content-Type": "application/x-www-form-urlencoded", "User-Agent": f"raeto-python/{__version__}"})
    with _open(req, timeout) as r:
        return json.loads(r.read())