#!/usr/bin/env python3
"""Incrementally add or remove nodes and edges in the context graph.

Usage:
    python3 .claude/skills/flydocs-workflow/scripts/graph_update.py \\
        add-node <ID> --type <TYPE> [--label STR] [--path STR]

    python3 .claude/skills/flydocs-workflow/scripts/graph_update.py \\
        remove-node <ID>

    python3 .claude/skills/flydocs-workflow/scripts/graph_update.py \\
        add-edge <FROM> <TO> <REL> [--weight N] [--manual]

    python3 .claude/skills/flydocs-workflow/scripts/graph_update.py \\
        remove-edge <FROM> <TO> <REL>
"""

import argparse
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent))
from graph_utils import (
    find_project_root,
    load_graph,
    save_graph,
    output_json,
    fail,
    VALID_NODE_TYPES,
    VALID_REL_TYPES,
)


def cmd_add_node(args, graph):
    """Add a node to the graph."""
    node_id = args.id
    if args.type not in VALID_NODE_TYPES:
        fail(f"Invalid node type: {args.type}. Valid: {', '.join(sorted(VALID_NODE_TYPES))}")

    node = {"type": args.type}
    if args.label:
        node["label"] = args.label
    else:
        # Default label from ID
        node["label"] = node_id.split(":", 1)[-1] if ":" in node_id else node_id
    if args.path:
        node["path"] = args.path
    if args.status:
        node["status"] = args.status
    if args.date:
        node["date"] = args.date

    node["manual"] = True

    graph["nodes"][node_id] = node
    return {"success": True, "action": "add-node", "node": node_id}


def cmd_remove_node(args, graph):
    """Remove a node and its connected edges from the graph."""
    node_id = args.id
    if node_id not in graph["nodes"]:
        fail(f"Node not found: {node_id}")

    del graph["nodes"][node_id]

    # Remove connected edges
    before = len(graph["edges"])
    graph["edges"] = [
        e for e in graph["edges"]
        if e["from"] != node_id and e["to"] != node_id
    ]
    removed_edges = before - len(graph["edges"])

    return {
        "success": True,
        "action": "remove-node",
        "removed": node_id,
        "removedEdges": removed_edges,
    }


def cmd_add_edge(args, graph):
    """Add an edge to the graph."""
    if args.rel not in VALID_REL_TYPES:
        fail(f"Invalid relationship: {args.rel}. Valid: {', '.join(sorted(VALID_REL_TYPES))}")

    # Validate nodes exist
    if args.source not in graph["nodes"]:
        fail(f"Source node not found: {args.source}")
    if args.target not in graph["nodes"]:
        fail(f"Target node not found: {args.target}")

    # Check for duplicate
    for edge in graph["edges"]:
        if edge["from"] == args.source and edge["to"] == args.target and edge["rel"] == args.rel:
            fail(f"Edge already exists: {args.source} --{args.rel}--> {args.target}")

    edge = {
        "from": args.source,
        "to": args.target,
        "rel": args.rel,
        "weight": args.weight,
    }
    if args.manual:
        edge["manual"] = True

    graph["edges"].append(edge)

    return {
        "success": True,
        "action": "add-edge",
        "edge": f"{args.source} --{args.rel}--> {args.target}",
    }


def cmd_remove_edge(args, graph):
    """Remove an edge from the graph."""
    before = len(graph["edges"])
    graph["edges"] = [
        e for e in graph["edges"]
        if not (e["from"] == args.source and e["to"] == args.target and e["rel"] == args.rel)
    ]

    if len(graph["edges"]) == before:
        fail(f"Edge not found: {args.source} --{args.rel}--> {args.target}")

    return {
        "success": True,
        "action": "remove-edge",
        "removed": f"{args.source} --{args.rel}--> {args.target}",
    }


def main():
    parser = argparse.ArgumentParser(description="Update the context graph")
    parser.add_argument(
        "--root", type=str, default=None,
        help="Project root (default: auto-detect)"
    )
    subparsers = parser.add_subparsers(dest="command", required=True)

    # add-node
    p_add_node = subparsers.add_parser("add-node", help="Add a node")
    p_add_node.add_argument("id", help="Node ID (e.g., module:auth)")
    p_add_node.add_argument("--type", required=True, help="Node type")
    p_add_node.add_argument("--label", help="Human-readable label")
    p_add_node.add_argument("--path", help="File path relative to project root")
    p_add_node.add_argument("--status", help="Node status")
    p_add_node.add_argument("--date", help="Date for temporal nodes (ISO format)")

    # remove-node
    p_rm_node = subparsers.add_parser("remove-node", help="Remove a node")
    p_rm_node.add_argument("id", help="Node ID to remove")

    # add-edge
    p_add_edge = subparsers.add_parser("add-edge", help="Add an edge")
    p_add_edge.add_argument("source", help="Source node ID")
    p_add_edge.add_argument("target", help="Target node ID")
    p_add_edge.add_argument("rel", help="Relationship type (e.g., EXTENDS)")
    p_add_edge.add_argument("--weight", type=float, default=0.8, help="Edge weight 0.0-1.0")
    p_add_edge.add_argument("--manual", action="store_true", help="Mark as manually added")

    # remove-edge
    p_rm_edge = subparsers.add_parser("remove-edge", help="Remove an edge")
    p_rm_edge.add_argument("source", help="Source node ID")
    p_rm_edge.add_argument("target", help="Target node ID")
    p_rm_edge.add_argument("rel", help="Relationship type")

    args = parser.parse_args()

    root = Path(args.root) if args.root else find_project_root()
    if not root:
        fail("Could not find project root (no .flydocs/ directory found)")

    graph = load_graph(root)

    if args.command == "add-node":
        result = cmd_add_node(args, graph)
    elif args.command == "remove-node":
        result = cmd_remove_node(args, graph)
    elif args.command == "add-edge":
        result = cmd_add_edge(args, graph)
    elif args.command == "remove-edge":
        result = cmd_remove_edge(args, graph)
    else:
        fail(f"Unknown command: {args.command}")

    save_graph(root, graph)
    output_json(result)


if __name__ == "__main__":
    main()