# SPDX-License-Identifier: AGPL-3.0-only
# Copyright (c) 2026 sol pbc

"""Resident loopback proxy for `sol link serve`."""

from __future__ import annotations

import argparse
import errno
import json
import logging
import queue
import sys
from collections.abc import Iterable
from concurrent.futures import Future
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path
from typing import NamedTuple, Protocol, cast
from urllib.parse import urlsplit

from solstone.think.link.bundle import load_client_identity
from solstone.think.link.client import RECOMMENDED_CHUNK, BodySource
from solstone.think.link.dialer import (
    TunnelClient,
    TunnelLifecycleError,
    TunnelRequestError,
    TunnelResponseHead,
)
from solstone.think.link.observer_paths import observer_bundle_dir, observer_spl_root
from solstone.think.link.paths import relay_url

LOG = logging.getLogger(__name__)
LOOPBACK_HOST = "127.0.0.1"
DEFAULT_PORT = 5015
RESPONSE_QUEUE_MAX_ITEMS = 8
RESPONSE_HEAD_TIMEOUT_SECONDS = 30.0
STATUS_PATH = "/_solstone/link/status"
REQUEST_HOP_BY_HOP = {
    "connection",
    "keep-alive",
    "proxy-authenticate",
    "proxy-authorization",
    "te",
    "trailers",
    "transfer-encoding",
    "upgrade",
}
RESPONSE_HOP_BY_HOP = {
    "connection",
    "keep-alive",
    "proxy-authenticate",
    "proxy-authorization",
    "te",
    "trailers",
    "upgrade",
}


class _ProxyTunnel(Protocol):
    def proxy_stream_request(
        self,
        method: str,
        path: str,
        *,
        headers: dict[str, str] | None = None,
        body: bytes | BodySource = b"",
        chunks: queue.Queue[TunnelResponseHead | bytes | Exception | None],
    ) -> Future[None]: ...

    def status(self) -> dict[str, object]: ...


class _BundleSelection(NamedTuple):
    label: str
    bundle_dir: Path


class _ProxyServer(ThreadingHTTPServer):
    tunnel: _ProxyTunnel


def add_arguments(parser: argparse.ArgumentParser) -> None:
    parser.add_argument("--label", help="Observer link bundle label")
    parser.add_argument(
        "--port",
        type=int,
        default=DEFAULT_PORT,
        help=f"Loopback port to serve on (default: {DEFAULT_PORT})",
    )
    parser.add_argument("--relay-url", help="Override the spl relay URL")
    parser.add_argument(
        "--direct",
        action="store_true",
        help="PL-direct only: dial the journal over the LAN secure listener, "
        "never the spl relay. Use when the home is reachable directly "
        "(same LAN/VPN) to avoid any relay dependency.",
    )


def main(args: argparse.Namespace) -> int:
    port_error = _port_error(args.port)
    if port_error is not None:
        return _fail(port_error, code=2)

    try:
        selection = _resolve_bundle_dir(args.label)
        identity = load_client_identity(selection.bundle_dir)
    except ValueError as exc:
        return _fail(str(exc), code=1)

    relay = (
        None if getattr(args, "direct", False) else _resolve_relay_url(args.relay_url)
    )
    tunnel = TunnelClient(identity, relay)
    tunnel.start()
    try:
        server = _build_server(args.port, tunnel)
    except OSError as exc:
        tunnel.close()
        return _fail(_bind_error(args.port, exc), code=1)

    LOG.info(
        "forwarding %s:%d -> home %s over pl",
        LOOPBACK_HOST,
        args.port,
        selection.label,
    )
    try:
        server.serve_forever()
    except KeyboardInterrupt:
        return 0
    finally:
        server.shutdown()
        server.server_close()
        tunnel.close()
    return 0


