"""Spens mitmproxy addon.

Provides four capabilities for the interceptor sidecar:

1. **LLM trace capture** -- captures completion requests and responses (including
   streaming SSE chunks) to JSONL trace files, compatible with the spens log viewer.
   SSE responses are streamed to the agent in real-time via a tee callable while
   chunks are simultaneously captured for auditing.
2. **Request log** -- logs every request (method, URL, status) without bodies.
3. **Secret substitution** -- the agent is given a *placeholder* string (e.g. the
   literal value ``"ANTHROPIC_API_KEY"``) as its API key.  When the agent sends a
   request, the placeholder appears in the Authorization header or in a URL
   query parameter.  The interceptor replaces the placeholder with the real
   secret value pulled from an environment variable, so the real secret never
   enters the agent container.  Each rule declares ``for_domains`` (domain
   patterns such as ``*api.anthropic.com*``, matched with suffix semantics:
   the domain itself or its subdomains) and the placeholder is only
   replaced on requests whose *hostname* matches those patterns -- a
   placeholder sent to any other destination is left unreplaced, so a real
   secret can never be exfiltrated to an unauthorized host.  Any real secret
   value is masked (redacted) out of URLs before they are written to log or
   trace files, so a secret substituted into a query parameter never lands on
   disk.
4. **Domain filtering** -- enforces per-hostname method whitelist rules and
   *fails closed*: a request whose hostname matches no rule is blocked.
   CONNECT tunnel targets are domain-checked the same way (fail closed)
   before the tunnel is established, and every CONNECT decision is written
   to the request log.  The interceptor also runs mitmproxy with
   ``rawtcp=false`` so non-HTTP protocols inside a tunnel (SSH etc.) are
   blocked rather than blindly forwarded.

Configuration is read from ``/app/spens_interceptor_config.json`` (mounted at runtime).
"""

from __future__ import annotations

import json
import os
import time
import uuid
from pathlib import Path
from typing import Any

from mitmproxy import ctx, http
from mitmproxy_helpers import (
    StreamingSSEParser,
    TraceWriter,
    decode_stream_bytes,
    host_allowed,
    is_allowed,
    is_secret_authorized,
    mask_secrets,
    parse_json,
    parse_sse,
    should_capture,
    timestamp,
    write_ready_marker,
)

CONFIG_PATH = "/app/spens_interceptor_config.json"
TRACES_DIR = "/app/traces"
# The readiness marker is written into the cert volume (/root/.mitmproxy in
# the interceptor, mounted read-only at /certs in the agent) so both the
# runner and the agent entrypoint can wait on it.
READY_MARKER_DIR = "/root/.mitmproxy"
TRACE_FILENAME = "captured.jsonl"
REQUEST_LOG_FILENAME = "request_log.jsonl"

DEFAULT_CAPTURE_PATTERNS = [
    "*api.openai.com*",
    "*api.anthropic.com*",
]


