import codecs
import email.utils
import http.client
import json
import os
import re
import socket
import threading
import time
import urllib.error
import urllib.parse
import urllib.request
import uuid


RETRYABLE_STATUSES = {408, 425, 429, 500, 502, 503, 504}


class VibeCodeError(RuntimeError):
    def __init__(self, message, status=None, data=None):
        super().__init__(message)
        self.status = status
        self.data = data
        self.retry_after = 0


class VibeCodeTimeoutError(VibeCodeError):
    pass


class VibeCodeClient:
    def __init__(
        self,
        api_key=None,
        base_url="https://api.vibecodepro.ru",
        timeout=120,
        retries=2,
        retry_delay=0.25,
    ):
        self.api_key = api_key or os.environ.get("VIBECODE_API_KEY")
        if not self.api_key:
            raise VibeCodeError("VIBECODE_API_KEY is required")
        self.base_url = base_url.rstrip("/")
        self.timeout = timeout
        self.retries = retries
        self.retry_delay = retry_delay

    def request(self, path, body=None, method="POST", headers=None, timeout=None, retries=None):
        timeout = self.timeout if timeout is None else timeout
        retries = self.retries if retries is None else retries
        payload = None if method == "GET" else json.dumps(body or {}).encode("utf-8")
        request_headers = {
            "Authorization": "Bearer " + self.api_key,
            "Content-Type": "application/json",
            **(headers or {}),
        }
        if method != "GET" and retries > 0 and not _has_header(request_headers, "idempotency-key"):
            request_headers["Idempotency-Key"] = "vc_idem_" + str(uuid.uuid4())
        request = urllib.request.Request(
            self.base_url + path,
            data=payload,
            method=method,
            headers=request_headers,
        )
        last_error = None
        for attempt in range(retries + 1):
            try:
                with urllib.request.urlopen(request, timeout=timeout) as res:
                    text = res.read().decode("utf-8")
                    return _json_or_none(text)
            except urllib.error.HTTPError as exc:
                last_error = _http_error(exc)
                if not _should_retry_http(last_error, exc.code) or attempt >= retries:
                    raise last_error from exc
            except (socket.timeout, TimeoutError) as exc:
                last_error = VibeCodeTimeoutError("VibeCode request timed out")
                if attempt >= retries:
                    raise last_error from exc
            except urllib.error.URLError as exc:
                last_error = _url_error(exc)
                if attempt >= retries:
                    raise last_error from exc
            except (http.client.HTTPException, OSError) as exc:
                # R2-PUB-04: a dropped connection (RemoteDisconnected, IncompleteRead, ConnectionResetError,
                # BrokenPipeError, ssl.SSLError) is a transport error like any other: wrap and retry it.
                last_error = VibeCodeError("VibeCode connection failed: " + (str(exc) or type(exc).__name__))
                if attempt >= retries:
                    raise last_error from exc
            # R2-PUB-10: honour Retry-After (capped at 30 s) on 429/503 instead of the fixed backoff only.
            time.sleep(max(self.retry_delay * (2 ** attempt), getattr(last_error, "retry_after", 0) or 0))
        raise last_error or VibeCodeError("VibeCode request failed")

    def stream(self, path, body=None, headers=None, timeout=None, total_timeout=None):
        """Yield the parsed JSON payload of each server-sent event.

        `timeout` is the per-socket-operation timeout (connect / each read);
        `total_timeout` (seconds, optional) bounds the whole stream.
        """
        timeout = self.timeout if timeout is None else timeout
        deadline = None if total_timeout is None else time.monotonic() + total_timeout
        payload = json.dumps({**(body or {}), "stream": True}).encode("utf-8")
        request = urllib.request.Request(
            self.base_url + path,
            data=payload,
            method="POST",
            headers={
                "Authorization": "Bearer " + self.api_key,
                "Content-Type": "application/json",
                "Accept": "text/event-stream",
                **(headers or {}),
            },
        )
        try:
            # urlopen includes connecting and waiting for response headers. It must
            # receive the remaining budget, before the first response body read.
            try:
                response = _open_stream(request, timeout, deadline)
            except urllib.error.HTTPError as exc:
                with exc:
                    raise _http_error(exc, timeout=timeout, deadline=deadline) from exc
            with response as res:
                _stream_timeout(timeout, deadline)
                content_type = res.headers.get("Content-Type", "") or ""
                if re.search(r"application/json|text/html", content_type, re.I):
                    # A 200 that is not an event stream (JSON error body, HTML proxy page).
                    text = _stream_body(res, timeout, deadline).decode("utf-8", "replace")
                    data = _json_or_none(text) or {}
                    raise VibeCodeError(_error_message(data, res.status), status=res.status, data=data)
                parser = _SseParser()
                decoder = codecs.getincrementaldecoder("utf-8")(errors="replace")
                reader = getattr(res, "read1", None) or res.read
                sock = _response_socket(res)
                while True:
                    operation_timeout = _stream_timeout(timeout, deadline)
                    if sock is not None:
                        sock.settimeout(operation_timeout)
                    chunk = reader(65536)
                    _stream_timeout(timeout, deadline)
                    if not chunk:
                        for event in parser.finish():
                            if isinstance(event, VibeCodeError):
                                raise event
                            if event == "[DONE]":
                                return
                            yield event
                        return
                    # R2-PUB-10: events parsed before an in-band error in the same chunk are delivered first.
                    for event in parser.push(decoder.decode(chunk)):
                        if isinstance(event, VibeCodeError):
                            raise event
                        if event == "[DONE]":
                            return
                        yield event
        except (socket.timeout, TimeoutError) as exc:
            raise VibeCodeTimeoutError("VibeCode stream timed out") from exc
        except urllib.error.URLError as exc:
            # PUB-21: request() wrapped URLError; stream() leaked the raw exception.
            raise _url_error(exc) from exc
        except (http.client.HTTPException, OSError) as exc:
            # R2-PUB-04: a stream cut mid-body (IncompleteRead, reset) surfaces as VibeCodeError; never retried.
            raise VibeCodeError("VibeCode stream failed: " + (str(exc) or type(exc).__name__)) from exc

    def chat(self, messages, model="gpt-6.1-terra", **kwargs):
        return self.request("/v1/chat/completions", {"model": model, "messages": messages, **kwargs})

    def chat_stream(self, messages, model="gpt-6.1-terra", **kwargs):
        return self.stream("/v1/chat/completions", {"model": model, "messages": messages, **kwargs})

    def responses(self, input, model="gpt-6.1-sol", **kwargs):
        return self.request("/v1/responses", {"model": model, "input": input, **kwargs})

    def responses_stream(self, input, model="gpt-6.1-sol", **kwargs):
        return self.stream("/v1/responses", {"model": model, "input": input, **kwargs})

    def messages(self, messages, model="claude-sonnet-5-5", max_tokens=1024, **kwargs):
        return self.request("/v1/messages", {"model": model, "messages": messages, "max_tokens": max_tokens, **kwargs})

    def messages_stream(self, messages, model="claude-sonnet-5-5", max_tokens=1024, **kwargs):
        return self.stream("/v1/messages", {"model": model, "messages": messages, "max_tokens": max_tokens, **kwargs})

    def status(self):
        return self.request("/v1/status", method="GET")

    def models(self):
        return self.request("/v1/models", method="GET")

    def balance(self):
        return self.request("/v1/balance", method="GET")

    def usage(self, api_key_id=None, **query):
        # GW-14: the documented /v1/usage (any API key with usage:read). Query: limit, cursor, from_, to,
        # model, status, key_id; a positional value is a key id, as the former usage(api_key_id) took.
        if api_key_id is not None:
            query["key_id"] = api_key_id
        return self.request("/v1/usage" + _query_string(query), method="GET")

    def usage_breakdown(self, **query):
        return self.request("/v1/usage/breakdown" + _query_string(query), method="GET")

    def usage_chart(self, **query):
        return self.request("/v1/usage/chart" + _query_string(query), method="GET")

    def model_stats(self):
        return self.request("/v1/models/stats", method="GET")

    def billing(self, path="/billing/orders"):
        return self.request(path, method="GET")

    def doctor(self, tool="codex", model="gpt-6.1-terra"):
        query = urllib.parse.urlencode({"tool": tool, "model": model})
        return self.request(f"/setup/doctor?{query}", method="GET")