def _resolve_bundle_dir(label: str | None) -> _BundleSelection:
    if label:
        bundle_dir = observer_bundle_dir(label)
        try:
            load_client_identity(bundle_dir)
        except ValueError as exc:
            raise ValueError(
                f"invalid link bundle for label '{label}' at {bundle_dir}: {exc}. "
                "Run `sol link join` to pair this device."
            ) from exc
        return _BundleSelection(label, bundle_dir)

    bundles = _valid_observer_bundles()
    if not bundles:
        raise ValueError(
            f"no observer link bundles found under {observer_spl_root()}. "
            "Run `sol link join` to pair this device."
        )
    if len(bundles) > 1:
        labels = ", ".join(sorted(bundles))
        raise ValueError(
            f"multiple observer link bundles found: {labels}. "
            "Pass --label to choose one."
        )
    selected_label, bundle_dir = next(iter(bundles.items()))
    return _BundleSelection(selected_label, bundle_dir)


def _valid_observer_bundles() -> dict[str, Path]:
    root = observer_spl_root()
    if not root.is_dir():
        return {}
    bundles: dict[str, Path] = {}
    for bundle_dir in sorted(root.iterdir()):
        if not bundle_dir.is_dir():
            continue
        try:
            load_client_identity(bundle_dir)
        except ValueError:
            continue
        bundles[bundle_dir.name] = bundle_dir
    return bundles


def _resolve_relay_url(value: str | None) -> str:
    if value and value.strip():
        return value.strip().rstrip("/")
    return relay_url()


def _build_server(port: int, tunnel: _ProxyTunnel) -> _ProxyServer:
    server = _ProxyServer((LOOPBACK_HOST, port), _ProxyHandler)
    server.tunnel = tunnel
    return server


def _port_error(port: int) -> str | None:
    if port < 1 or port > 65535:
        return "--port must be between 1 and 65535"
    return None


def _bind_error(port: int, exc: OSError) -> str:
    if exc.errno == errno.EADDRINUSE:
        return (
            f"cannot bind {LOOPBACK_HOST}:{port}: address already in use. "
            "Another `sol link serve` or Convey may already be using that port."
        )
    return f"cannot bind {LOOPBACK_HOST}:{port}: {exc}"


def _fail(message: str, *, code: int) -> int:
    print(message, file=sys.stderr)
    return code


