"""Pre/post-training hook system.

Runs quality checks before training starts and verification after training
completes. Works with any training framework.
"""

from __future__ import annotations

import logging
import subprocess
import sys
from dataclasses import dataclass, field

logger = logging.getLogger(__name__)


@dataclass
class PreTrainingResult:
    """Result of pre-training quality check."""

    passed: bool
    grade: str
    score: int
    summary: str
    cleaned_dataset: str | None = None  # Path if auto_clean was used
    curated_dataset: str | None = None  # Path if auto_curate was used
    warnings: list[str] = field(default_factory=list)


def pre_training_check(
    dataset: str,
    policy: str,
    min_grade: str = "C",
    auto_clean: bool = False,
    auto_curate: int | None = None,
    source_format: str | None = None,
) -> PreTrainingResult:
    """Run quality checks before training starts.

    Args:
        dataset: Dataset path or repo ID.
        policy: Target policy architecture.
        min_grade: Minimum grade to pass.
        auto_clean: Automatically remove bad episodes.
        auto_curate: Automatically select best N episodes.
        source_format: Override dataset format detection.

    Returns:
        PreTrainingResult with pass/fail and optional cleaned dataset path.
    """
    from orbit.utils.display import console

    # Step 1: Run quality gate
    console.print(f"[dim]Running quality gate: {dataset} → {policy}[/dim]")

    try:
        from orbit.analyzer.dataset_loader import load_dataset_info, load_episode_data
        from orbit.analyzer.coverage import compute_coverage, _extract_array
        from orbit.analyzer.episode_ranker import compute_episode_ranking
        from orbit.analyzer.policy_fit import PolicyFitAnalyzer
        from orbit.analyzer.success_predictor import DatasetReadinessGrader
        from orbit.analyzer.task_context import TaskContextAnalyzer

        info = load_dataset_info(dataset)
        episode_data = load_episode_data(dataset, max_episodes=200)

        if episode_data is None or episode_data.empty:
            return PreTrainingResult(
                passed=False, grade="F", score=0,
                summary="Could not load episode data",
            )

        # Extract action episodes
        action_episodes = []
        state_episodes = []
        for _, ep_df in episode_data.groupby("episode_index"):
            if "action" in ep_df.columns:
                arr = _extract_array(ep_df["action"])
                if arr is not None and arr.size > 0:
                    action_episodes.append(arr)
            if "state" in ep_df.columns:
                arr = _extract_array(ep_df["state"])
                if arr is not None and arr.size > 0:
                    state_episodes.append(arr)

        if not action_episodes:
            return PreTrainingResult(
                passed=False, grade="F", score=0,
                summary="No action data found in dataset",
            )

        # Consistency
        consistency_report = None
        if len(action_episodes) >= 2:
            try:
                from orbit.analyzer.consistency_metrics import compute_consistency_metrics
                consistency_report = compute_consistency_metrics(
                    action_episodes,
                    state_episodes=state_episodes if state_episodes else None,
                )
            except Exception:
                pass

        # Analysis
        coverage = compute_coverage(info, episode_data)
        episode_ranking_result = compute_episode_ranking(action_episodes)

        task_context = TaskContextAnalyzer().analyze(
            action_episodes,
            metadata={"robot_type": info.robot_type, "task_description": dataset},
        )

        pf_analyzer = PolicyFitAnalyzer()
        policy_fit_result = pf_analyzer.analyze(
            action_episodes,
            policy_type=policy,
            episode_count=info.num_episodes,
            metadata={"is_bimanual": task_context.is_bimanual},
            task_complexity=task_context.task_complexity_score,
            consistency_score=coverage.episode_consistency,
            multimodality_score=task_context.trajectory_multimodality,
        )

        grader = DatasetReadinessGrader()
        readiness = grader.grade(
            quality_score=50.0,
            episode_consistency=coverage.episode_consistency,
            signal_health={
                "dead_dimensions": [],
                "blockers": [],
                "warnings": [],
                "total_joints": task_context.action_dim,
                "clipping_joints": [],
            },
            policy_fit_score=policy_fit_result.fit_score,
            episode_ranking={
                "total": len(action_episodes),
                "remove_count": episode_ranking_result.remove_count,
                "review_count": episode_ranking_result.review_count,
            },
            action_divergence=0.0,
            episode_count=info.num_episodes,
            policy_type=policy_fit_result.policy_type,
            consistency=consistency_report,
        )

        grade_order = {"A": 4, "B": 3, "C": 2, "D": 1, "F": 0}
        passed = grade_order.get(readiness.grade, 0) >= grade_order.get(min_grade, 2)

    except Exception as e:
        return PreTrainingResult(
            passed=False, grade="F", score=0,
            summary=f"Quality check failed: {e}",
        )

    # Step 2: Auto-clean if requested and failed
    cleaned_dataset = None
    if not passed and auto_clean:
        console.print("[dim]Auto-cleaning dataset...[/dim]")
        try:
            clean_args = ["orbit", "clean", dataset, "--policy", policy, "--aggressive"]
            result = subprocess.run(clean_args, capture_output=True, text=True, timeout=300)
            if result.returncode == 0:
                cleaned_dataset = dataset  # cleaned in-place
                console.print("[green]Dataset cleaned successfully[/green]")
                # Re-check after cleaning
                return pre_training_check(
                    dataset=dataset,
                    policy=policy,
                    min_grade=min_grade,
                    auto_clean=False,  # don't recurse
                    auto_curate=auto_curate,
                    source_format=source_format,
                )
        except Exception as e:
            logger.warning("Auto-clean failed: %s", e)

    # Step 3: Auto-curate if requested
    curated_dataset = None
    if auto_curate is not None and auto_curate > 0:
        console.print(f"[dim]Curating top {auto_curate} episodes...[/dim]")
        try:
            curate_args = [
                "orbit", "curate", dataset,
                "--budget", str(auto_curate),
            ]
            result = subprocess.run(curate_args, capture_output=True, text=True, timeout=300)
            if result.returncode == 0:
                curated_dataset = dataset
        except Exception as e:
            logger.warning("Auto-curate failed: %s", e)

    warnings = []
    if not passed:
        warnings.extend(readiness.problems[:3])

    return PreTrainingResult(
        passed=passed,
        grade=readiness.grade,
        score=readiness.score,
        summary=readiness.summary,
        cleaned_dataset=cleaned_dataset,
        curated_dataset=curated_dataset,
        warnings=warnings,
    )