class SpensAddon:
    """mitmproxy addon implementing spens interception features."""

    def __init__(self) -> None:
        self.config: dict[str, Any] = {}
        self.capture_patterns: list[str] = list(DEFAULT_CAPTURE_PATTERNS)
        self.exclude_patterns: list[str] = []
        self.domain_rules: list[dict[str, Any]] = []
        self.inject_headers_cfg: list[dict[str, Any]] = []
        self.traces_dir = Path(TRACES_DIR)
        self.trace_file: Path = Path()
        self.request_log_file: Path = Path()
        self._trace_writer: TraceWriter | None = None
        self._warned_missing_env: set[str] = set()
        self._warned_unauthorized: set[tuple[str, str]] = set()
        self._secret_values: set[str] = set()

    def load(self, loader) -> None:
        loader.add_option(
            "spens_config",
            str,
            CONFIG_PATH,
            "Path to spens interceptor config JSON",
        )

    def configure(self, updates: set[str]) -> None:
        config_path = ctx.options.spens_config
        self.config = self._load_config(config_path)
        self.capture_patterns = list(DEFAULT_CAPTURE_PATTERNS)
        self.capture_patterns.extend(self.config.get("addition_capture_urls", []))
        self.exclude_patterns = self.config.get("exclude_capture_urls", [])
        self.domain_rules = self.config.get("domain_rules", [])
        self.inject_headers_cfg = self.config.get("inject_headers", [])
        # Collect the real secret values so they can be masked out of any URL
        # (a substituted secret may end up in a query parameter) before it is
        # written to a log or trace file on disk.
        self._secret_values = {
            os.environ.get(rule.get("env_var", ""), "")
            for rule in self.inject_headers_cfg
            if rule.get("env_var") and os.environ.get(rule.get("env_var", ""))
        }
        self.traces_dir = Path(self.config.get("traces_dir", TRACES_DIR))
        self.traces_dir.mkdir(parents=True, exist_ok=True)
        self.trace_file = self.traces_dir / TRACE_FILENAME
        self.request_log_file = self.traces_dir / REQUEST_LOG_FILENAME
        self._trace_writer = TraceWriter(self.trace_file)
        ctx.log.info(
            f"spens addon configured: {len(self.capture_patterns)} capture patterns, "
            f"{len(self.exclude_patterns)} exclude patterns, "
            f"{len(self.domain_rules)} domain rules, "
            f"{len(self.inject_headers_cfg)} secret substitutions"
        )
        for rule in self.inject_headers_cfg:
            env_var = rule.get("env_var", "")
            placeholder = rule.get("placeholder", "")
            if env_var and not os.environ.get(env_var, ""):
                ctx.log.warn(
                    f"spens: inject_headers rule (placeholder '{placeholder}') references env "
                    f"var '{env_var}' which is NOT set in this interceptor container -- the "
                    "placeholder will NOT be substituted. Check the 'env_var' name in "
                    ".spens.config.json and that it is exported in the shell running spens."
                )
            if placeholder and not rule.get("for_domains"):
                ctx.log.warn(
                    f"spens: inject_headers rule (placeholder '{placeholder}') has no "
                    "'for_domains' patterns -- secrets are only substituted on authorized "
                    "URLs, so this placeholder will NEVER be substituted. Add a for_domains "
                    "list (e.g. [\"*api.anthropic.com*\"]) to the rule in .spens.config.json."
                )

    def running(self) -> None:
        """Signal readiness once mitmproxy is fully up and listening.

        This fires only after the addon is loaded and configured and the
        proxy server is listening, so anyone waiting on the marker knows
        interception (capture, secret substitution, domain rules) works.
        """
        try:
            marker = write_ready_marker(Path(READY_MARKER_DIR))
            ctx.log.info(f"spens addon ready, wrote readiness marker: {marker}")
        except OSError as exc:
            ctx.log.warn(f"spens: could not write readiness marker: {exc}")

    def _load_config(self, path: str) -> dict[str, Any]:
        try:
            with open(path, encoding="utf-8") as fh:
                return json.load(fh)
        except (OSError, json.JSONDecodeError):
            ctx.log.warn(f"spens: could not load config from {path}")
            return {}

    # -- helpers --------------------------------------------------------

    def _substitute_secrets(self, flow: http.HTTPFlow) -> None:
        """Replace placeholder strings with real secret values.

        For each rule the agent was told its key is ``placeholder`` (a literal
        string).  That placeholder appears in request headers and/or URL query
        parameters.  We replace every occurrence with the real value from the
        environment variable named ``env_var`` -- but only when the request
        URL matches one of the rule's ``for_domains`` glob patterns.  This
        scopes the secret to authorized destinations: if an agent is tricked
        into sending the placeholder to an unexpected host, the real secret is
        never attached to that request.
        """
        url = flow.request.pretty_url
        for rule in self.inject_headers_cfg:
            placeholder = rule.get("placeholder", "")
            env_var = rule.get("env_var", "")
            if not placeholder or not env_var:
                continue

            for_domains = rule.get("for_domains", [])
            if not is_secret_authorized(url, for_domains):
                # The URL is not authorized for this secret.  Check whether the
                # placeholder actually appears in the request and, if so, warn
                # once per placeholder so the operator can spot the leak
                # attempt (or the missing for_domains pattern) in the logs.
                if self._placeholder_present(flow, placeholder):
                    warn_key = (placeholder, url)
                    if warn_key not in self._warned_unauthorized:
                        self._warned_unauthorized.add(warn_key)
                        ctx.log.warn(
                            f"spens: placeholder '{placeholder}' found in request to "
                            f"'{url}', which is not authorized by the rule's "
                            "for_domains patterns -- the placeholder will be sent "
                            "upstream unreplaced"
                        )
                continue

            real_value = os.environ.get(env_var, "")
            if not real_value:
                if env_var not in self._warned_missing_env:
                    self._warned_missing_env.add(env_var)
                    ctx.log.warn(
                        f"spens: skipping secret substitution for placeholder '{placeholder}': "
                        f"env var '{env_var}' is not set in the interceptor; the placeholder "
                        "will be sent upstream unreplaced"
                    )
                continue

            for key in list(flow.request.headers.keys()):
                old_val = flow.request.headers[key]
                if placeholder in old_val:
                    flow.request.headers[key] = old_val.replace(placeholder, real_value)

            for qkey in list(flow.request.query.keys()):
                old_val = flow.request.query[qkey]
                if placeholder in old_val:
                    flow.request.query[qkey] = old_val.replace(placeholder, real_value)

    @staticmethod
    def _placeholder_present(flow: http.HTTPFlow, placeholder: str) -> bool:
        """Return True if ``placeholder`` occurs in any header or query value."""
        return (
            any(placeholder in value for value in flow.request.headers.values())
            or any(placeholder in value for value in flow.request.query.values())
        )

    def _safe_url(self, flow: http.HTTPFlow) -> str:
        """Return the request URL with any real secret values redacted."""
        return mask_secrets(flow.request.pretty_url, self._secret_values)

    def _log_request(self, flow: http.HTTPFlow) -> None:
        self._log_entry(
            flow.request.method,
            self._safe_url(flow),
            flow.response.status_code if flow.response else None,
        )

    def _log_entry(self, method: str, url: str, status_code: int | None) -> None:
        """Append one entry to the request log (no bodies, secrets masked)."""
        entry = {
            "timestamp": timestamp(),
            "method": method,
            "url": url,
            "status_code": status_code,
        }
        with open(self.request_log_file, "a", encoding="utf-8") as fh:
            fh.write(json.dumps(entry) + "\n")

    # -- mitmproxy hooks ------------------------------------------------

    def http_connect(self, flow: http.HTTPFlow) -> None:
        """Domain-check CONNECT targets *before* a tunnel is established.

        mitmproxy only shows a CONNECT tunnel's *inner* traffic to the
        ``request`` hook when it is HTTP.  Without this hook a CONNECT to
        any host would be tunneled blindly: a non-HTTP payload (e.g. SSH,
        which mitmproxy's layer selector routes to raw TCP forwarding when
        the ``rawtcp`` option is on) would reach the upstream with no domain
        check and no request-log entry.  We therefore enforce the domain
        policy here, fail closed, and log every CONNECT decision.

        Only the *hostname* is checked (via ``host_allowed``): the CONNECT
        method itself is not compared against each rule's ``allow`` list,
        because the methods that matter are those of the HTTP requests sent
        through the tunnel, which are still checked in ``request``.
        Denial works by setting a non-2xx response on the flow: mitmproxy
        then returns it to the client and never opens the tunnel.
        """
        # CONNECT requests carry the target as an authority-form URL
        # ("host:port"), which ``host_allowed``/``hostname`` handle.
        target = flow.request.pretty_url
        if not host_allowed(target, self.domain_rules):
            ctx.log.warn(
                f"spens: CONNECT to non-allowed destination denied: {target}"
            )
            flow.response = http.Response.make(
                403,
                b'{"error": "blocked by spens domain policy"}',
                {"Content-Type": "application/json"},
            )
            self._log_entry("CONNECT", mask_secrets(target, self._secret_values), 403)
        else:
            # mitmproxy itself answers the CONNECT with 200 once the hook
            # returns without a response; record the decision either way so
            # tunnels are auditable (the connect response does not pass
            # through the ``response`` hook).
            self._log_entry("CONNECT", mask_secrets(target, self._secret_values), 200)

    def request(self, flow: http.HTTPFlow) -> None:
        url = flow.request.pretty_url
        method = flow.request.method

        if not is_allowed(url, method, self.domain_rules):
            flow.response = http.Response.make(
                403,
                b'{"error": "blocked by spens domain policy"}',
                {"Content-Type": "application/json"},
            )
            return

        self._substitute_secrets(flow)

        if should_capture(url, self.capture_patterns, self.exclude_patterns):
            rid = uuid.uuid4().hex[:16]
            flow.metadata["spens_rid"] = rid
            flow.metadata["spens_start"] = time.time()

            # Ask the upstream for an *uncompressed* body.  The streaming tee
            # parses the SSE bytes it forwards; a gzip/br/zstd-compressed
            # event stream would be opaque to it (and to the audit trail),
            # silently losing every response event.  Clients handle an
            # identity response fine regardless of what they advertised.
            flow.request.headers["accept-encoding"] = "identity"

            body = parse_json(flow.request.get_text())
            self._trace_writer.add({
                "type": "request",
                "id": rid,
                "timestamp": timestamp(),
                "method": method,
                "url": self._safe_url(flow),
                "body": body,
            })
            self._trace_writer.flush()

    def responseheaders(self, flow: http.HTTPFlow) -> None:
        """Set response streaming mode based on capture status and content type.

        - Non-captured URLs: stream directly (no buffering, no processing).
        - Captured SSE responses: stream via a tee callable that captures
          chunks while forwarding them to the agent immediately.
        - Captured non-SSE responses: buffer for full body capture.
        """
        url = flow.request.pretty_url

        if not should_capture(url, self.capture_patterns, self.exclude_patterns):
            flow.response.stream = True
            return

        rid = flow.metadata.get("spens_rid")
        if not rid:
            flow.response.stream = True
            return

        content_type = flow.response.headers.get("content-type", "")
        is_sse = "text/event-stream" in content_type

        if is_sse:
            parser = StreamingSSEParser()
            trace_buffer: list[dict[str, Any]] = []
            # Raw bytes as forwarded (possibly still content-encoded): kept
            # so ``response``/``error`` can retry parsing with decompression
            # if the live parse produced nothing.
            raw_chunks: list[bytes] = []

            def tee(chunk: bytes) -> bytes:
                raw_chunks.append(chunk)
                events = parser.feed(chunk)
                for event in events:
                    trace_buffer.append({
                        "type": "response_chunk",
                        "request_id": rid,
                        "timestamp": timestamp(),
                        "content": event,
                    })
                return chunk

            flow.response.stream = tee
            flow.metadata["spens_trace_buffer"] = trace_buffer
            flow.metadata["spens_raw_chunks"] = raw_chunks
        else:
            flow.response.stream = False

    def _sse_entries(
        self, flow: http.HTTPFlow, rid: str
    ) -> list[dict[str, Any]]:
        """Return the captured SSE chunk rows for a finished stream.

        Primary source is the tee's live-parsed event buffer.  If that is
        empty -- the stream arrived content-encoded despite the identity
        ``accept-encoding`` we send, or the tee never saw parseable bytes --
        fall back to re-parsing the accumulated raw bytes after best-effort
        decompression, and warn loudly when even that yields nothing, so a
        capture gap is visible in the interceptor log instead of only as a
        silent zero-usage row in the session summary.
        """
        trace_buffer = flow.metadata.get("spens_trace_buffer")
        if trace_buffer:
            return list(trace_buffer)

        raw = b"".join(flow.metadata.get("spens_raw_chunks") or [])
        if not raw:
            # Streaming responses may not retain a body on the flow, but try
            # whatever is there before giving up.
            raw = flow.response.content or b""
        if not raw:
            self._warn_no_sse_events(flow, "no body bytes were seen by the capture tee")
            return []

        encoding = flow.response.headers.get("content-encoding", "")
        decoded = decode_stream_bytes(raw, encoding)
        entries = [
            {
                "type": "response_chunk",
                "request_id": rid,
                "timestamp": timestamp(),
                "content": event,
            }
            for event in parse_sse(decoded)
        ]
        if entries:
            return entries
        self._warn_no_sse_events(
            flow,
            f"{len(raw)} body bytes did not parse as SSE "
            f"(content-encoding: {encoding or 'identity'})",
        )
        return []

    def _warn_no_sse_events(self, flow: http.HTTPFlow, detail: str) -> None:
        ctx.log.warn(
            f"spens: captured NO SSE events for {self._safe_url(flow)} -- {detail}. "
            "The response trace for this request is lost; please report this "
            "so the capture path can be fixed."
        )

    def response(self, flow: http.HTTPFlow) -> None:
        self._log_request(flow)

        rid = flow.metadata.get("spens_rid")
        if not rid:
            return

        start = flow.metadata.get("spens_start", time.time())
        latency_ms = (time.time() - start) * 1000

        content_type = flow.response.headers.get("content-type", "")
        is_sse = "text/event-stream" in content_type

        if is_sse:
            for entry in self._sse_entries(flow, rid):
                self._trace_writer.add(entry)
            self._trace_writer.add({
                "type": "response_meta",
                "request_id": rid,
                "timestamp": timestamp(),
                "total_latency_ms": latency_ms,
                "status_code": flow.response.status_code,
            })
            self._trace_writer.flush()
        else:
            body = parse_json(flow.response.get_text())
            self._trace_writer.add({
                "type": "response",
                "request_id": rid,
                "timestamp": timestamp(),
                "status_code": flow.response.status_code,
                "body": body,
                "latency_ms": latency_ms,
            })
            self._trace_writer.flush()

    def error(self, flow: http.HTTPFlow) -> None:
        self._log_request(flow)

        rid = flow.metadata.get("spens_rid")
        if not rid:
            return

        # The client or upstream dropped mid-stream: keep whatever events
        # were captured (and retry the raw bytes), so a cancelled request
        # still leaves a partial trace rather than none at all.
        if flow.response is not None and "text/event-stream" in (
            flow.response.headers.get("content-type", "")
        ):
            for entry in self._sse_entries(flow, rid):
                self._trace_writer.add(entry)
            self._trace_writer.flush()


addons = [SpensAddon()]
