#!/usr/bin/env python3
"""Small e-Gov API v2 evidence client. Python 3.10+, standard library, MIT.

Search, save a law, compare two explicit dates, or extract one article offline.
https://codeagent.jp/posts/egov-law-api-v2-python/
"""
import argparse
import base64
import binascii
from datetime import date, datetime, timezone
import difflib
import hashlib
import json
from pathlib import Path
import re
import sys
from urllib.error import HTTPError, URLError
from urllib.parse import quote, urlencode
from urllib.request import Request, urlopen
import xml.etree.ElementTree as ET

BASE = "https://laws.e-gov.go.jp/api/2"
MAX_BYTES = 12 * 1024 * 1024  # Small-example limit, not an API limit.
CONTAINERS = {"Part", "Chapter", "Section", "Subsection", "Division"}


def fetch_json(path, params):
    url = f"{BASE}/{path}?{urlencode(params)}"
    request = Request(url, headers={"Accept": "application/json"})
    with urlopen(request, timeout=30) as response:
        if response.headers.get_content_type() != "application/json":
            raise ValueError("Expected application/json; check the service response")
        raw = response.read(MAX_BYTES + 1)
    if len(raw) > MAX_BYTES:
        raise ValueError("Response exceeds this example's 12 MiB limit")
    data = json.loads(raw)
    if not isinstance(data, dict):
        raise ValueError("Expected a JSON object; check the current API schema")
    return url, data


def local_name(node):
    return node.tag.rsplit("}", 1)[-1]


def validated_xml(data):
    if not isinstance(data, dict):
        raise ValueError("Expected response to be a JSON object")
    for key in ("law_info", "revision_info"):
        if not isinstance(data.get(key), dict):
            raise ValueError(f"Expected {key} to be a JSON object")
    if not isinstance(data["law_info"].get("law_id"), str) or not data["law_info"]["law_id"]:
        raise ValueError("Missing law ID; response was not saved as a success")
    if not isinstance(data["revision_info"].get("law_revision_id"), str) or not data["revision_info"]["law_revision_id"]:
        raise ValueError("Missing revision ID; check the API response")
    if not data["revision_info"]["law_revision_id"].startswith(data["law_info"]["law_id"] + "_"):
        raise ValueError("Revision ID does not belong to the response law ID")
    xml_bytes = base64.b64decode(data["law_full_text"], validate=True)
    # Only the documented UTF-8 law XML is accepted by this small example.
    xml_text = xml_bytes.decode("utf-8-sig")
    if "<!DOCTYPE" in xml_text or "<!ENTITY" in xml_text:
        raise ValueError("DTD/entity declarations are outside this example's scope")
    root = ET.fromstring(xml_bytes)
    if local_name(root) != "Law":
        raise ValueError("Expected a full Law XML document")
    return xml_bytes, root


def now():
    return datetime.now(timezone.utc)


def envelope(url, data, asof):
    return {"retrieved_at": now().isoformat(), "request_url": url,
            "requested_asof": asof, "response": data}


def write_json(path, data):
    with path.open("x", encoding="utf-8") as file:
        json.dump(data, file, ensure_ascii=False, indent=2)
        file.write("\n")


def write_snapshot(stem, snapshot, xml_bytes):
    write_json(stem.with_suffix(".json"), snapshot)
    with stem.with_suffix(".xml").open("xb") as file:
        file.write(xml_bytes)


def get_snapshot(law_id, asof):
    params = {"response_format": "json", "law_full_text_format": "xml"}
    if asof:
        params["asof"] = asof
    url, data = fetch_json("law_data/" + quote(law_id, safe=""), params)
    xml_bytes, _ = validated_xml(data)
    if data["law_info"]["law_id"] != law_id:
        raise ValueError("Returned law ID does not match the requested ID")
    return envelope(url, data, asof), xml_bytes


def positive_limit(value):
    number = int(value)
    if not 1 <= number <= 20:
        raise argparse.ArgumentTypeError("Use 1..20 for this small example")
    return number


def law_id_only(value):
    if not re.fullmatch(r"[0-9A-Z]{15}", value):
        raise argparse.ArgumentTypeError("Use a 15-character law ID, not a revision ID")
    return value


def positive_index(value):
    number = int(value)
    if number < 1:
        raise argparse.ArgumentTypeError("Supplement index must be 1 or greater")
    return number


def xml_lines(xml_bytes):
    # Formatting is for a review diff only. Original XML bytes are saved separately.
    root = ET.fromstring(xml_bytes)
    for node in root.iter():
        attributes = sorted(node.attrib.items())
        node.attrib.clear()
        node.attrib.update(attributes)
    ET.indent(root, space="  ")
    return [line + "\n" for line in ET.tostring(root, encoding="unicode").splitlines()]


