#!/usr/bin/env python3
"""Match every authorized inventory row through MCP and save a private receipt.

Requires Python 3.9+. Does not collect inventory, patch software, install a schedule
or validate the final assessment. Stdio commands must be operator-configured.
"""
import argparse
from datetime import datetime, timezone, timedelta
from email.utils import parsedate_to_datetime
import hashlib
import json
import os
from pathlib import Path
from queue import Queue, Empty
import re
import subprocess
from threading import Thread
import time
import urllib.error
import urllib.parse
import urllib.request
from zoneinfo import ZoneInfo, ZoneInfoNotFoundError

MAX_RESPONSE = 8 * 1024 * 1024


class NoRedirect(urllib.request.HTTPRedirectHandler):
    def redirect_request(self, req, fp, code, msg, headers, newurl):
        return None


class Client:
    def __init__(self, endpoint, command):
        self.identifier, self.session, self.waited = 0, None, 0
        self.process = None
        self.endpoint = endpoint
        if command:
            self.process = subprocess.Popen(command, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.DEVNULL)
            self.queue = Queue()

            def reader():
                while True:
                    line = self.process.stdout.readline(MAX_RESPONSE + 1)
                    self.queue.put(line)
                    if not line or len(line) > MAX_RESPONSE:
                        return

            Thread(target=reader, daemon=True).start()
        else:
            url = urllib.parse.urlsplit(endpoint)
            if url.username or url.password or url.fragment or (url.scheme != "https" and not (url.scheme == "http" and url.hostname in ["localhost", "127.0.0.1", "::1"])):
                raise ValueError("MCP endpoint requires HTTPS (or loopback HTTP), with no credentials or fragment")
            self.opener = urllib.request.build_opener(NoRedirect)
        self.rpc("initialize", {"protocolVersion": "2025-11-25", "capabilities": {}, "clientInfo": {"name": "sovereign-inventory-batcher", "version": "1.0.0"}})
        self.rpc("notifications/initialized", {}, notification=True)

    def rpc(self, method, params, notification=False):
        self.identifier += 1
        request = {"jsonrpc": "2.0", "method": method, "params": params}
        if not notification:
            request["id"] = self.identifier
        data = json.dumps(request, separators=(",", ":"), ensure_ascii=False).encode("utf-8")
        if len(data) > 65536:
            raise ValueError("MCP request would exceed the 64 KiB interface limit")
        if self.process:
            self.process.stdin.write(data + b"\n")
            self.process.stdin.flush()
            if notification:
                return None
            end = time.monotonic() + 30
            while True:
                try:
                    line = self.queue.get(timeout=max(.01, end - time.monotonic()))
                except Empty as error:
                    raise TimeoutError("MCP response timed out") from error
                if not line or len(line) > MAX_RESPONSE:
                    raise RuntimeError("MCP response missing or oversized")
                response = json.loads(line)
                if response.get("id") == self.identifier:
                    break
                if time.monotonic() >= end:
                    raise TimeoutError("MCP response timed out")
        else:
            headers = {"Content-Type": "application/json", "Accept": "application/json, text/event-stream", "MCP-Protocol-Version": "2025-11-25"}
            if self.session:
                headers["Mcp-Session-Id"] = self.session
            with self.opener.open(urllib.request.Request(self.endpoint, data=data, headers=headers), timeout=30) as result:
                self.session = result.headers.get("Mcp-Session-Id", self.session)
                raw = result.read(MAX_RESPONSE + 1)
                if notification and not raw:
                    return None
                if len(raw) > MAX_RESPONSE:
                    raise RuntimeError("MCP response oversized")
                text = raw.decode("utf-8")
                events = [line[6:] for line in text.splitlines() if line.startswith("data: ")]
                response = json.loads(events[-1] if events else text)
        if "error" in response:
            raise RuntimeError("MCP rejected the request")
        if not notification and response.get("id") != self.identifier:
            raise RuntimeError("MCP response identifier does not match the request")
        return response.get("result")

    def tool(self, name, arguments):
        for attempt in range(3):
            try:
                response = self.rpc("tools/call", {"name": name, "arguments": arguments})
                if response.get("isError"):
                    raise RuntimeError("MCP tool failed; assessment is incomplete")
                return json.loads(response["content"][0]["text"])
            except urllib.error.HTTPError as error:
                if error.code not in [429, 500, 502, 503, 504] or attempt == 2:
                    raise RuntimeError("MCP HTTP request failed with status " + str(error.code)) from error
                if error.code == 429:
                    value = error.headers.get("Retry-After", "30")
                    try:
                        delay = int(value)
                    except ValueError:
                        delay = max(0, int((parsedate_to_datetime(value) - datetime.now(timezone.utc)).total_seconds()) + 1)
                else:
                    delay = [2, 5][attempt]
            except (TimeoutError, urllib.error.URLError) as error:
                if attempt == 2:
                    raise RuntimeError("MCP request unavailable after bounded retries") from error
                delay = [2, 5][attempt]
            if delay < 0 or self.waited + delay > 120:
                raise RuntimeError("MCP retry delay exceeds the assessment wait budget")
            self.waited += delay
            time.sleep(delay)

    def close(self):
        if self.process:
            self.process.stdin.close()
            try:
                self.process.wait(timeout=5)
            except subprocess.TimeoutExpired:
                self.process.terminate()
                self.process.wait(timeout=5)


