from __future__ import annotations

import argparse
import hashlib
import json
import os
import re
import secrets
import struct
import subprocess
import time
from datetime import UTC, datetime
from pathlib import Path
from typing import Any

from pcapng import FileScanner
from pcapng.blocks import EnhancedPacket
from preflight import (
    pymobiledevice3_binary,
    pymobiledevice3_version,
    usb_device_count,
)

REPO_ROOT = Path(__file__).resolve().parents[4]
SESSION_ROOT = REPO_ROOT / "sessions" / "iphone-k1-observation"
SAFE_LABEL = re.compile(r"^[a-z0-9][a-z0-9-]{0,31}$")


def utc_now() -> str:
    return datetime.now(UTC).isoformat()


def sha256_file(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as stream:
        for chunk in iter(lambda: stream.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def git_output(*args: str) -> bytes:
    completed = subprocess.run(
        ["git", *args],
        cwd=REPO_ROOT,
        check=True,
        capture_output=True,
    )
    return completed.stdout


def source_snapshot() -> dict[str, Any]:
    revision = git_output("rev-parse", "HEAD").decode().strip()
    status = git_output("status", "--porcelain=v1", "--untracked-files=all")
    dirty = bool(status.strip())
    digest = hashlib.sha256()
    digest.update(revision.encode())
    digest.update(b"\0")
    digest.update(status)
    if dirty:
        digest.update(git_output("diff", "--binary", "HEAD"))
        untracked = git_output("ls-files", "--others", "--exclude-standard", "-z")
        for raw_name in filter(None, untracked.split(b"\0")):
            relative = raw_name.decode("utf-8", errors="surrogateescape")
            path = REPO_ROOT / relative
            digest.update(raw_name)
            digest.update(b"\0")
            if path.is_file():
                digest.update(sha256_file(path).encode())
    return {
        "git_revision": revision,
        "source_tree_dirty": dirty,
        "source_tree_fingerprint_sha256": digest.hexdigest(),
    }


def terminate(process: subprocess.Popen[bytes]) -> int:
    if process.poll() is None:
        process.terminate()
        try:
            process.wait(timeout=10)
        except subprocess.TimeoutExpired:
            process.kill()
            process.wait(timeout=5)
    return int(process.returncode or 0)


def artifact_record(path: Path, kind: str, session_dir: Path) -> dict[str, Any] | None:
    if not path.is_file() or path.stat().st_size == 0:
        return None
    path.chmod(0o600)
    return {
        "kind": kind,
        "relative_path": str(path.relative_to(session_dir)),
        "size_bytes": path.stat().st_size,
        "sha256": sha256_file(path),
    }


def convert_pcapng_to_pcap(source: Path, destination: Path) -> int:
    """Create a classic Ethernet PCAP that macOS's bundled tcpdump can read."""
    packet_count = 0
    with source.open("rb") as source_stream, destination.open("xb") as destination_stream:
        destination_stream.write(
            struct.pack(
                "<IHHIIII",
                0xA1B2C3D4,
                2,
                4,
                0,
                0,
                262_144,
                1,
            )
        )
        for block in FileScanner(source_stream):
            if not isinstance(block, EnhancedPacket):
                continue
            packet_data = bytes(block.packet_data)
            captured_length = int(block.captured_len)
            packet_length = int(block.packet_len)
            if captured_length != len(packet_data) or packet_length < captured_length:
                raise ValueError("invalid enhanced-packet lengths")

            timestamp = float(block.timestamp)
            seconds = int(timestamp)
            microseconds = int(round((timestamp - seconds) * 1_000_000))
            if microseconds >= 1_000_000:
                seconds += 1
                microseconds -= 1_000_000
            destination_stream.write(
                struct.pack(
                    "<IIII",
                    seconds,
                    microseconds,
                    captured_length,
                    packet_length,
                )
            )
            destination_stream.write(packet_data)
            packet_count += 1
    destination.chmod(0o600)
    return packet_count


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Bounded, redacted USB iPhone capture for the Mission Core K1 lab"
    )
    parser.add_argument("--duration", type=int, default=120, help="capture seconds (5..600)")
    parser.add_argument("--label", default="lixelgo-observe", help="safe session label")
    parser.add_argument("--process", help="optional iOS process filter for network capture")
    parser.add_argument("--interface", help="optional iOS interface filter for network capture")
    parser.add_argument(
        "--bluetooth",
        action="store_true",
        help="also capture iPhone Bluetooth HCI; requires Apple's logging profile",
    )
    parser.add_argument(
        "--acknowledge-ios-bluetooth-profile",
        action="store_true",
        help="confirm the operator already installed Apple's Bluetooth logging profile",
    )
    parser.add_argument(
        "--allow-dirty",
        action="store_true",
        help="record an explicit dirty-tree fingerprint instead of refusing the capture",
    )
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    if not 5 <= args.duration <= 600:
        raise SystemExit("--duration must be between 5 and 600 seconds")
    if not SAFE_LABEL.fullmatch(args.label):
        raise SystemExit("--label must use lowercase ASCII letters, digits, and hyphens")
    if args.bluetooth and not args.acknowledge_ios_bluetooth_profile:
        raise SystemExit(
            "Bluetooth capture requires --acknowledge-ios-bluetooth-profile; "
            "the script never installs the profile"
        )

    os.umask(0o077)
    binary = pymobiledevice3_binary()
    if usb_device_count(binary) != 1:
        raise SystemExit("exactly one trusted USB iPhone must be connected")
    snapshot = source_snapshot()
    if snapshot["source_tree_dirty"] and not args.allow_dirty:
        raise SystemExit(
            "repository is dirty; commit the milestone or pass --allow-dirty "
            "to record its fingerprint"
        )

    session_id = (
        datetime.now(UTC).strftime("%Y%m%dT%H%M%SZ") + f"_{args.label}_{secrets.token_hex(2)}"
    )
    session_dir = SESSION_ROOT / session_id
    capture_dir = session_dir / "captures"
    capture_dir.mkdir(parents=True, mode=0o700)
    network_pcapng_path = capture_dir / "iphone-network.pcapng"
    network_pcap_path = capture_dir / "iphone-network.pcap"
    bluetooth_path = capture_dir / "iphone-bluetooth.pcapng"
    stderr_path = session_dir / "capture.stderr.log"
    manifest_path = session_dir / "manifest.redacted.json"

    manifest: dict[str, Any] = {
        "schema": "missioncore.xgrids-k1/iphone-observation-manifest/v1alpha1",
        "session_id": session_id,
        "status": "starting",
        "started_at_utc": utc_now(),
        "started_monotonic_ns": time.monotonic_ns(),
        "requested_duration_seconds": args.duration,
        "capture_backend": "pymobiledevice3-9.36.0/com.apple.pcapd",
        "pymobiledevice3_version": pymobiledevice3_version(binary),
        "network_capture": True,
        "bluetooth_capture": args.bluetooth,
        "network_process_filter": args.process,
        "network_interface_filter": args.interface,
        "device_identity_recorded": False,
        "credentials_recorded_in_manifest": False,
        "source": snapshot,
        "artifacts": [],
    }
    manifest_path.write_text(json.dumps(manifest, indent=2) + "\n", encoding="utf-8")
    manifest_path.chmod(0o600)

    network_command = [binary, "--no-color", "pcap", "--out", str(network_pcapng_path)]
    if args.process:
        network_command.extend(["--process", args.process])
    if args.interface:
        network_command.extend(["--interface", args.interface])
    commands = [("network", network_command)]
    if args.bluetooth:
        commands.append(
            (
                "bluetooth",
                [binary, "--no-color", "btlogger", str(bluetooth_path), "--format", "pcapng"],
            )
        )

    processes: list[tuple[str, subprocess.Popen[bytes]]] = []
    exit_codes: dict[str, int] = {}
    interrupted = False
    failure: str | None = None
    with stderr_path.open("wb") as error_log:
        stderr_path.chmod(0o600)
        try:
            for name, command in commands:
                process = subprocess.Popen(
                    command,
                    cwd=REPO_ROOT,
                    stdout=subprocess.DEVNULL,
                    stderr=error_log,
                )
                processes.append((name, process))
            time.sleep(1)
            early = [name for name, process in processes if process.poll() is not None]
            if early:
                raise RuntimeError(f"capture service exited early: {', '.join(early)}")

            deadline = time.monotonic() + args.duration
            while time.monotonic() < deadline:
                ended = [name for name, process in processes if process.poll() is not None]
                if ended:
                    raise RuntimeError(f"capture service stopped unexpectedly: {', '.join(ended)}")
                time.sleep(min(0.25, deadline - time.monotonic()))
        except KeyboardInterrupt:
            interrupted = True
        except RuntimeError as exc:
            failure = str(exc)
        finally:
            for name, process in processes:
                exit_codes[name] = terminate(process)

    network_packet_count: int | None = None
    if failure is None:
        try:
            network_packet_count = convert_pcapng_to_pcap(
                network_pcapng_path,
                network_pcap_path,
            )
            if network_packet_count == 0:
                failure = "network capture validation found no packets"
        except Exception as exc:  # noqa: BLE001 - third-party parser errors vary by version
            failure = f"network capture validation failed: {type(exc).__name__}"

    artifacts = [
        artifact_record(network_pcapng_path, "iphone-network-pcapng-raw", session_dir),
        artifact_record(network_pcap_path, "iphone-network-pcap-compatible", session_dir),
        artifact_record(bluetooth_path, "iphone-bluetooth-hci-pcapng", session_dir),
        artifact_record(stderr_path, "capture-stderr", session_dir),
    ]
    manifest.update(
        {
            "status": "failed" if failure else "completed",
            "interrupted_by_operator": interrupted,
            "failure": failure,
            "ended_at_utc": utc_now(),
            "ended_monotonic_ns": time.monotonic_ns(),
            "process_exit_codes": exit_codes,
            "network_packet_count": network_packet_count,
            "artifacts": [item for item in artifacts if item is not None],
        }
    )
    manifest_path.write_text(json.dumps(manifest, indent=2) + "\n", encoding="utf-8")
    manifest_path.chmod(0o600)
    print(f"session: {session_dir}")
    print(f"status: {manifest['status']}")
    return 1 if failure else 0


if __name__ == "__main__":
    raise SystemExit(main())