#!/usr/bin/env python3
"""Protocol-level OAuth, push, and SMTP mocks for the isolated local drill."""

from __future__ import annotations

import base64
import hashlib
import json
import os
import signal
import socketserver
import threading
import time
import urllib.parse
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from typing import Any


HTTP_PORT = 8080
SMTP_PORT = 2525
SMTP_RETRY_PORT = 2526
SMTP_TIMEOUT_PORT = 2527
DELAY_SECONDS = 0.35


class BoundaryState:
    def __init__(self) -> None:
        self.lock = threading.Lock()
        self.oauth_mode = "success"
        self.oauth_authorize_requests = 0
        self.oauth_token_requests = 0
        self.oauth_user_requests = 0
        self.oauth_code_counter = 0
        self.oauth_codes: dict[str, dict[str, Any]] = {}
        self.oauth_tokens: set[str] = set()
        self.push_mode = "success"
        self.push_attempts = 0
        self.push_receipts: dict[str, str] = {}
        self.push_deliveries = 0
        self.smtp_connections = {
            str(SMTP_PORT): 0,
            str(SMTP_RETRY_PORT): 0,
            str(SMTP_TIMEOUT_PORT): 0,
        }
        self.smtp_messages = 0
        self.smtp_rejections = 0
        self.smtp_digests: list[str] = []

    def control(self, component: str, mode: str) -> None:
        with self.lock:
            if component == "oauth":
                self.oauth_mode = mode
                self.oauth_authorize_requests = 0
                self.oauth_token_requests = 0
                self.oauth_user_requests = 0
                self.oauth_codes = {}
                self.oauth_tokens = set()
            elif component == "push":
                self.push_mode = mode
                self.push_attempts = 0
                self.push_receipts = {}
                self.push_deliveries = 0
            elif component == "smtp":
                self.smtp_connections = {
                    str(SMTP_PORT): 0,
                    str(SMTP_RETRY_PORT): 0,
                    str(SMTP_TIMEOUT_PORT): 0,
                }
                self.smtp_messages = 0
                self.smtp_rejections = 0
                self.smtp_digests = []
            else:
                raise ValueError("unsupported component")

    def snapshot(self) -> dict[str, Any]:
        with self.lock:
            return {
                "oauth": {
                    "mode": self.oauth_mode,
                    "authorize_requests": self.oauth_authorize_requests,
                    "token_requests": self.oauth_token_requests,
                    "user_requests": self.oauth_user_requests,
                    "issued_codes": len(self.oauth_codes),
                    "consumed_codes": sum(
                        1 for code in self.oauth_codes.values() if code["used"]
                    ),
                },
                "push": {
                    "mode": self.push_mode,
                    "attempts": self.push_attempts,
                    "deliveries": self.push_deliveries,
                    "idempotency_keys": len(self.push_receipts),
                },
                "smtp": {
                    "connections": dict(self.smtp_connections),
                    "messages": self.smtp_messages,
                    "rejections": self.smtp_rejections,
                    "message_digests": list(self.smtp_digests),
                },
            }


STATE = BoundaryState()
OAUTH_CLIENT_ID = os.environ["MOCK_OAUTH_CLIENT_ID"]
OAUTH_CLIENT_SECRET = os.environ["MOCK_OAUTH_CLIENT_SECRET"]
PUSH_BEARER_TOKEN = os.environ["MOCK_PUSH_BEARER_TOKEN"]


def json_bytes(value: Any) -> bytes:
    return json.dumps(value, separators=(",", ":"), sort_keys=True).encode()


class BoundaryHTTPServer(ThreadingHTTPServer):
    daemon_threads = True