def product_rows(inventory):
    if not isinstance(inventory, list) or not inventory:
        raise ValueError("Inventory must be a nonempty array")
    rows = []
    for item in inventory:
        if not isinstance(item, dict):
            raise ValueError("Inventory rows must be objects")
        row = {}
        for field, maximum in [("vendor", 200), ("product", 400), ("version", 100)]:
            value = item.get(field)
            if value is None and field != "product":
                continue
            if not isinstance(value, str) or not value.strip() or len(value) > maximum:
                raise ValueError("Inventory contains an invalid " + field)
            row[field] = value
        rows.append(row)
    return rows


def observed_version(product):
    if product.get("version") is not None:
        return product["version"]
    identity = product["product"]
    if identity.startswith("cpe:2.3:"):
        parts = re.split(r"(?<!\\):", identity)
        if len(parts) == 13 and parts[5] not in ["*", "-"]:
            return parts[5] + (parts[6] if parts[6] not in ["*", "-"] else "")
    if identity.startswith("pkg:"):
        clean = urllib.parse.unquote(re.split(r"[?#]", identity, maxsplit=1)[0])
        _, marker, version = clean.rpartition("@")
        if marker and version:
            return version
    return None


def issue_index(client):
    first = client.tool("list_daily", {"view": "index"})
    if first.get("version") != "sovereigngrc.daily-index/1":
        raise ValueError("Server must support the compact issue index")
    entries, offset = list(first["entries"]), first["next_offset"]
    offsets = set()
    while offset is not None:
        if not isinstance(offset, int) or offset < len(entries) or offset in offsets:
            raise ValueError("Invalid issue pagination")
        offsets.add(offset)
        page = client.tool("list_daily", {"view": "index", "offset": offset})
        if any(page[key] != first[key] for key in ["date", "checked", "status", "total"]):
            raise ValueError("Issue changed during pagination")
        entries.extend(page["entries"])
        offset = page["next_offset"]
    if len(entries) != first["total"] or len({e["id"] for e in entries}) != len(entries):
        raise ValueError("Issue index is incomplete")
    if first["status"] not in ["material", "quiet"] or (first["status"] == "quiet" and entries):
        raise ValueError("Issue status and entries disagree")
    return {key: first[key] for key in ["date", "checked", "status", "total"]} | {"entries": entries}