def run_training_pipeline(
    gate_args: str | None,
    training_command: list[str],
    auto_clean: bool = False,
    skip_gate: bool = False,
    dry_run: bool = False,
) -> None:
    """Run the full training pipeline: gate → user's command → verify.

    Uses passthrough design — ORBIT wraps ANY training command without
    needing to know the framework's CLI.

    Args:
        gate_args: Quality gate arguments, e.g. "act --min-grade B".
            Parsed as "<policy> [--min-grade <grade>]".
        training_command: The user's training command (everything after --).
        auto_clean: Auto-fix data before training.
        skip_gate: Skip the quality gate.
        dry_run: Show what would happen without actually training.
    """
    from orbit.utils.display import console

    train_cmd_str = " ".join(training_command)

    console.print(f"\n[bold]ORBIT Training Pipeline[/bold]")
    console.print(f"  Command: {train_cmd_str}")
    console.print()

    # Step 1: Quality gate
    if not skip_gate and gate_args:
        console.print("[bold]Step 1: Quality Gate[/bold]")

        # Parse gate_args: "<policy> [--min-grade <grade>] [dataset]"
        parts = gate_args.split()
        policy = parts[0] if parts else "act"
        min_grade = "C"
        dataset = None

        i = 1
        while i < len(parts):
            if parts[i] == "--min-grade" and i + 1 < len(parts):
                min_grade = parts[i + 1]
                i += 2
            else:
                dataset = parts[i]
                i += 1

        # Try to auto-detect dataset from training command
        if dataset is None:
            dataset = _detect_dataset_from_command(training_command)

        if dataset is None:
            console.print(
                "[yellow]Could not detect dataset from training command.[/yellow]\n"
                "  Specify explicitly: orbit train --gate 'act --min-grade B /path/to/data' -- ..."
            )
            console.print("[dim]Skipping gate — no dataset to check[/dim]")
        else:
            result = pre_training_check(
                dataset=dataset,
                policy=policy,
                min_grade=min_grade,
                auto_clean=auto_clean,
            )

            status = "[green]PASS[/green]" if result.passed else "[red]FAIL[/red]"
            console.print(f"  {status} — Grade: {result.grade} (score: {result.score})")
            console.print(f"  {result.summary}")

            if not result.passed:
                console.print(f"\n  [red]Training aborted.[/red] Fix data issues first:")
                for w in result.warnings:
                    console.print(f"    - {w}")
                console.print(f"\n  Run: orbit analyze {dataset} --deep")
                sys.exit(1)
    elif not skip_gate and not gate_args:
        console.print("[dim]No --gate args specified, skipping quality gate[/dim]")
    else:
        console.print("[dim]Quality gate skipped (--skip-gate)[/dim]")

    # Step 2: Execute training command
    console.print(f"\n[bold]Step 2: Training[/bold]")
    console.print(f"  [green]{train_cmd_str}[/green]")

    if dry_run:
        console.print(f"\n  [dim]Dry run — would run: {train_cmd_str}[/dim]")
        return

    console.print("[dim]Starting training... (Ctrl+C to stop)[/dim]")

    try:
        proc = subprocess.run(training_command, timeout=86400)
        if proc.returncode != 0:
            console.print(f"\n[red]Training failed with exit code {proc.returncode}[/red]")
            console.print(f"  Run: orbit debug ./outputs/train/ --ai")
            sys.exit(proc.returncode)
    except subprocess.TimeoutExpired:
        console.print("[yellow]Training timed out after 24h[/yellow]")
    except KeyboardInterrupt:
        console.print("\n[yellow]Training interrupted by user[/yellow]")

    # Step 3: Post-training verification
    console.print(f"\n[bold]Step 3: Verification[/bold]")
    training_dir = _detect_latest_output()
    if training_dir:
        console.print(f"[dim]Checking {training_dir}...[/dim]")
        from orbit.verify import print_verification, verify_training

        vresult = verify_training(training_dir)
        print_verification(vresult)
    else:
        console.print("[dim]Could not auto-detect training output directory[/dim]")
        console.print("  Run manually: orbit verify ./outputs/train/<your-run>/")
    console.print()


def _detect_dataset_from_command(cmd: list[str]) -> str | None:
    """Try to detect the dataset path/repo from a training command."""
    for i, arg in enumerate(cmd):
        # LeRobot style: --dataset.repo_id=...
        if "dataset" in arg and "=" in arg:
            return arg.split("=", 1)[1]
        # Generic: --dataset ... or --data_root_dir ...
        if arg in ("--dataset", "--data_root_dir", "--dataset-path", "--dataset_path"):
            if i + 1 < len(cmd):
                return cmd[i + 1]
    return None


def _detect_latest_output() -> str | None:
    """Try to find the latest training output directory."""
    import os
    from pathlib import Path

    candidates = [
        Path("outputs/train"),
        Path("outputs"),
        Path("output"),
        Path("checkpoints"),
    ]
    for base in candidates:
        if not base.exists():
            continue
        # Find most recently modified subdirectory
        subdirs = [d for d in base.iterdir() if d.is_dir()]
        if subdirs:
            latest = max(subdirs, key=lambda d: d.stat().st_mtime)
            return str(latest)
    return None