def _json_or_none(text):
    if not text:
        return None
    try:
        return json.loads(text)
    except Exception:
        return {"detail": text}


def _error_message(data, status):
    error = data.get("error") if isinstance(data.get("error"), dict) else {}
    message = str(data.get("detail") or error.get("message") or f"VibeCode HTTP {status}")
    # PUB-21: an HTML 502 page from a proxy must not become a megabyte-long message.
    return message if len(message) <= 300 else message[:300] + "\u2026"


def _http_error(exc, timeout=None, deadline=None):
    body = exc.read() if deadline is None else _stream_body(exc, timeout, deadline)
    text = body.decode("utf-8", "replace")
    data = _json_or_none(text) or {}
    error = VibeCodeError(_error_message(data, exc.code), status=exc.code, data=data)
    error.retry_after = _retry_after_seconds(exc.code, exc.headers)
    return error


def _url_error(exc):
    # R2-PUB-10: a connect timeout arrives as URLError(reason=timeout).
    if isinstance(exc.reason, (socket.timeout, TimeoutError)):
        return VibeCodeTimeoutError("VibeCode request timed out")
    return VibeCodeError(str(exc.reason or exc))


def _retry_after_seconds(status, headers):
    if status not in (429, 503) or headers is None:
        return 0
    value = headers.get("Retry-After")
    if not value:
        return 0
    try:
        seconds = float(value)
    except ValueError:
        try:
            seconds = email.utils.parsedate_to_datetime(value).timestamp() - time.time()
        except (TypeError, ValueError, IndexError, OverflowError):
            return 0
    return min(max(seconds, 0), 30)


