from __future__ import annotations

import argparse
import json

from .config import load_source_config
from .contracts import generate_contracts
from .export import ensure_profile, start_export
from .manifest import find_export, write_manifest
from .source_db import inspect_tables


def main() -> None:
    parser = argparse.ArgumentParser(description="Export TiDB Cloud snapshots to S3")
    parser.add_argument("--source", required=True, help="Path to source.yaml")
    subparsers = parser.add_subparsers(dest="command", required=True)

    subparsers.add_parser("generate-contracts")
    subparsers.add_parser("inspect-source")
    subparsers.add_parser("ensure-profile")  # bootstrap the ticloud profile from per-source secrets

    export_parser = subparsers.add_parser("start-export")
    export_parser.add_argument("--run-date")
    export_parser.add_argument("--run-id")

    status_parser = subparsers.add_parser("export-status")
    status_parser.add_argument("--export-id", required=True)

    manifest_parser = subparsers.add_parser("write-manifest")
    manifest_parser.add_argument("--run-date", required=True)
    manifest_parser.add_argument("--run-id", required=True)
    manifest_parser.add_argument("--export-id", required=True)
    manifest_parser.add_argument("--no-upload", action="store_true")

    args = parser.parse_args()
    config = load_source_config(args.source)

    if args.command == "generate-contracts":
        paths = generate_contracts(config)
        print(json.dumps({"contracts": [str(path) for path in paths]}, indent=2))
    elif args.command == "inspect-source":
        tables = inspect_tables(config)
        print(json.dumps({"tables": [{"name": table.name, "primary_key": table.primary_key} for table in tables]}, indent=2))
    elif args.command == "ensure-profile":
        print(json.dumps(ensure_profile(config), indent=2))
    elif args.command == "start-export":
        print(json.dumps(start_export(config, args.run_date, args.run_id), indent=2))
    elif args.command == "export-status":
        export = find_export(config, args.export_id)
        print(json.dumps({"export_id": args.export_id, "state": export.get("state"), "export": export}, indent=2))
    elif args.command == "write-manifest":
        path = write_manifest(config, args.run_date, args.run_id, args.export_id, upload=not args.no_upload)
        print(json.dumps({"manifest": str(path)}, indent=2))


if __name__ == "__main__":
    main()