"""Command line interface for coremate."""

from __future__ import annotations

import argparse
import json
from dataclasses import asdict
from typing import Sequence

from .ai import (
    apply_cuda_env_defaults,
    build_training_plan,
    detect_ai_stack,
    optimize_torch_cuda,
    recommend_training_preset,
    recommend_torch_training_kwargs,
    set_global_seed,
    tune_torch_runtime,
)


def _build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        prog="coremate",
        description="CoreMate - Torch and CUDA helper toolkit for AI workflows.",
    )
    subparsers = parser.add_subparsers(dest="command", required=True)

    ai_parser = subparsers.add_parser("ai", help="AI and CUDA helper commands.")
    ai_subparsers = ai_parser.add_subparsers(dest="ai_command", required=True)

    ai_subparsers.add_parser("report", help="Show torch/CUDA/accelerator report.")

    ai_recommend = ai_subparsers.add_parser("recommend", help="Recommend AI training preset.")
    ai_recommend.add_argument(
        "--task",
        default="general",
        choices=["general", "cv", "nlp", "llm", "inference"],
        help="Workload profile to optimize for.",
    )
    ai_recommend.add_argument(
        "--aggressive",
        action="store_true",
        help="Use more aggressive memory-saving settings.",
    )

    ai_plan = ai_subparsers.add_parser("plan", help="Build an actionable training plan.")
    ai_plan.add_argument(
        "--task",
        default="general",
        choices=["general", "cv", "nlp", "llm", "inference"],
        help="Workload profile to optimize for.",
    )
    ai_plan.add_argument(
        "--model-scale",
        default="base",
        choices=["tiny", "small", "base", "large", "xl"],
        help="Approximate model scale for learning-rate and batch heuristics.",
    )
    ai_plan.add_argument(
        "--target-global-batch",
        type=int,
        default=None,
        help="Optional target global batch size.",
    )
    ai_plan.add_argument(
        "--aggressive",
        action="store_true",
        help="Use more aggressive memory-saving settings.",
    )

    ai_seed = ai_subparsers.add_parser("seed", help="Set deterministic seed helpers.")
    ai_seed.add_argument("--seed", type=int, required=True, help="Seed value.")
    ai_seed.add_argument(
        "--deterministic",
        action="store_true",
        help="Enable deterministic behavior where possible.",
    )

    ai_tune = ai_subparsers.add_parser(
        "tune", help="Apply practical runtime tuning for PyTorch."
    )
    ai_tune.add_argument("--seed", type=int, default=None, help="Optional seed value.")
    ai_tune.add_argument(
        "--deterministic",
        action="store_true",
        help="Enable deterministic behavior where possible.",
    )
    ai_tune.add_argument(
        "--benchmark",
        action=argparse.BooleanOptionalAction,
        default=True,
        help="Enable cudnn benchmark mode for repeated input shapes.",
    )
    ai_tune.add_argument(
        "--allow-tf32",
        action=argparse.BooleanOptionalAction,
        default=True,
        help="Allow TF32 kernels on supported GPUs.",
    )
    ai_tune.add_argument(
        "--matmul-precision",
        default="high",
        choices=["highest", "high", "medium"],
        help="Torch float32 matmul precision policy.",
    )
    ai_tune.add_argument(
        "--num-threads",
        type=int,
        default=None,
        help="Optional torch thread count.",
    )

    ai_env = ai_subparsers.add_parser(
        "env", help="Set CUDA/PyTorch environment defaults in current process."
    )
    ai_env.add_argument(
        "--max-split-size-mb",
        type=int,
        default=256,
        help="Memory block split size for PYTORCH_CUDA_ALLOC_CONF.",
    )
    ai_env.add_argument(
        "--expandable-segments",
        action=argparse.BooleanOptionalAction,
        default=True,
        help="Enable or disable expandable_segments in allocator config.",
    )
    ai_env.add_argument(
        "--cuda-module-loading",
        default="LAZY",
        choices=["LAZY", "EAGER"],
        help="CUDA module loading mode.",
    )

    ai_optimize = ai_subparsers.add_parser(
        "optimize",
        help="Apply CUDA env and Torch tuning in one command.",
    )
    ai_optimize.add_argument(
        "--task",
        default="general",
        choices=["general", "cv", "nlp", "llm", "inference"],
        help="Workload profile to optimize for.",
    )
    ai_optimize.add_argument(
        "--model-scale",
        default="base",
        choices=["tiny", "small", "base", "large", "xl"],
        help="Approximate model scale for learning-rate and batch heuristics.",
    )
    ai_optimize.add_argument(
        "--target-global-batch",
        type=int,
        default=None,
        help="Optional target global batch size.",
    )
    ai_optimize.add_argument(
        "--aggressive",
        action="store_true",
        help="Use more aggressive memory-saving settings.",
    )
    ai_optimize.add_argument("--seed", type=int, default=None, help="Optional seed value.")
    ai_optimize.add_argument(
        "--deterministic",
        action="store_true",
        help="Enable deterministic behavior where possible.",
    )
    ai_optimize.add_argument(
        "--benchmark",
        action=argparse.BooleanOptionalAction,
        default=True,
        help="Enable cudnn benchmark mode for repeated input shapes.",
    )
    ai_optimize.add_argument(
        "--allow-tf32",
        action=argparse.BooleanOptionalAction,
        default=True,
        help="Allow TF32 kernels on supported GPUs.",
    )
    ai_optimize.add_argument(
        "--matmul-precision",
        default="high",
        choices=["highest", "high", "medium"],
        help="Torch float32 matmul precision policy.",
    )
    ai_optimize.add_argument(
        "--num-threads",
        type=int,
        default=None,
        help="Optional torch thread count.",
    )
    ai_optimize.add_argument(
        "--max-split-size-mb",
        type=int,
        default=256,
        help="Memory block split size for PYTORCH_CUDA_ALLOC_CONF.",
    )
    ai_optimize.add_argument(
        "--expandable-segments",
        action=argparse.BooleanOptionalAction,
        default=True,
        help="Enable or disable expandable_segments in allocator config.",
    )
    ai_optimize.add_argument(
        "--cuda-module-loading",
        default="LAZY",
        choices=["LAZY", "EAGER"],
        help="CUDA module loading mode.",
    )

    return parser