def _stream_timeout(timeout, deadline):
    if deadline is None:
        return timeout
    remaining = deadline - time.monotonic()
    if remaining <= 0:
        raise VibeCodeTimeoutError("VibeCode stream timed out")
    return remaining if timeout is None else min(timeout, remaining)


# Only blocking OS DNS calls run in background threads. Capacity is shared across
# clients; timed-out callers never enqueue unlimited work or open late sockets.
_STREAM_DNS_SLOTS = threading.BoundedSemaphore(4)


def _stream_addresses(host, port, budget):
    wait = budget.remaining()
    if not _STREAM_DNS_SLOTS.acquire(timeout=wait):
        raise VibeCodeTimeoutError("VibeCode stream timed out")
    ready = threading.Event()
    result = {}

    def resolve():
        try:
            result["addresses"] = socket.getaddrinfo(host, port, 0, socket.SOCK_STREAM)
        except BaseException as error:
            result["error"] = error
        finally:
            _STREAM_DNS_SLOTS.release()
            ready.set()

    try:
        budget.remaining()
        worker = threading.Thread(target=resolve, name="VibeCodeDNS", daemon=True)
        worker.start()
    except BaseException:
        _STREAM_DNS_SLOTS.release()
        raise
    if not ready.wait(budget.remaining()):
        raise VibeCodeTimeoutError("VibeCode stream timed out")
    budget.remaining()  # An expired lookup must never proceed to connect/send.
    if "error" in result:
        raise result["error"]
    return result["addresses"]


class _StreamBudget:
    def __init__(self, timeout, deadline):
        self.timeout = timeout
        self.deadline = deadline
        self.expired = threading.Event()
        self.lock = threading.Lock()
        self.sockets = set()
        self.timer = threading.Timer(_stream_timeout(None, deadline), self.cancel)
        self.timer.name = "VibeCodeDeadline"
        self.timer.daemon = True

    def remaining(self):
        if self.expired.is_set():
            raise VibeCodeTimeoutError("VibeCode stream timed out")
        return _stream_timeout(self.timeout, self.deadline)

    def attach(self, sock):
        with self.lock:
            self.sockets.add(sock)
        self.prepare(sock)

    def prepare(self, sock):
        try:
            sock.settimeout(self.remaining())
        except BaseException:
            sock.close()
            raise

    def cancel(self):
        self.expired.set()
        with self.lock:
            sockets = list(self.sockets)
        for sock in sockets:
            # close alone does not interrupt HTTPResponse's makefile reader.
            try:
                sock.shutdown(socket.SHUT_RDWR)
            except OSError:
                pass
            sock.close()

    def stop(self):
        self.timer.cancel()
        self.timer.join()
        with self.lock:
            self.sockets.clear()


