#!/usr/bin/env python3
from __future__ import annotations

import argparse
from collections import Counter, defaultdict

from vault_common import all_records, json_print, known_note_names, source_values, wikilinks


def broken_links(records: list[dict[str, object]]) -> list[dict[str, object]]:
    known = known_note_names(records)
    result = []
    for record in records:
        links = wikilinks(record["body"])
        missing = sorted({link for link in links if link not in known})
        if missing:
            result.append({"path": record["path"], "missing": missing})
    return result


def orphans(records: list[dict[str, object]]) -> list[dict[str, object]]:
    incoming: Counter[str] = Counter()
    aliases: dict[str, str] = {}
    path_by_name: dict[str, str] = {}

    for record in records:
        path_by_name[record["name"]] = record["path"]
        aliases[record["name"]] = record["name"]
        aliases[record["path"]] = record["name"]
        for alias in record["aliases"]:
            aliases[alias] = record["name"]

    for record in records:
        for link in wikilinks(record["body"]):
            target = aliases.get(link)
            if target and target != record["name"]:
                incoming[target] += 1

    result = []
    for name, path in sorted(path_by_name.items(), key=lambda item: item[1]):
        if incoming[name] == 0 and not path.startswith("99 System/Archive/"):
            result.append({"path": path, "name": name, "incoming": 0})
    return result


def source_links(records: list[dict[str, object]]) -> dict[str, object]:
    known = known_note_names(records)
    linked = []
    missing = []
    plain = []

    for record in records:
        for source in source_values(record["frontmatter"].get("source")):
            links = wikilinks(source)
            if links:
                for link in links:
                    item = {"path": record["path"], "source": source, "target": link}
                    if link in known:
                        linked.append(item)
                    else:
                        missing.append(item)
            else:
                plain.append({"path": record["path"], "source": source})

    return {"linked": linked, "missing": missing, "plain": plain}


def graph_components(records: list[dict[str, object]]) -> list[dict[str, object]]:
    known = known_note_names(records)
    name_by_link = {}
    for record in records:
        name_by_link[record["name"]] = record["name"]
        name_by_link[record["path"]] = record["name"]
        for alias in record["aliases"]:
            name_by_link[alias] = record["name"]

    graph: dict[str, set[str]] = defaultdict(set)
    path_by_name = {record["name"]: record["path"] for record in records}
    for record in records:
        name = record["name"]
        graph.setdefault(name, set())
        for link in wikilinks(record["body"]):
            if link not in known:
                continue
            target = name_by_link.get(link)
            if target and target != name:
                graph[name].add(target)
                graph[target].add(name)

    seen = set()
    components = []
    for name in sorted(graph):
        if name in seen:
            continue
        stack = [name]
        component = []
        seen.add(name)
        while stack:
            current = stack.pop()
            component.append(current)
            for next_name in graph[current]:
                if next_name not in seen:
                    seen.add(next_name)
                    stack.append(next_name)
        components.append(
            {
                "size": len(component),
                "notes": [path_by_name[item] for item in sorted(component)[:20]],
                "truncated": len(component) > 20,
            }
        )

    return sorted(components, key=lambda item: item["size"])


def main() -> int:
    parser = argparse.ArgumentParser(description="Agent-facing vault audit helper")
    sub = parser.add_subparsers(dest="command", required=True)
    sub.add_parser("broken-links")
    sub.add_parser("orphans")
    sub.add_parser("source-links")
    sub.add_parser("components")
    sub.add_parser("all")
    args = parser.parse_args()

    records = all_records()
    if args.command == "broken-links":
        json_print({"ok": True, "broken_links": broken_links(records)})
    elif args.command == "orphans":
        json_print({"ok": True, "orphans": orphans(records)})
    elif args.command == "source-links":
        json_print({"ok": True, "source_links": source_links(records)})
    elif args.command == "components":
        json_print({"ok": True, "components": graph_components(records)})
    elif args.command == "all":
        json_print(
            {
                "ok": True,
                "broken_links": broken_links(records),
                "orphans": orphans(records),
                "source_links": source_links(records),
                "components": graph_components(records),
            }
        )

    return 0


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