def compare_snapshots(before, after, before_xml, after_xml):
    if before["response"]["law_info"]["law_id"] != after["response"]["law_info"]["law_id"]:
        raise ValueError("Cannot compare different law IDs")
    def summary(snapshot, xml_bytes):
        revision = snapshot["response"]["revision_info"]
        return {"asof": snapshot["requested_asof"],
                "law_revision_id": revision["law_revision_id"],
                "law_title": revision.get("law_title"),
                "request_url": snapshot["request_url"],
                "retrieved_at": snapshot["retrieved_at"],
                "xml_sha256": hashlib.sha256(xml_bytes).hexdigest()}
    left, right = summary(before, before_xml), summary(after, after_xml)
    delta = "".join(difflib.unified_diff(
        xml_lines(before_xml), xml_lines(after_xml),
        fromfile=left["law_revision_id"], tofile=right["law_revision_id"]))
    result = {"schema_version": "egov-comparison/1",
              "law_id": before["response"]["law_info"]["law_id"],
              "before": left, "after": right,
              "same_revision": left["law_revision_id"] == right["law_revision_id"],
              "same_xml_bytes": before_xml == after_xml,
              "same_formatted_xml": not delta,
              "scope": "XML representation comparison; not a legal interpretation"}
    return result, delta


def compare_command(args, parser):
    if args.before >= args.after:
        parser.error("--before must be earlier than --after")
    # Both responses must validate before any comparison files are written.
    before, before_xml = get_snapshot(args.law_id, args.before.isoformat())
    after, after_xml = get_snapshot(args.law_id, args.after.isoformat())
    result, delta = compare_snapshots(before, after, before_xml, after_xml)
    args.out.mkdir(parents=True, exist_ok=True)
    folder = args.out / ("compare-" + now().strftime("%Y%m%dT%H%M%S%fZ"))
    folder.mkdir()
    write_snapshot(folder / "before", before, before_xml)
    write_snapshot(folder / "after", after, after_xml)
    with (folder / "xml.diff").open("x", encoding="utf-8", newline="\n") as file:
        file.write(delta)
    # Written last: summary presence signals that all evidence files were saved.
    write_json(folder / "comparison.json", result)
    print("before:", result["before"]["law_revision_id"])
    print("after:", result["after"]["law_revision_id"])
    print("same_revision:", str(result["same_revision"]).lower())
    print("same_xml_bytes:", str(result["same_xml_bytes"]).lower())
    print("saved:", folder)


def child(parent, tag):
    return next((node for node in parent if local_name(node) == tag), None)


def collect_articles(parent, path):
    for position, node in enumerate(parent, 1):
        tag = local_name(node)
        node_path = f"{path}/{tag}[{position}]"
        if tag == "Article":
            yield node, node_path
        elif tag in CONTAINERS:
            yield from collect_articles(node, node_path)


def xml_node(node):
    """Keep text, inline nodes (Ruby etc.), tail text and all attributes in order."""
    content = []
    if node.text is not None:
        content.append(node.text)
    for item in node:
        content.append(xml_node(item))
        if item.tail is not None:
            content.append(item.tail)
    return {"element": node.tag, "attributes": dict(node.attrib), "content": content}


