"""Pure-Python helpers for the spens mitmproxy addon.

These functions have no mitmproxy dependency so they can be unit-tested
without installing mitmproxy.
"""

from __future__ import annotations

import fnmatch
import json
from datetime import UTC, datetime
from pathlib import Path
from typing import Any
from urllib.parse import urlsplit

REDACTED = "***REDACTED***"


def timestamp() -> str:
    return datetime.now(UTC).isoformat().replace("+00:00", "Z")


def parse_json(data: bytes | str | None) -> Any:
    if data is None:
        return None
    if isinstance(data, bytes):
        try:
            data = data.decode("utf-8")
        except UnicodeDecodeError:
            return None
    try:
        return json.loads(data)
    except (json.JSONDecodeError, ValueError):
        return None


def parse_sse(content: bytes) -> list[dict[str, Any]]:
    """Parse a complete SSE buffer into a list of JSON objects."""
    chunks: list[dict[str, Any]] = []
    try:
        text = content.decode("utf-8", errors="replace")
    except Exception:
        return chunks
    for line in text.split("\n"):
        line = line.strip()
        if line.startswith("data: "):
            data = line[6:]
            if data == "[DONE]":
                continue
            parsed = parse_json(data)
            if isinstance(parsed, dict):
                chunks.append(parsed)
    return chunks


class StreamingSSEParser:
    """Incremental SSE parser that handles chunks split across boundaries."""

    def __init__(self) -> None:
        # Bytes, not str: a multi-byte UTF-8 sequence can be split across
        # chunk boundaries, and per-chunk str decoding would replace each
        # half with U+FFFD.  Only complete lines are decoded, as a whole.
        self._buffer: bytes = b""

    def feed(self, chunk: bytes) -> list[dict[str, Any]]:
        self._buffer += chunk

        events: list[dict[str, Any]] = []
        lines = self._buffer.split(b"\n")
        self._buffer = lines[-1]

        for line in lines[:-1]:
            line = line.strip()
            if not line.startswith(b"data: "):
                continue
            data = line[6:]
            if data == b"[DONE]":
                continue
            parsed = parse_json(data)
            if isinstance(parsed, dict):
                events.append(parsed)

        return events


def decode_stream_bytes(raw: bytes, content_encoding: str | None) -> bytes:
    """Best-effort decompression of a captured response body.

    A client that advertises ``accept-encoding: gzip/br/zstd`` can make the
    upstream compress its (SSE) responses; the streaming tee then forwards
    compressed bytes that the SSE parser cannot read, silently losing the
    whole response trace.  The declared encoding is tried first, then the
    other common ones; if nothing decodes, the bytes are returned unchanged
    (they may simply be uncompressed).
    """
    import gzip
    import zlib

    def _try_zstd(data: bytes) -> bytes | None:
        try:
            import zstandard  # type: ignore[import-not-found]
        except ImportError:
            return None
        try:
            return zstandard.ZstdDecompressor().stream_reader(data).read()  # type: ignore[union-attr]
        except Exception:
            return None

    def _try_brotli(data: bytes) -> bytes | None:
        try:
            import brotli  # type: ignore[import-not-found]
        except ImportError:
            return None
        try:
            return brotli.decompress(data)  # type: ignore[union-attr]
        except Exception:
            return None

    declared = (content_encoding or "").strip().lower()
    order: list[str] = []
    if declared and declared != "identity":
        order.append(declared)
    order += [e for e in ("gzip", "br", "zstd", "deflate") if e not in order]

    for enc in order:
        try:
            if enc in ("gzip", "x-gzip"):
                return gzip.decompress(raw)
            if enc == "deflate":
                try:
                    return zlib.decompress(raw)
                except zlib.error:
                    return zlib.decompress(raw, -15)
            if enc == "br":
                out = _try_brotli(raw)
                if out is not None:
                    return out
            if enc == "zstd":
                out = _try_zstd(raw)
                if out is not None:
                    return out
        except Exception:
            continue
    return raw


def matches(url: str, pattern: str) -> bool:
    return fnmatch.fnmatch(url, pattern)


def hostname(url: str) -> str:
    """Extract the lower-cased hostname from a URL.

    Falls back to the raw string if the URL cannot be parsed (e.g. it is
    already a bare hostname).  A bare ``host:port`` string (an HTTP
    authority-form target, as mitmproxy presents CONNECT requests) is also
    handled: ``urlsplit`` mis-parses such strings (it treats the host as a
    URI scheme and yields no hostname), so the port is stripped manually.
    Bare IPv6 literals (``[::1]:22`` / ``::1``) keep their colons and simply
    fail to match any domain pattern -- fail closed.
    """
    try:
        host = urlsplit(url).hostname
    except ValueError:
        host = None
    if host:
        return host.lower()
    raw = url.lower()
    if raw.count(":") == 1:
        maybe_host, _, port = raw.rpartition(":")
        if maybe_host and port.isdigit():
            return maybe_host
    return raw