def main(argv: Sequence[str] | None = None) -> int:
    parser = _build_parser()
    args = parser.parse_args(list(argv) if argv is not None else None)

    if args.command == "ai":
        if args.ai_command == "report":
            print(json.dumps(detect_ai_stack(), indent=2))
            return 0

        if args.ai_command == "recommend":
            preset = recommend_training_preset(task=args.task, aggressive=bool(args.aggressive))
            print(json.dumps(asdict(preset), indent=2))
            return 0

        if args.ai_command == "plan":
            plan = build_training_plan(
                task=args.task,
                model_scale=args.model_scale,
                aggressive=bool(args.aggressive),
                target_global_batch_size=args.target_global_batch,
            )
            print(
                json.dumps(
                    {
                        "plan": asdict(plan),
                        "training_kwargs": recommend_torch_training_kwargs(plan),
                    },
                    indent=2,
                )
            )
            return 0

        if args.ai_command == "seed":
            print(
                json.dumps(
                    set_global_seed(seed=args.seed, deterministic=bool(args.deterministic)),
                    indent=2,
                )
            )
            return 0

        if args.ai_command == "tune":
            try:
                tuned = tune_torch_runtime(
                    seed=args.seed,
                    deterministic=bool(args.deterministic),
                    benchmark=bool(args.benchmark),
                    allow_tf32=bool(args.allow_tf32),
                    matmul_precision=args.matmul_precision,
                    num_threads=args.num_threads,
                )
            except RuntimeError as error:
                print(json.dumps({"error": str(error)}, indent=2))
                return 2
            print(json.dumps(tuned, indent=2))
            return 0

        if args.ai_command == "env":
            print(
                json.dumps(
                    apply_cuda_env_defaults(
                        max_split_size_mb=args.max_split_size_mb,
                        expandable_segments=bool(args.expandable_segments),
                        cuda_module_loading=args.cuda_module_loading,
                    ),
                    indent=2,
                )
            )
            return 0

        if args.ai_command == "optimize":
            result = optimize_torch_cuda(
                task=args.task,
                model_scale=args.model_scale,
                target_global_batch_size=args.target_global_batch,
                aggressive=bool(args.aggressive),
                seed=args.seed,
                deterministic=bool(args.deterministic),
                benchmark=bool(args.benchmark),
                allow_tf32=bool(args.allow_tf32),
                matmul_precision=args.matmul_precision,
                num_threads=args.num_threads,
                max_split_size_mb=args.max_split_size_mb,
                expandable_segments=bool(args.expandable_segments),
                cuda_module_loading=args.cuda_module_loading,
            )
            print(json.dumps(result, indent=2))
            return 0

    parser.print_help()
    return 1


if __name__ == "__main__":
    raise SystemExit(main())
