"""Progress tracking: compare dataset against a collection plan."""

from __future__ import annotations

import json
from dataclasses import dataclass, field

from orbit.analyzer.coverage import FullCoverageReport
from orbit.analyzer.dataset_loader import DatasetInfo
from orbit.planner.playbook import Estimates, Phase, Playbook, RobotConfig


@dataclass
class ChecklistItem:
    """A single checklist item for progress tracking."""

    name: str
    status: str  # "complete" | "in_progress" | "not_started"
    current: float
    target: float
    description: str


@dataclass
class TrackingReport:
    """Progress tracking report comparing dataset against plan."""

    total_target_episodes: int
    current_episodes: int
    progress_percent: float
    phase: str  # "core" or "diversity"
    coverage_checklist: list[ChecklistItem] = field(default_factory=list)
    ready_to_train: bool = False
    next_session_priorities: list[str] = field(default_factory=list)


# Mapping from plan requirement keys to coverage metrics
_REQUIREMENT_TO_METRIC = {
    "position_quadrants": "position_diversity",
    "position_variety": "position_diversity",
    "object_positions": "position_diversity",
    "block_positions": "position_diversity",
    "container_positions": "position_diversity",
    "start_positions": "position_diversity",
    "surface_coverage": "position_diversity",
    "object_variety": "action_diversity",
    "grasp_angles": "action_diversity",
    "approach_angles": "action_diversity",
    "alignment_angles": "action_diversity",
    "direction_variety": "action_diversity",
    "pressure_variety": "action_diversity",
    "pour_amount": "action_diversity",
    "stack_heights": "action_diversity",
    "placement_variety": "temporal_coverage",
    "placement_precision": "temporal_coverage",
    "speed": "episode_consistency",
    "speed_variety": "episode_consistency",
    "failures": "failure_recovery",
    "misalignment_recovery": "failure_recovery",
    "edge_cases": "temporal_coverage",
    "lighting": "temporal_coverage",
    "rotated_objects": "action_diversity",
    "rotated_targets": "action_diversity",
}

# Target thresholds for coverage metrics
_METRIC_TARGETS = {
    "position_diversity": 0.60,
    "action_diversity": 0.50,
    "episode_consistency": 0.70,
    "temporal_coverage": 0.50,
    "failure_recovery": 3,  # number of recovery episodes
}


def _estimate_failure_count(coverage: FullCoverageReport, info: DatasetInfo) -> int:
    """Estimate number of failure recovery episodes from coverage metrics."""
    # Heuristic: if episode consistency is very high, likely no failure demos
    if coverage.episode_consistency > 0.95:
        return 0
    # If there are outlier episodes, some may be failure recoveries
    outlier_count = len(coverage.outlier_episodes)
    if outlier_count > 0:
        return min(outlier_count, 5)
    # Otherwise estimate based on consistency deviation
    if coverage.episode_consistency < 0.85:
        return max(1, int((0.95 - coverage.episode_consistency) * 20))
    return 1


def _build_checklist(
    plan: Playbook,
    info: DatasetInfo,
    coverage: FullCoverageReport,
) -> list[ChecklistItem]:
    """Build a checklist comparing plan requirements against coverage metrics."""
    items: list[ChecklistItem] = []
    seen_metrics: set[str] = set()

    # Collect all requirements from all phases
    all_requirements: dict[str, str] = {}
    for phase in plan.phases:
        all_requirements.update(phase.requirements)

    failure_count = _estimate_failure_count(coverage, info)

    for req_key, req_desc in all_requirements.items():
        metric_name = _REQUIREMENT_TO_METRIC.get(req_key)
        if metric_name is None or metric_name in seen_metrics:
            continue
        seen_metrics.add(metric_name)

        if metric_name == "failure_recovery":
            # Parse target from description (e.g., "Include 5+ near-failure...")
            failure_target = 3
            for word in req_desc.split():
                if word.endswith("+") and word[:-1].isdigit():
                    failure_target = int(word[:-1])
                    break
            current = float(failure_count)
            status = (
                "complete"
                if current >= failure_target
                else ("in_progress" if current > 0 else "not_started")
            )
            items.append(
                ChecklistItem(
                    name="Near-failure recoveries",
                    status=status,
                    current=current,
                    target=float(failure_target),
                    description=req_desc,
                )
            )
        else:
            target = float(_METRIC_TARGETS.get(metric_name, 0.50))
            current = getattr(coverage, metric_name, 0.0)
            status = (
                "complete"
                if current >= target
                else ("in_progress" if current >= target * 0.5 else "not_started")
            )
            display_name = metric_name.replace("_", " ").title()
            items.append(
                ChecklistItem(
                    name=display_name,
                    status=status,
                    current=round(current, 2),
                    target=target,
                    description=req_desc,
                )
            )

    return items