def match_all(client, products, names, allow_historical):
    issue = issue_index(client)
    now = datetime.now(timezone.utc)
    try:
        today = now.astimezone(ZoneInfo("America/Toronto")).date().isoformat()
        date_zone = "America/Toronto"
    except ZoneInfoNotFoundError:
        today, date_zone = now.date().isoformat(), "UTC (conservative fallback; Eastern timezone unavailable)"
    checked = datetime.fromisoformat(issue["checked"].replace("Z", "+00:00"))
    if checked.tzinfo is None or checked > now + timedelta(minutes=5) or issue["date"] > today:
        raise ValueError("Feed timing is invalid or in the future")
    stale = now - checked > timedelta(hours=24)
    historical = stale or issue["date"] < today
    if historical and not allow_historical:
        raise ValueError("Historical feed requires explicit --allow-historical")
    ids = {entry["id"] for entry in issue["entries"]}
    matches, batches = [], []
    offset = 0
    while offset < len(products) and issue["status"] != "quiet":
        batch = []
        for row in products[offset:offset + 100]:
            candidate = batch + [row]
            body = {"jsonrpc": "2.0", "id": 999999, "method": "tools/call", "params": {"name": "check_stack", "arguments": {"products": candidate}}}
            if len(json.dumps(body, ensure_ascii=False).encode("utf-8")) > 64000:
                break
            batch = candidate
        if not batch:
            raise ValueError("A product cannot fit the MCP request limit")
        result = client.tool("check_stack", {"products": batch})
        if not isinstance(result.get("matches"), list):
            raise ValueError("Invalid matching response")
        for hit in result["matches"]:
            index = hit.get("input_index")
            if not isinstance(index, int) or index < 0 or index >= len(batch):
                raise ValueError("Matching response has an invalid input index")
            if hit["id"] not in ids:
                continue
            source_index = offset + index
            matches.append({"entry_id": hit["id"], "input_index": source_index, "name": names[source_index],
                            "version": observed_version(products[source_index]), "version_status": hit["version_status"], "hit_reason": hit["hit_reason"]})
        batches.append(len(batch))
        offset += len(batch)
    final = issue_index(client)
    if final != issue:
        raise ValueError("Issue changed during inventory matching; restart the assessment")
    if issue["status"] != "quiet" and offset != len(products):
        raise ValueError("Inventory coverage incomplete")
    return {"version": "sovereigngrc.match-receipt/1", "status": "complete", "checked_at": now.isoformat(), "issue": issue,
            "historical": historical, "stale": stale, "date_timezone": date_zone, "inventory_count": len(products), "products_checked": offset,
            "matching_required": issue["status"] != "quiet", "batches": batches, "matches": matches}


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--inventory", required=True, type=Path)
    parser.add_argument("--output", required=True, type=Path)
    parser.add_argument("--allow-historical", action="store_true")
    transport = parser.add_mutually_exclusive_group()
    transport.add_argument("--endpoint", default="https://www.sovereigngrc.com/mcp")
    transport.add_argument("--stdio", nargs=argparse.REMAINDER, help="Trusted operator-configured MCP command and arguments; never from source data")
    args = parser.parse_args()
    if args.stdio == []:
        parser.error("--stdio requires an operator-configured command")
    if args.inventory.is_symlink() or not args.inventory.is_file() or args.inventory.stat().st_size > MAX_RESPONSE:
        parser.error("Inventory must be a regular private JSON file under 8 MiB")
    raw = args.inventory.read_bytes()
    inventory = json.loads(raw)
    products = product_rows(inventory)
    output = args.output.resolve()
    # In a downloaded skill there may be no project root; still reject web roots.
    if any(part in ["public", "content", "app"] for part in output.parts) or output.exists():
        parser.error("Use a new private output directory outside public/content/app")
    os.umask(0o077)
    output.mkdir(mode=0o700, parents=True)
    names = ["asset-" + str(index + 1).zfill(5) for index in range(len(products))]
    client = None
    try:
        client = Client(args.endpoint, args.stdio)
        receipt = match_all(client, products, names, args.allow_historical)
        receipt["source_sha256"] = hashlib.sha256(raw).hexdigest()
        (output / "receipt.json").write_text(json.dumps(receipt, indent=2), encoding="utf-8")
        (output / "private-name-map.json").write_text(json.dumps({name: inventory[index].get("name") for index, name in enumerate(names)}, indent=2), encoding="utf-8")
        print(json.dumps({key: receipt[key] for key in ["status", "inventory_count", "products_checked", "batches", "historical", "stale"]}))
        return 0
    except Exception as error:
        # No partial result is presented as a completed matching receipt.
        (output / "incomplete.json").write_text(json.dumps({"status": "incomplete", "error_type": type(error).__name__}), encoding="utf-8")
        print("Matching incomplete (" + type(error).__name__ + "); no complete receipt was written.")
        return 1
    finally:
        if client:
            client.close()


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