class BoundaryHTTPHandler(BaseHTTPRequestHandler):
    server_version = "WhoNeedHelpBoundaryMock/1"

    def log_message(self, _format: str, *_args: Any) -> None:
        return

    def send_json(self, status: int, value: Any) -> None:
        payload = json_bytes(value)
        try:
            self.send_response(status)
            self.send_header("content-type", "application/json")
            self.send_header("content-length", str(len(payload)))
            self.end_headers()
            self.wfile.write(payload)
        except (BrokenPipeError, ConnectionResetError):
            return

    def read_body(self) -> bytes:
        length = int(self.headers.get("content-length", "0"))
        if length < 0 or length > 65_536:
            raise ValueError("invalid body length")
        return self.rfile.read(length)

    def do_GET(self) -> None:  # noqa: N802
        parsed = urllib.parse.urlsplit(self.path)

        if parsed.path == "/healthz":
            self.send_json(200, {"status": "ok"})
        elif parsed.path == "/state":
            self.send_json(200, STATE.snapshot())
        elif parsed.path == "/oauth/authorize":
            self.oauth_authorize(parsed)
        elif parsed.path == "/oauth/user":
            self.oauth_user()
        else:
            self.send_json(404, {"error": "not_found"})

    def do_POST(self) -> None:  # noqa: N802
        parsed = urllib.parse.urlsplit(self.path)

        try:
            body = self.read_body()
        except ValueError:
            self.send_json(413, {"error": "invalid_body"})
            return

        if parsed.path == "/control":
            self.control(body)
        elif parsed.path == "/oauth/token":
            self.oauth_token(body)
        elif parsed.path == "/push":
            self.push(body)
        else:
            self.send_json(404, {"error": "not_found"})

    def control(self, body: bytes) -> None:
        try:
            command = json.loads(body)
            component = command["component"]
            mode = command.get("mode", "success")
            STATE.control(component, mode)
        except (json.JSONDecodeError, KeyError, TypeError, ValueError):
            self.send_json(400, {"error": "invalid_control"})
            return

        self.send_json(200, {"status": "reset", "component": component, "mode": mode})

    def oauth_authorize(self, parsed: urllib.parse.SplitResult) -> None:
        params = urllib.parse.parse_qs(parsed.query)
        required = {
            "client_id",
            "redirect_uri",
            "state",
            "code_challenge",
            "code_challenge_method",
        }

        if not required.issubset(params) or params["client_id"][0] != OAUTH_CLIENT_ID:
            self.send_json(400, {"error": "invalid_authorization_request"})
            return

        if params["code_challenge_method"][0] != "S256":
            self.send_json(400, {"error": "unsupported_challenge_method"})
            return

        redirect_uri = params["redirect_uri"][0]
        state = params["state"][0]

        with STATE.lock:
            STATE.oauth_authorize_requests += 1
            mode = STATE.oauth_mode

            if mode != "deny_authorize":
                STATE.oauth_code_counter += 1
                code = f"boundary-code-{STATE.oauth_code_counter}"
                STATE.oauth_codes[code] = {
                    "challenge": params["code_challenge"][0],
                    "redirect_uri": redirect_uri,
                    "used": False,
                }

        if mode == "deny_authorize":
            callback_params = {
                "error": "access_denied",
                "error_description": "local mock denial",
                "state": state,
            }
        else:
            callback_params = {"code": code, "state": state}

        separator = "&" if urllib.parse.urlsplit(redirect_uri).query else "?"
        location = f"{redirect_uri}{separator}{urllib.parse.urlencode(callback_params)}"
        self.send_response(302)
        self.send_header("location", location)
        self.send_header("content-length", "0")
        self.end_headers()

    def oauth_token(self, body: bytes) -> None:
        params = {
            key: values[0]
            for key, values in urllib.parse.parse_qs(body.decode()).items()
        }

        with STATE.lock:
            STATE.oauth_token_requests += 1
            attempt = STATE.oauth_token_requests
            mode = STATE.oauth_mode

        if mode == "token_temporary_once" and attempt == 1:
            self.send_json(503, {"error": "temporarily_unavailable"})
            return

        if mode == "token_timeout_once" and attempt == 1:
            time.sleep(DELAY_SECONDS)

        if (
            params.get("client_id") != OAUTH_CLIENT_ID
            or params.get("client_secret") != OAUTH_CLIENT_SECRET
            or params.get("grant_type") != "authorization_code"
        ):
            self.send_json(401, {"error": "invalid_client"})
            return

        code_value = params.get("code")
        verifier = params.get("code_verifier", "")
        challenge = base64.urlsafe_b64encode(
            hashlib.sha256(verifier.encode()).digest()
        ).rstrip(b"=").decode()

        with STATE.lock:
            code = STATE.oauth_codes.get(code_value or "")

            if (
                code is None
                or code["used"]
                or code["redirect_uri"] != params.get("redirect_uri")
                or code["challenge"] != challenge
            ):
                valid = False
            else:
                valid = True
                code["used"] = True
                token = f"boundary-access-{STATE.oauth_code_counter}-{attempt}"
                STATE.oauth_tokens.add(token)

        if not valid:
            self.send_json(400, {"error": "invalid_grant"})
            return

        self.send_json(
            200,
            {"access_token": token, "scope": "", "token_type": "bearer"},
        )

    def oauth_user(self) -> None:
        authorization = self.headers.get("authorization", "")
        token = authorization.removeprefix("Bearer ")

        with STATE.lock:
            STATE.oauth_user_requests += 1
            valid = token in STATE.oauth_tokens

        if not valid:
            self.send_json(401, {"error": "unauthorized"})
            return

        self.send_json(
            200,
            {
                "id": 4242,
                "login": "local-neighbor",
                "name": "Local Neighbor",
                "html_url": "https://github.com/local-neighbor",
                "avatar_url": "https://avatars.example.invalid/4242",
            },
        )

    def push(self, body: bytes) -> None:
        try:
            notification = json.loads(body)
        except json.JSONDecodeError:
            self.send_json(400, {"error": "invalid_json"})
            return

        idempotency_key = self.headers.get("idempotency-key", "")
        authorization_ok = (
            self.headers.get("authorization") == f"Bearer {PUSH_BEARER_TOKEN}"
        )

        if not authorization_ok or not idempotency_key:
            self.send_json(401, {"error": "unauthorized"})
            return

        if not all(
            isinstance(notification.get(key), str) and notification[key]
            for key in ("recipient", "title", "body", "idempotency_key")
        ) or notification["idempotency_key"] != idempotency_key:
            self.send_json(400, {"error": "invalid_notification"})
            return

        with STATE.lock:
            STATE.push_attempts += 1
            attempt = STATE.push_attempts
            mode = STATE.push_mode

        if mode == "reject":
            self.send_json(400, {"error": "invalid_recipient"})
            return

        if mode == "temporary_once" and attempt == 1:
            self.send_json(503, {"error": "temporarily_unavailable"})
            return

        with STATE.lock:
            receipt = STATE.push_receipts.get(idempotency_key)
            duplicate = receipt is not None

            if receipt is None:
                receipt = f"push-receipt-{len(STATE.push_receipts) + 1}"
                STATE.push_receipts[idempotency_key] = receipt
                STATE.push_deliveries += 1

        if mode == "timeout_after_accept" and attempt == 1:
            time.sleep(DELAY_SECONDS)

        self.send_json(202, {"id": receipt, "duplicate": duplicate})