def _generate_priorities(
    checklist: list[ChecklistItem],
    current_episodes: int,
    total_target: int,
) -> list[str]:
    """Generate top 3 next-session priorities from the largest gaps."""
    gaps: list[tuple[float, str]] = []

    for item in checklist:
        if item.status == "complete":
            continue
        if item.name == "Near-failure recoveries":
            gap = item.target - item.current
            gaps.append(
                (
                    gap / max(1, item.target),
                    f"Record {int(gap)} more near-failure recovery demonstrations",
                )
            )
        else:
            gap = item.target - item.current
            gaps.append(
                (
                    gap / max(0.01, item.target),
                    f"Collect episodes focusing on {item.name.lower()} "
                    f"(currently {item.current:.2f}, target {item.target:.2f})",
                )
            )

    # Add episode count gap if applicable
    if current_episodes < total_target:
        ep_gap = total_target - current_episodes
        gaps.append(
            (
                ep_gap / max(1, total_target),
                f"Continue until {total_target} total episodes reached ({ep_gap} more needed)",
            )
        )

    # Sort by relative gap descending, take top 3
    gaps.sort(key=lambda x: x[0], reverse=True)
    return [desc for _, desc in gaps[:3]]


def compute_progress(
    plan: Playbook,
    info: DatasetInfo,
    coverage: FullCoverageReport,
) -> TrackingReport:
    """Compare a dataset against a collection plan.

    Args:
        plan: The collection playbook.
        info: Dataset metadata.
        coverage: Full coverage analysis results.

    Returns:
        TrackingReport with progress, checklist, and priorities.
    """
    total_target = plan.total_target_episodes
    current = info.num_episodes
    progress = min(100.0, (current / max(1, total_target)) * 100)

    # Determine current phase
    core_target = plan.phases[0].target_episodes if plan.phases else total_target
    phase = "core" if current < core_target else "diversity"

    checklist = _build_checklist(plan, info, coverage)

    # Ready to train if we have enough episodes and all checklist items pass
    all_complete = all(item.status == "complete" for item in checklist)
    ready = current >= total_target and all_complete

    priorities = _generate_priorities(checklist, current, total_target)

    return TrackingReport(
        total_target_episodes=total_target,
        current_episodes=current,
        progress_percent=round(progress, 1),
        phase=phase,
        coverage_checklist=checklist,
        ready_to_train=ready,
        next_session_priorities=priorities,
    )


def load_plan_from_file(path: str) -> Playbook:
    """Load a Playbook from a saved JSON file.

    Args:
        path: Path to the JSON plan file.

    Returns:
        A Playbook instance.

    Raises:
        FileNotFoundError: If the plan file doesn't exist.
        ValueError: If the JSON is invalid.
    """
    with open(path) as f:
        data = json.load(f)

    phases = [
        Phase(
            name=p["name"],
            target_episodes=p["target_episodes"],
            requirements=p["requirements"],
        )
        for p in data["phases"]
    ]

    estimates = Estimates(
        eighty_percent_success=data["estimates"]["eighty_percent_success"],
        ninety_five_percent_success=data["estimates"]["ninety_five_percent_success"],
    )

    robot_config = None
    if data.get("robot_config"):
        rc = data["robot_config"]
        robot_config = RobotConfig(
            robot_type=rc["robot_type"],
            workspace_dims=rc["workspace_dims"],
            recommended_camera_position=rc["recommended_camera_position"],
            recommended_fps=rc["recommended_fps"],
            recommended_resolution=rc["recommended_resolution"],
            arm_count=rc["arm_count"],
            notes=rc["notes"],
        )

    return Playbook(
        task_description=data["task_description"],
        task_type=data["task_type"],
        robot_type=data["robot_type"],
        policy_type=data["policy_type"],
        phases=phases,
        estimates=estimates,
        robot_config=robot_config,
        total_target_episodes=data["total_target_episodes"],
        policy_notes=data.get("policy_notes", []),
    )
