"""Playbook generator for data collection planning."""

from __future__ import annotations

import json
from dataclasses import asdict, dataclass, field
from pathlib import Path

TEMPLATES_DIR = Path(__file__).parent / "templates"
ROBOT_CONFIGS_DIR = Path(__file__).parent / "robot_configs"


@dataclass
class Phase:
    """A single phase of the data collection playbook."""

    name: str
    target_episodes: int
    requirements: dict[str, str]


@dataclass
class Estimates:
    """Episode estimates for reaching success thresholds."""

    eighty_percent_success: str
    ninety_five_percent_success: str


@dataclass
class RobotConfig:
    """Robot-specific configuration and tips."""

    robot_type: str
    workspace_dims: list[float]
    recommended_camera_position: str
    recommended_fps: int
    recommended_resolution: list[int]
    arm_count: int
    notes: str


@dataclass
class Playbook:
    """A complete data collection playbook."""

    task_description: str
    task_type: str
    robot_type: str
    policy_type: str
    phases: list[Phase]
    estimates: Estimates
    robot_config: RobotConfig | None = None
    total_target_episodes: int = 0
    policy_notes: list[str] = field(default_factory=list)

    def to_dict(self) -> dict:
        """Convert to a JSON-serializable dict."""
        d = asdict(self)
        return d

    def save(self, path: str) -> None:
        """Save the playbook to a JSON file."""
        with open(path, "w") as f:
            json.dump(self.to_dict(), f, indent=2)


def _load_templates() -> list[dict]:
    """Load all task templates from the templates directory."""
    templates = []
    for p in sorted(TEMPLATES_DIR.glob("*.json")):
        with open(p) as f:
            templates.append(json.load(f))
    return templates


def _match_template(task_description: str, templates: list[dict]) -> dict:
    """Match a task description to the best template using keyword scoring."""
    desc_lower = task_description.lower()
    best_template = None
    best_score = 0

    for template in templates:
        if template["task_type"] == "generic":
            continue
        score = sum(1 for kw in template["keywords"] if kw in desc_lower)
        if score > best_score:
            best_score = score
            best_template = template

    if best_template is None or best_score == 0:
        # Fall back to generic
        for template in templates:
            if template["task_type"] == "generic":
                return template
        # Should never happen, but just in case
        raise ValueError("No generic template found")

    return best_template


def _load_robot_config(robot_type: str) -> RobotConfig | None:
    """Load robot configuration from config files."""
    config_path = ROBOT_CONFIGS_DIR / f"{robot_type}.json"
    if not config_path.exists():
        return None
    with open(config_path) as f:
        data = json.load(f)
    return RobotConfig(
        robot_type=data["robot_type"],
        workspace_dims=data["workspace_dims"],
        recommended_camera_position=data["recommended_camera_position"],
        recommended_fps=data["recommended_fps"],
        recommended_resolution=data["recommended_resolution"],
        arm_count=data["arm_count"],
        notes=data["notes"],
    )


def _apply_policy_adjustments(
    phases: list[Phase],
    estimates: Estimates,
    policy_type: str,
) -> tuple[list[Phase], Estimates, list[str]]:
    """Adjust phase targets and provide notes based on policy type."""
    policy_lower = policy_type.lower()
    notes: list[str] = []

    if policy_lower == "act":
        notes.append("ACT: Recommend 50+ demonstrations minimum")
        notes.append("ACT: Use 30fps recording for smooth action prediction")
        notes.append("ACT: High-resolution cameras recommended (480x640+)")
        notes.append("ACT: Consistent camera placement is critical")
        # ACT works well with fewer demos but needs quality
        for phase in phases:
            phase.target_episodes = max(50, phase.target_episodes)

    elif policy_lower in ("diffusion", "diffusion_policy", "diffusion-policy"):
        notes.append("Diffusion Policy: Recommend 100+ demonstrations")
        notes.append("Diffusion Policy: Can work with lower FPS (10fps OK)")
        notes.append("Diffusion Policy: Benefits from diverse demonstrations")
        notes.append("Diffusion Policy: Multi-modal action distributions handled well")
        # Diffusion policy benefits from more data
        for phase in phases:
            phase.target_episodes = int(phase.target_episodes * 1.5)
        estimates = Estimates(
            eighty_percent_success="150-250 episodes",
            ninety_five_percent_success="350-500 episodes",
        )

    elif policy_lower in ("smolvla", "vla"):
        notes.append("SmolVLA: Recommend 50+ demonstrations")
        notes.append("SmolVLA: Single camera is sufficient")
        notes.append("SmolVLA: Language-conditioned — use clear, consistent task descriptions")
        notes.append("SmolVLA: Pre-trained visual features help with fewer demos")
        # SmolVLA can work with fewer demos
        for phase in phases:
            phase.target_episodes = max(40, phase.target_episodes)

    else:
        notes.append(f"Policy '{policy_type}': Using default episode targets")
        notes.append("Adjust targets based on your policy's data requirements")

    return phases, estimates, notes


def _apply_robot_adjustments(
    phases: list[Phase],
    robot_config: RobotConfig | None,
) -> list[Phase]:
    """Adjust phases based on robot configuration."""
    if robot_config is None:
        return phases

    # Bimanual robots need more demonstrations
    if robot_config.arm_count >= 2:
        for phase in phases:
            phase.target_episodes = int(phase.target_episodes * 1.5)
            phase.requirements["bimanual"] = (
                "Coordinate both arms — practice bimanual handoffs and synchronized motions"
            )

    return phases


def generate_playbook(
    task_description: str,
    robot_type: str,
    policy_type: str,
) -> Playbook:
    """Generate a data collection playbook.

    Args:
        task_description: Natural language description of the task.
        robot_type: Robot identifier (e.g., "so100", "aloha").
        policy_type: Policy type (e.g., "act", "diffusion_policy", "smolvla").

    Returns:
        A Playbook with customized phases and recommendations.
    """
    templates = _load_templates()
    template = _match_template(task_description, templates)

    # Build phases from template
    phases = [
        Phase(
            name=p["name"],
            target_episodes=p["target_episodes"],
            requirements=dict(p["requirements"]),
        )
        for p in template["phases"]
    ]
    estimates = Estimates(
        eighty_percent_success=template["estimates"]["80_percent_success"],
        ninety_five_percent_success=template["estimates"]["95_percent_success"],
    )

    # Load robot config and apply adjustments
    robot_config = _load_robot_config(robot_type)
    phases = _apply_robot_adjustments(phases, robot_config)

    # Apply policy-specific adjustments
    phases, estimates, policy_notes = _apply_policy_adjustments(
        phases,
        estimates,
        policy_type,
    )

    total_episodes = sum(p.target_episodes for p in phases)

    return Playbook(
        task_description=task_description,
        task_type=template["task_type"],
        robot_type=robot_type,
        policy_type=policy_type,
        phases=phases,
        estimates=estimates,
        robot_config=robot_config,
        total_target_episodes=total_episodes,
        policy_notes=policy_notes,
    )