def extract_article(snapshot, number, supplement=None):
    if not isinstance(snapshot, dict):
        raise ValueError("Expected snapshot to be a JSON object")
    for field in ("request_url", "retrieved_at"):
        if not isinstance(snapshot.get(field), str) or not snapshot[field]:
            raise ValueError(f"Missing source metadata: {field}")
    if "requested_asof" not in snapshot:
        raise ValueError("Missing source metadata: requested_asof")
    data = snapshot["response"]
    xml_bytes, root = validated_xml(data)
    body = child(root, "LawBody")
    if body is None:
        raise ValueError("Missing LawBody")
    if supplement is None:
        provision = child(body, "MainProvision")
        path = "Law/LawBody/MainProvision"
        scope = "main"
    else:
        supplements = [node for node in body if local_name(node) == "SupplProvision"]
        if not 1 <= supplement <= len(supplements):
            raise ValueError("SUPPLEMENT_NOT_FOUND: choose a source XML supplement index")
        provision = supplements[supplement - 1]
        path = f"Law/LawBody/SupplProvision[{supplement}]"
        scope = "supplementary"
    if provision is None:
        raise ValueError("PROVISION_NOT_FOUND")
    matches = [(node, path) for node, path in collect_articles(provision, path)
               if node.get("Num") == number]
    if not matches:
        raise ValueError("ARTICLE_NOT_FOUND: check Num in this revision and provision")
    if len(matches) != 1:
        raise ValueError("AMBIGUOUS_ARTICLE: multiple matching Article elements")
    article, path = matches[0]
    revision = data["revision_info"]["law_revision_id"]
    return {"schema_version": "egov-article/1", "scope": scope,
            "article_number": number, "article": xml_node(article),
            "provision_attributes": dict(provision.attrib),
            "fixture_notice": snapshot.get("fixture_notice"),
            "source": {"law_id": data["law_info"]["law_id"],
                       "law_revision_id": revision,
                       "law_title": data["revision_info"].get("law_title"),
                       "requested_asof": snapshot["requested_asof"],
                       "retrieved_at": snapshot["retrieved_at"],
                       "request_url": snapshot["request_url"],
                       "source_url": (None if snapshot.get("fixture_notice") else
                                      "https://laws.e-gov.go.jp/law/" +
                                      quote(data["law_info"]["law_id"], safe="") + "/" +
                                      quote(revision.split("_", 1)[1], safe="")),
                       "xml_sha256": hashlib.sha256(xml_bytes).hexdigest(),
                       "xml_path": path},
            "notes": ["One Article subtree only; not a complete law or MCP server.",
                      "XML text, content order and attributes are retained, including indentation.",
                      "xml_path uses child positions for review; it is not an XPath expression."]}


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--out", type=Path, default=Path("egov-results"))
    commands = parser.add_subparsers(dest="command", required=True)
    search = commands.add_parser("search", help="Search one page; no auto-download")
    search.add_argument("title")
    search.add_argument("--limit", type=positive_limit, default=3)
    search.add_argument("--offset", type=int, default=0)
    get = commands.add_parser("get", help="Download one explicitly chosen law ID")
    get.add_argument("law_id", type=law_id_only)
    get.add_argument("--asof", type=date.fromisoformat)
    compare = commands.add_parser("compare", help="Save two dates and a review-only XML diff")
    compare.add_argument("law_id", type=law_id_only)
    compare.add_argument("--before", required=True, type=date.fromisoformat)
    compare.add_argument("--after", required=True, type=date.fromisoformat)
    excerpt = commands.add_parser("excerpt", help="Extract one article from a saved JSON, offline")
    excerpt.add_argument("snapshot", type=Path)
    excerpt.add_argument("--article", required=True, help="Exact XML Num, e.g. 1 or 21_2")
    excerpt.add_argument("--supplement", type=positive_index, help="1-based direct SupplProvision index")
    args = parser.parse_args()
    if args.command == "compare":
        compare_command(args, parser)
        return
    if args.command == "excerpt":
        with args.snapshot.open("rb") as file:
            raw = file.read(MAX_BYTES + 1)
        if len(raw) > MAX_BYTES:
            raise ValueError("Snapshot exceeds this example's 12 MiB limit")
        result = extract_article(json.loads(raw), args.article, args.supplement)
        args.out.mkdir(parents=True, exist_ok=True)
        path = args.out / ("excerpt-" + now().strftime("%Y%m%dT%H%M%S%fZ") + ".json")
        write_json(path, result)
        print("article:", result["scope"], result["article_number"])
        print("revision:", result["source"]["law_revision_id"])
        print("saved:", path)
        return
    if args.command == "get":
        snapshot, xml_bytes = get_snapshot(args.law_id, args.asof.isoformat() if args.asof else None)
        args.out.mkdir(parents=True, exist_ok=True)
        stem = args.out / ("get-" + now().strftime("%Y%m%dT%H%M%S%fZ"))
        write_snapshot(stem, snapshot, xml_bytes)
        data = snapshot["response"]
        print(data["law_info"]["law_id"], data["revision_info"].get("law_title", ""))
        print("revision:", data["revision_info"]["law_revision_id"])
    else:
        if args.offset < 0:
            parser.error("offset must be nonnegative")
        url, data = fetch_json("laws", {"response_format": "json", "law_title": args.title,
                                       "limit": args.limit, "offset": args.offset})
        if not isinstance(data.get("laws"), list):
            raise ValueError("Missing laws list; check the API response")
        args.out.mkdir(parents=True, exist_ok=True)
        stem = args.out / ("search-" + now().strftime("%Y%m%dT%H%M%S%fZ"))
        write_json(stem.with_suffix(".json"), envelope(url, data, None))
        print(f"count={data.get('count')} total_count={data.get('total_count')} "
              f"next_offset={data.get('next_offset')}")
        for law in data["laws"]:
            print(law["law_info"]["law_id"], law["revision_info"]["law_title"])
    print("saved:", stem.with_suffix(".json"))


if __name__ == "__main__":
    try:
        main()
    except HTTPError as error:
        print(f"HTTP {error.code}: check parameters, service status and API docs; "
              "no automatic retry", file=sys.stderr)
        sys.exit(1)
    except (URLError, TimeoutError, ValueError, KeyError, TypeError, OSError,
            ET.ParseError, binascii.Error) as error:
        print(f"Stopped: {type(error).__name__}: {error}", file=sys.stderr)
        sys.exit(1)