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

    from vela import Vela
    vela = Vela(api_key="sk-vela-...")
    print(vela.ask("Say hi"))

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

__all__ = ["Vela", "VelaError", "create_pkce", "authorize_url", "exchange_code"]
__version__ = "1.0.0"
DEFAULT_BASE = "https://velastocks.online/api"


class VelaError(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 VelaError(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 VelaError(msg, e.code, err if isinstance(err, str) else None, j or raw) from None


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

    def __init__(self, api_key, base_url=DEFAULT_BASE, timeout=120):
        if not api_key:
            raise VelaError("api_key is required. Create one at https://velastocks.online/dashboard?tab=dev")
        self.api_key, self.base, self.timeout, self.last = api_key, _base(base_url), timeout, None

    def _req(self, path, body=None):
        data = json.dumps(body).encode() if body is not None else None
        req = urllib.request.Request(self.base + path, data=data, method="POST" if body is not None else "GET",
                                     headers={"Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json", "User-Agent": f"vela-python/{__version__}"})
        return _open(req, self.timeout)

    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-vela-cost")
            self.last = {"model": r.headers.get("x-vela-model"), "tier": r.headers.get("x-vela-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 vela.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())


# ---------------------------------------------------------------- Sign in with Vela (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://velastocks.online"):
    """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"vela-python/{__version__}"})
    with _open(req, timeout) as r:
        return json.loads(r.read())