class _ProxyHandler(BaseHTTPRequestHandler):
    protocol_version = "HTTP/1.1"

    def do_GET(self) -> None:
        self._handle()

    def do_POST(self) -> None:
        self._handle()

    def do_PUT(self) -> None:
        self._handle()

    def do_DELETE(self) -> None:
        self._handle()

    def do_PATCH(self) -> None:
        self._handle()

    def do_HEAD(self) -> None:
        self._handle()

    def do_OPTIONS(self) -> None:
        self._handle()

    def _handle(self) -> None:
        if _origin_path(self.path) == STATUS_PATH:
            self._send_status()
            return

        body = self._read_body_source()
        if body is None:
            return

        chunks: queue.Queue[TunnelResponseHead | bytes | Exception | None] = (
            queue.Queue(maxsize=RESPONSE_QUEUE_MAX_ITEMS)
        )
        server = cast(_ProxyServer, self.server)
        future = server.tunnel.proxy_stream_request(
            self.command,
            _origin_path(self.path),
            headers=_forward_request_headers(self.headers.items(), body.length),
            body=body,
            chunks=chunks,
        )

        try:
            first = chunks.get(timeout=RESPONSE_HEAD_TIMEOUT_SECONDS)
        except queue.Empty:
            future.cancel()
            self._send_gateway_error(TimeoutError("response head timed out"))
            return

        if isinstance(first, Exception):
            self._send_gateway_error(first)
            return
        if first is None:
            future.cancel()
            self._send_gateway_error(RuntimeError("empty response from link tunnel"))
            return
        if not isinstance(first, TunnelResponseHead):
            future.cancel()
            self._send_gateway_error(RuntimeError("bad response from link tunnel"))
            return

        self.send_response_only(first.status)
        for name, value in first.headers.items():
            if name.lower() in RESPONSE_HOP_BY_HOP:
                continue
            self.send_header(name, value)
        self.end_headers()

        write_body = self.command != "HEAD"
        while True:
            item = chunks.get()
            if item is None:
                return
            if isinstance(item, Exception):
                future.cancel()
                LOG.warning("link proxy stream failed after response head: %s", item)
                self.close_connection = True
                return
            if not isinstance(item, bytes):
                future.cancel()
                LOG.warning("link proxy stream returned unexpected item: %r", item)
                self.close_connection = True
                return
            if not write_body:
                continue
            try:
                self.wfile.write(item)
                self.wfile.flush()
            except (BrokenPipeError, ConnectionError, OSError):
                future.cancel()
                self.close_connection = True
                return

    def _read_body_source(self) -> BodySource | None:
        if self.headers.get("Transfer-Encoding") is not None:
            LOG.warning("link proxy rejected request with Transfer-Encoding")
            self._send_bad_request("Transfer-Encoding is not supported")
            self.close_connection = True
            return None

        content_length = self.headers.get("Content-Length")
        if content_length is None:
            return BodySource(0, ())
        try:
            length = int(content_length)
        except ValueError:
            self._send_bad_request("Bad Content-Length")
            self.close_connection = True
            return None
        if length < 0:
            self._send_bad_request("Bad Content-Length")
            self.close_connection = True
            return None
        return BodySource(length, self._body_chunks(length))

    def _body_chunks(self, length: int) -> Iterable[bytes]:
        remaining = length
        while remaining > 0:
            chunk = self.rfile.read(min(RECOMMENDED_CHUNK, remaining))
            if not chunk:
                raise ConnectionError("client closed request body early")
            remaining -= len(chunk)
            yield chunk

    def _send_bad_request(self, message: str) -> None:
        try:
            self.send_error(400, message)
        except (BrokenPipeError, ConnectionError, OSError) as exc:
            LOG.debug("link proxy could not send bad request response: %s", exc)

    def _send_gateway_error(self, exc: Exception) -> None:
        if isinstance(exc, TunnelLifecycleError):
            self._send_lifecycle_error(exc)
            return
        if isinstance(exc, TunnelRequestError):
            detail = exc.reason
        else:
            detail = type(exc).__name__
        try:
            self.send_error(502, f"Link tunnel failed: {detail}")
        except (BrokenPipeError, ConnectionError, OSError) as send_exc:
            LOG.debug("link proxy could not send gateway error response: %s", send_exc)
        self.close_connection = True

    def _send_lifecycle_error(self, exc: TunnelLifecycleError) -> None:
        body = json.dumps(exc.as_dict(), sort_keys=True).encode("utf-8")
        try:
            self.send_response_only(503)
            self.send_header("Content-Type", "application/json")
            self.send_header("Content-Length", str(len(body)))
            self.end_headers()
            if self.command != "HEAD":
                self.wfile.write(body)
                self.wfile.flush()
        except (BrokenPipeError, ConnectionError, OSError) as send_exc:
            LOG.debug(
                "link proxy could not send lifecycle error response: %s", send_exc
            )
        self.close_connection = True

    def _send_status(self) -> None:
        server = cast(_ProxyServer, self.server)
        status = server.tunnel.status()
        body = json.dumps(status, sort_keys=True).encode("utf-8")
        try:
            self.send_response_only(200)
            self.send_header("Content-Type", "application/json")
            self.send_header("Content-Length", str(len(body)))
            self.end_headers()
            if self.command != "HEAD":
                self.wfile.write(body)
                self.wfile.flush()
        except (BrokenPipeError, ConnectionError, OSError) as exc:
            LOG.debug("link proxy could not send status response: %s", exc)
        self.close_connection = True

    def log_message(self, fmt: str, *args: object) -> None:
        LOG.debug("%s - %s", self.address_string(), fmt % args)


def _origin_path(raw_path: str) -> str:
    if raw_path.startswith("http://") or raw_path.startswith("https://"):
        parsed = urlsplit(raw_path)
        path = parsed.path or "/"
        if parsed.query:
            return f"{path}?{parsed.query}"
        return path
    return raw_path or "/"


def _forward_request_headers(
    headers: Iterable[tuple[str, str]],
    body_length: int,
) -> dict[str, str]:
    forwarded: dict[str, str] = {}
    for name, value in headers:
        if name.lower() == "authorization":
            continue
        if name.lower() in REQUEST_HOP_BY_HOP:
            continue
        forwarded[name] = value
    forwarded["Content-Length"] = str(body_length)
    return forwarded
