"""Any4Robot command line interface."""

from __future__ import annotations

import argparse
import sys
from collections import defaultdict
from pathlib import Path
from typing import Any, Dict, List

from any4robot.qualitycheck import check_lerobot_v2_1_quality, delete_invalid_episodes
from any4robot.lerobot_editor.constants import header, highlight, info, success, warning, error


def _format_ratio(value: Any) -> str:
    if value is None:
        return "-"
    try:
        return f"{value:.3f}"
    except Exception:
        return str(value)


def _print_task_distribution(task_counts: Dict[str, int], total: int) -> None:
    print("\n" + "=" * 70)
    print(header("TASK DISTRIBUTION"))
    print("=" * 70)

    if not task_counts:
        print(warning("No task counts available."))
        return

    sorted_tasks = sorted(task_counts.items(), key=lambda item: -item[1])
    print(f"{'Task':50s} | {'Count':>8s} | {'Ratio':>8s}")
    print("-" * 70)
    for task, count in sorted_tasks:
        ratio = count / total if total else 0
        print(f"{task[:50]:50s} | {count:8d} | {ratio:7.2%}")


def _group_invalid_by_dataset(invalid_episodes: List[Dict[str, Any]]) -> Dict[str, List[Dict[str, Any]]]:
    grouped: Dict[str, List[Dict[str, Any]]] = defaultdict(list)
    for entry in invalid_episodes:
        dataset = entry.get("dataset_path", "UNKNOWN")
        grouped[dataset].append(entry)
    return grouped


def _print_invalid_report(report: Dict[str, Any]) -> None:
    invalid_episodes = report.get("invalid_episodes", [])
    grouped = _group_invalid_by_dataset(invalid_episodes)

    print("\n" + "=" * 70)
    print(header("INVALID EPISODES"))
    print("=" * 70)

    if not invalid_episodes:
        print(success("No invalid episodes found."))
        return

    for dataset_path, entries in grouped.items():
        dataset_name = Path(dataset_path).name
        print("\n" + highlight(f"Dataset: {dataset_name}"))
        print(info(f"Path: {dataset_path}"))
        print(info(f"Invalid episodes: {len(entries)}"))

        entries_sorted = sorted(entries, key=lambda item: item.get("episode_index", -1))
        for entry in entries_sorted:
            episode_index = entry.get("episode_index")
            task = entry.get("task")
            length = entry.get("length")
            avg_len = entry.get("avg_length")
            ratio = entry.get("ratio")
            errors = entry.get("errors", [])
            error_str = ", ".join(errors) if errors else "UNKNOWN"

            print(
                f"  {highlight(str(entry.get('episode_name', episode_index)))} "
                f"task={task} length={length} avg={avg_len} ratio={_format_ratio(ratio)}"
            )
            print(f"    {warning(error_str)}")


def _print_summary(report: Dict[str, Any]) -> None:
    print("\n" + "=" * 70)
    print(header("SUMMARY"))
    print("=" * 70)

    print(f"Total episodes      : {report.get('total_episodes', 0)}")
    print(f"Total invalid       : {report.get('total_invalid', 0)}")
    print(f"Invalid ratio       : {report.get('invalid_ratio', 0.0):.6f}")
    print(f"Min ratio threshold : {report.get('min_ratio', 0.0)}")
    print(f"Max ratio threshold : {report.get('max_ratio', 0.0)}")


def _print_dataset_errors(report: Dict[str, Any]) -> None:
    dataset_errors = report.get("dataset_errors", [])
    if not dataset_errors:
        return

    print("\n" + "=" * 70)
    print(header("DATASET ERRORS"))
    print("=" * 70)

    for entry in dataset_errors:
        dataset = entry.get("dataset", "UNKNOWN")
        err = entry.get("error", "UNKNOWN")
        print(error(f"{dataset}: {err}"))


def _run_check(args: argparse.Namespace) -> int:
    report = check_lerobot_v2_1_quality(
        args.path,
        min_ratio=args.thresh,
        max_ratio=args.max_ratio,
    )

    print(header("Any4Robot Quality Check"))
    print(info(f"Root: {args.path}"))

    _print_summary(report)
    _print_dataset_errors(report)
    _print_invalid_report(report)
    _print_task_distribution(report.get("task_counts", {}), report.get("total_episodes", 0))

    if args.delete:
        deletion = delete_invalid_episodes(
            args.path,
            report=report,
            dry_run=args.dry_run,
        )
        deleted = deletion.get("deleted", [])
        failed = deletion.get("failed", [])

        print("\n" + "=" * 70)
        print(header("DELETE RESULTS"))
        print("=" * 70)

        if args.dry_run:
            print(warning("Dry run enabled. No files deleted."))

        if deleted:
            print(success(f"Deleted episodes: {len(deleted)}"))
        if failed:
            print(error(f"Failed to delete: {len(failed)}"))

    return 0


def _build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Any4Robot CLI",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  any4robot check /data/folder
  any4robot check /data/folder --delete --thresh 0.25
  any4robot check /data/folder --delete --thresh 0.25 --dry-run
        """,
    )

    subparsers = parser.add_subparsers(dest="command", required=True)

    check_parser = subparsers.add_parser("check", help="Check dataset quality")
    check_parser.add_argument("path", help="Dataset root or parent folder")
    check_parser.add_argument(
        "--thresh",
        type=float,
        default=0.25,
        help="Minimum length ratio vs task average to accept (default: 0.25)",
    )
    check_parser.add_argument(
        "--max-ratio",
        type=float,
        default=4.0,
        help="Maximum length ratio vs task average to accept (default: 4.0)",
    )
    check_parser.add_argument(
        "--delete",
        action="store_true",
        help="Delete invalid episodes after checking",
    )
    check_parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Show what would be deleted without deleting",
    )
    check_parser.set_defaults(func=_run_check)

    return parser


def main(argv: List[str] | None = None) -> int:
    parser = _build_parser()
    args = parser.parse_args(argv)
    return args.func(args)


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