class _StreamConnection:
    def __init__(self, *args, budget, **kwargs):
        super().__init__(*args, **kwargs)
        self._budget = budget
        self._create_connection = self._connect_with_budget

    def _connect_with_budget(self, address, timeout=None, source_address=None):
        addresses = _stream_addresses(address[0], address[1], self._budget)
        last_error = None
        for family, kind, protocol, _, sockaddr in addresses:
            self._budget.remaining()
            sock = socket.socket(family, kind, protocol)
            try:
                self._budget.attach(sock)
                if source_address:
                    sock.bind(source_address)
                sock.connect(sockaddr)
                self._budget.prepare(sock)
                return sock
            except OSError as error:
                sock.close()
                self._budget.remaining()
                last_error = error
        raise last_error or OSError("getaddrinfo returned no stream addresses")

    def send(self, data):
        self._budget.remaining()
        if self.sock is None:
            if not self.auto_open:
                raise http.client.NotConnected()
            self.connect()
        # Connect, proxy tunnelling and TLS must finish before any request bytes.
        self._budget.prepare(self.sock)
        return super().send(data)


class _StreamHTTPConnection(_StreamConnection, http.client.HTTPConnection):
    pass


class _StreamHTTPSConnection(_StreamConnection, http.client.HTTPSConnection):
    def connect(self):
        http.client.HTTPConnection.connect(self)
        self._budget.prepare(self.sock)
        hostname = self._tunnel_host or self.host
        self.sock = self._context.wrap_socket(self.sock, server_hostname=hostname, do_handshake_on_connect=False)
        self._budget.attach(self.sock)
        self.sock.do_handshake()
        self._budget.prepare(self.sock)


class _StreamHTTPHandler(urllib.request.HTTPHandler):
    def __init__(self, budget):
        super().__init__()
        self.budget = budget

    def http_open(self, req):
        def connection(*args, **kwargs):
            return _StreamHTTPConnection(*args, budget=self.budget, **kwargs)
        return self.do_open(connection, req)


class _StreamHTTPSHandler(urllib.request.HTTPSHandler):
    def __init__(self, budget):
        super().__init__()
        self.budget = budget

    def https_open(self, req):
        def connection(*args, **kwargs):
            return _StreamHTTPSConnection(*args, budget=self.budget, **kwargs)
        options = {"context": self._context}
        if hasattr(self, "_check_hostname"):
            options["check_hostname"] = self._check_hostname
        return self.do_open(connection, req, **options)


class _StreamRedirectHandler(urllib.request.HTTPRedirectHandler):
    def redirect_request(self, req, fp, code, msg, headers, newurl):
        if urllib.parse.urlsplit(newurl).scheme not in ("http", "https"):
            fp.close()
            raise VibeCodeError("VibeCode stream redirect requires HTTP or HTTPS")
        return super().redirect_request(req, fp, code, msg, headers, newurl)


def _open_stream(request, timeout, deadline):
    if deadline is None:
        return urllib.request.urlopen(request, timeout=timeout)
    budget = _StreamBudget(timeout, deadline)
    budget.timer.start()
    response = None
    try:
        # A private opener preserves standard proxy/TLS handling without changing
        # the process-wide urllib opener or restarting the budget on redirects.
        opener = urllib.request.build_opener(_StreamHTTPHandler(budget), _StreamHTTPSHandler(budget), _StreamRedirectHandler())
        response = opener.open(request, timeout=budget.remaining())
        budget.remaining()
        return response
    except BaseException as error:
        timed_out = budget.expired.is_set() or time.monotonic() >= deadline
        # HTTPError owns a readable response body, which stream() consumes under
        # the remaining budget after the pre-body watchdog has been joined.
        if isinstance(error, urllib.error.HTTPError) and not timed_out:
            raise
        budget.cancel()
        if response is not None:
            response.close()
        if timed_out:
            if isinstance(error, urllib.error.HTTPError):
                error.close()
            raise VibeCodeTimeoutError("VibeCode stream timed out") from error
        raise
    finally:
        budget.stop()


def _stream_body(res, timeout, deadline):
    """Read a non-SSE response without allowing a drip body to extend the budget."""
    reader = getattr(res, "read1", None) or res.read
    sock = _response_socket(res)
    chunks = []
    while True:
        operation_timeout = _stream_timeout(timeout, deadline)
        if sock is not None:
            sock.settimeout(operation_timeout)
        chunk = reader(65536)
        _stream_timeout(timeout, deadline)
        if not chunk:
            return b"".join(chunks)
        chunks.append(chunk)