class BoundarySMTPServer(socketserver.ThreadingTCPServer):
    allow_reuse_address = True
    daemon_threads = True

    def __init__(self, address: tuple[str, int], mode: str):
        self.mode = mode
        super().__init__(address, BoundarySMTPHandler)


class BoundarySMTPHandler(socketserver.StreamRequestHandler):
    def send_line(self, line: bytes) -> None:
        self.wfile.write(line + b"\r\n")
        self.wfile.flush()

    def handle(self) -> None:
        port = self.server.server_address[1]

        with STATE.lock:
            key = str(port)
            STATE.smtp_connections[key] += 1
            connection_number = STATE.smtp_connections[key]

        if self.server.mode == "timeout":
            time.sleep(DELAY_SECONDS)
            return

        if self.server.mode == "retry" and connection_number == 1:
            self.send_line(b"421 4.3.0 temporary local mock failure")
            return

        self.send_line(b"220 boundary.local ESMTP")
        recipients: list[str] = []

        while True:
            raw = self.rfile.readline(8192)
            if not raw:
                return

            command = raw.decode(errors="replace").strip()
            upper = command.upper()

            if upper.startswith(("EHLO ", "HELO ")):
                self.send_line(b"250-boundary.local")
                self.send_line(b"250 SIZE 1048576")
            elif upper.startswith("MAIL FROM:"):
                self.send_line(b"250 2.1.0 sender ok")
            elif upper.startswith("RCPT TO:"):
                recipient = command.split(":", 1)[1].strip("<>")

                if "reject@" in recipient:
                    with STATE.lock:
                        STATE.smtp_rejections += 1
                    self.send_line(b"550 5.1.1 recipient rejected")
                else:
                    recipients.append(recipient)
                    self.send_line(b"250 2.1.5 recipient ok")
            elif upper == "DATA":
                self.send_line(b"354 end with <CRLF>.<CRLF>")
                chunks: list[bytes] = []

                while True:
                    line = self.rfile.readline(65_536)
                    if not line or line == b".\r\n":
                        break
                    chunks.append(line)

                digest = hashlib.sha256(b"".join(chunks)).hexdigest()

                with STATE.lock:
                    STATE.smtp_messages += 1
                    STATE.smtp_digests.append(digest)

                self.send_line(b"250 2.0.0 queued-boundary")
            elif upper == "RSET":
                recipients = []
                self.send_line(b"250 2.0.0 reset")
            elif upper == "QUIT":
                self.send_line(b"221 2.0.0 bye")
                return
            else:
                self.send_line(b"500 5.5.2 unsupported command")


def serve() -> None:
    http_server = BoundaryHTTPServer(("0.0.0.0", HTTP_PORT), BoundaryHTTPHandler)
    smtp_servers = [
        BoundarySMTPServer(("0.0.0.0", SMTP_PORT), "normal"),
        BoundarySMTPServer(("0.0.0.0", SMTP_RETRY_PORT), "retry"),
        BoundarySMTPServer(("0.0.0.0", SMTP_TIMEOUT_PORT), "timeout"),
    ]
    servers = [http_server, *smtp_servers]
    threads = [
        threading.Thread(target=server.serve_forever, daemon=True) for server in servers
    ]

    for thread in threads:
        thread.start()

    stopped = threading.Event()

    def stop(_signum: int, _frame: Any) -> None:
        stopped.set()

    signal.signal(signal.SIGTERM, stop)
    signal.signal(signal.SIGINT, stop)
    stopped.wait()

    for server in servers:
        server.shutdown()
        server.server_close()


if __name__ == "__main__":
    serve()