def domain_matches(host: str, pattern: str) -> bool:
    """Match a domain pattern against an already-extracted hostname.

    This is spens' single definition of the domain-policy matching rule; it
    gates HTTP egress (via :func:`host_matches`) and DNS resolution (via the
    DNS forwarder's ``name_allowed``) alike, so the two channels cannot drift
    apart.

    Leading/trailing wildcards and dots are stripped from the pattern, and the
    resulting domain matches only itself or its subdomains.  For example a
    pattern of ``*.anthropic.com``, ``*anthropic.com*`` or ``anthropic.com``
    matches ``anthropic.com`` and ``api.anthropic.com`` but **not**
    ``api.anthropic.com.attacker.example`` (the domain must be a hostname
    *suffix at a label boundary*, not merely a substring), nor
    ``evil-anthropic.com``.

    A bare ``*`` pattern matches every hostname.  Wildcards inside the pattern
    (e.g. ``api.*.com``) are not supported and fail closed.
    """
    host = (host or "").strip().lower().rstrip(".")
    domain = pattern.strip().lower().strip("*").strip(".")
    if not domain:
        # A lone wildcard ("*" or "**") authorizes any hostname.
        return True
    return host == domain or host.endswith("." + domain)


def host_matches(url: str, pattern: str) -> bool:
    """Match a domain pattern against the hostname in ``url``.

    Matching only the hostname (rather than the full URL) prevents path- or
    query-based bypasses such as ``https://evil.com/?x=api.anthropic.com``.
    The pattern semantics are :func:`domain_matches`.
    """
    return domain_matches(hostname(url), pattern)


def mask_secrets(text: str | None, secrets: list[str] | set[str] | None) -> str | None:
    """Replace any real secret value occurrences in ``text`` with a redaction marker.

    Used before writing URLs (which may carry a substituted secret in a query
    parameter) to disk, so the real secret never lands in a log or trace file.
    """
    if not text or not secrets:
        return text
    for secret in secrets:
        if secret:
            text = text.replace(secret, REDACTED)
    return text


def should_capture(
    url: str,
    patterns: list[str],
    exclude_patterns: list[str] | None = None,
) -> bool:
    if exclude_patterns:
        for p in exclude_patterns:
            if matches(url, p):
                return False
    return any(matches(url, p) for p in patterns)


def is_secret_authorized(url: str, for_domains: list[str] | None) -> bool:
    """Return True if ``url`` matches at least one authorized domain pattern.

    Secret substitution rules use ``for_domains`` to scope where a placeholder
    may be replaced with the real secret value.  Patterns are domain-suffix
    patterns matched against the request's hostname only (e.g.
    ``*api.anthropic.com*`` matches ``api.anthropic.com`` and any subdomain of
    it), so a domain smuggled into a URL path, query, or as a prefix of an
    unrelated hostname (``api.anthropic.com.attacker.example``) cannot
    authorize substitution.

    An absent or empty ``for_domains`` list authorizes **no** URLs: by default
    a placeholder is never substituted, so a real secret cannot leak to an
    unexpected destination.
    """
    if not for_domains:
        return False
    return any(host_matches(url, pattern) for pattern in for_domains)


def host_allowed(url: str, domain_rules: list[dict[str, Any]]) -> bool:
    """Return whether *any* configured rule matches the URL's hostname.

    Used for CONNECT tunnel targets, where no HTTP method exists yet: the
    target hostname must match at least one rule's ``pattern`` (fail closed
    when rules are configured but none match).  Method-level enforcement is
    not applied here -- it still applies to every HTTP request sent through
    the tunnel via :func:`is_allowed` in the ``request`` hook.  When no rules
    are configured at all there is no policy to enforce, so the target is
    allowed (same semantics as ``is_allowed``).
    """
    if not domain_rules:
        return True
    return any(host_matches(url, rule.get("pattern", "")) for rule in domain_rules)


def is_allowed(url: str, method: str, domain_rules: list[dict[str, Any]]) -> bool:
    """Return whether ``method`` on ``url`` is permitted by the domain policy.

    Fails **closed**: if no rule matches the request's hostname the request is
    denied.  When no rules are configured at all there is no policy to enforce,
    so requests are allowed.  Patterns are domain-suffix patterns matched
    against the hostname only (see ``host_matches``).
    """
    if not domain_rules:
        return True
    for rule in domain_rules:
        if host_matches(url, rule.get("pattern", "")):
            allowed = rule.get("allow", ["*"])
            allowed_upper = {m.upper() for m in allowed}
            return "*" in allowed_upper or method.upper() in allowed_upper
    # No rule matched this hostname -- fail closed.
    return False


READY_MARKER_FILENAME = "spens_addon_ready"


def write_ready_marker(directory: Path) -> Path:
    """Write the addon readiness marker file and return its path.

    The marker lives in the shared cert volume so the spens runner (via
    ``docker exec``) and the agent entrypoint (via the read-only ``/certs``
    mount) can block until the addon is loaded, configured, and the proxy
    is actually listening -- instead of guessing from the CA cert file,
    which appears before the addon takes over port 9090.
    """
    directory.mkdir(parents=True, exist_ok=True)
    marker = directory / READY_MARKER_FILENAME
    marker.write_text(timestamp() + "\n", encoding="utf-8")
    return marker


class TraceWriter:
    """Buffers trace entries and flushes them in a single file operation."""

    def __init__(self, trace_file: Path) -> None:
        self._trace_file = trace_file
        self._buffer: list[str] = []

    def add(self, entry: dict[str, Any]) -> None:
        self._buffer.append(json.dumps(entry))

    def flush(self) -> None:
        if not self._buffer:
            return
        with open(self._trace_file, "a", encoding="utf-8") as fh:
            fh.write("\n".join(self._buffer) + "\n")
        self._buffer.clear()
