#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
collatz_document_base_states.py

If JSON has state info (e.g. per-palier state), document each state.
Outputs MD describing each state.

Usage: python collatz_document_base_states.py --base-projective JSON --output MD
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path
from typing import Any


def load_json(path: str) -> dict[str, Any]:
    """Load JSON file."""
    p = Path(path)
    if not p.exists():
        raise FileNotFoundError(f"File not found: {path}")
    return json.loads(p.read_text(encoding="utf-8"))


def is_state_info(obj: Any) -> bool:
    """Check if object looks like a state (has m, mod, etc.)."""
    if not isinstance(obj, dict):
        return False
    return "m" in obj or "mod" in obj


def document_states(data: dict[str, Any]) -> list[str]:
    """Extract and document each state from JSON."""
    lines: list[str] = ["# Base states documentation", ""]

    for key, val in sorted(data.items()):
        if not is_state_info(val):
            continue
        val = val if isinstance(val, dict) else {}
        lines.append(f"## State {key}")
        lines.append("")
        for k, v in sorted(val.items()):
            if isinstance(v, list) and len(v) > 10:
                lines.append(f"- **{k}**: {len(v)} items")
            else:
                lines.append(f"- **{k}**: {v}")
        lines.append("")

    if len(lines) <= 2:
        lines.append("No state information found in JSON.")

    return lines


def run(base_path: str, output_path: str) -> None:
    """Run and write output."""
    try:
        data = load_json(base_path)
        lines = document_states(data)
    except FileNotFoundError as e:
        lines = ["# Base states documentation", "", str(e)]
    except json.JSONDecodeError as e:
        lines = ["# Base states documentation", "", f"Invalid JSON: {e}"]

    Path(output_path).write_text("\n".join(lines) + "\n", encoding="utf-8")


def main() -> None:
    parser = argparse.ArgumentParser(description="Document base states from JSON")
    parser.add_argument("--base-projective", required=True, help="Path to base-projective JSON")
    parser.add_argument("--output", required=True, help="Path to output MD file")
    args = parser.parse_args()
    run(args.base_projective, args.output)


if __name__ == "__main__":
    main()