def _response_socket(res):
    fp = getattr(res, "fp", None)
    # HTTPError wraps an HTTPResponse in fp; a successful response exposes the
    # buffered socket directly. Both must receive each new remaining budget.
    raw = getattr(fp, "raw", None) or getattr(getattr(fp, "fp", None), "raw", None)
    sock = getattr(raw, "_sock", None)
    return sock if hasattr(sock, "settimeout") else None


def _utf8_len(text):
    return len(text.encode("utf-8", "surrogatepass"))


def _has_header(headers, name):
    normalized = name.lower()
    return any(str(key).lower() == normalized for key in (headers or {}))


def _query_string(query):
    # `from_` is spelled so because `from` is a Python keyword.
    pairs = {("from" if name == "from_" else name): value for name, value in query.items() if value is not None and value != ""}
    return "?" + urllib.parse.urlencode(pairs) if pairs else ""


def _should_retry_http(error, status):
    data = getattr(error, "data", None) or {}
    error_body = data.get("error") if isinstance(data.get("error"), dict) else {}
    code = data.get("code") or error_body.get("code")
    # A spending limit does not clear by waiting a few seconds.
    return code not in ("idempotency_key_reused", "spending_limit_exceeded") and status in RETRYABLE_STATUSES


class _SseParser:
    """Spec-conformant SSE parsing (PUB-19).

    Events end at a blank line; LF, CRLF and lone CR are line terminators; several
    ``data:`` lines of one event are joined with a newline; comments and unknown
    fields are ignored. Returns parsed JSON payloads (or the string "[DONE]"). An
    in-band error event or an oversized event is returned as a VibeCodeError after
    the events parsed before it (R2-PUB-10); parsing stops there.
    """

    _LINE_END = re.compile(r"\r\n|\r|\n")

    def __init__(self, max_event_bytes=2 * 1024 * 1024):
        self._buffer = ""
        self._data = []
        self._bytes = 0
        self._max = max_event_bytes
        self._failed = None

    def push(self, text):
        self._buffer += text
        return self._drain(final=False)

    def finish(self):
        return self._drain(final=True)

    def _drain(self, final):
        if self._failed is not None:
            return [self._failed]
        out = []
        while True:
            match = self._LINE_END.search(self._buffer)
            if not match:
                break
            # A trailing CR may be the first half of a CRLF split across chunks.
            if not final and match.group(0) == "\r" and match.start() == len(self._buffer) - 1:
                break
            line = self._buffer[: match.start()]
            self._buffer = self._buffer[match.end():]
            self._consume(line, out)
            if out and isinstance(out[-1], VibeCodeError):
                self._failed = out[-1]
                return out
        # Characters <= UTF-8 bytes <= 4 x characters: encode only near the limit.
        if len(self._buffer) > self._max or (len(self._buffer) * 4 > self._max and _utf8_len(self._buffer) > self._max):
            self._failed = VibeCodeError("VibeCode stream event exceeds the configured size limit")
            out.append(self._failed)
        return out

    def _consume(self, line, out):
        if line == "":
            event = self._dispatch()
            if event is not None:
                out.append(event)
            return
        if line.startswith(":"):
            return
        field, _, value = line.partition(":")
        if value.startswith(" "):
            value = value[1:]
        if field != "data":
            return
        # R2-PUB-10: the limit is in UTF-8 bytes, as documented (it was characters x 3).
        self._bytes += _utf8_len(value)
        if self._bytes > self._max:
            out.append(VibeCodeError("VibeCode stream event exceeds the configured size limit"))
            return
        self._data.append(value)

    def _dispatch(self):
        if not self._data:
            return None
        raw = "\n".join(self._data)
        self._data = []
        self._bytes = 0
        if raw == "[DONE]":
            return raw
        event = _json_or_none(raw)
        if isinstance(event, dict) and (isinstance(event.get("error"), dict) or event.get("type") in ("error", "response.failed")):
            detail = event.get("error") if isinstance(event.get("error"), dict) else (event.get("response") or {}).get("error") or {}
            return VibeCodeError(detail.get("message") or "VibeCode stream failed", data=event)
        return event