import yaml
import argparse
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--dag", default="spec/41_release_dag.yaml")
    parser.add_argument("--strict", action="store_true")
    args = parser.parse_args()

    dag_path = ROOT / args.dag
    if not dag_path.exists():
        print(f"DAG file not found: {dag_path}")
        return 1

    try:
        data = yaml.safe_load(dag_path.read_text(encoding="utf-8"))
    except Exception as e:
        print(f"Error parsing YAML: {e}")
        return 1

    nodes = data.get("dag", {}).get("nodes", {})
    all_outputs = set()
    for node in nodes.values():
        all_outputs.update(node.get("outputs") or [])

    missing_inputs = []
    for node_id, node in nodes.items():
        for inp in (node.get("inputs") or []):
            if inp in all_outputs:
                continue
            if not (ROOT / inp).exists():
                missing_inputs.append((node_id, inp))

    if missing_inputs:
        print(f"Found {len(missing_inputs)} missing inputs:")
        for nid, inp in missing_inputs:
            print(f"  Node [{nid}] requires [{inp}]")
        if args.strict:
            return 1
    else:
        print("All DAG inputs are resolved (exist or produced by upstream nodes).")
    
    return 0

if __name__ == "__main__":
    sys.exit(main())