"""Export WAA benchmark traces as training data.

This module provides functionality to filter and export successful WAA benchmark
traces in a format suitable for VLM fine-tuning. It converts benchmark execution
traces to the openadapt-ml Episode format.

Usage:
    # Via CLI
    uv run python -m openadapt_evals.benchmarks.vm_cli export-traces --status passed --output training_data/

    # Via Python
    from openadapt_evals.benchmarks.trace_export import export_traces, TraceExporter

    # Export all passing traces
    exporter = TraceExporter(
        benchmark_dir=Path("benchmark_results/waa_eval_20241214"),
        output_dir=Path("training_data"),
        status_filter="passed",
    )
    episodes = exporter.export()

    # Or use convenience function
    episodes = export_traces(
        benchmark_dir="benchmark_results/waa_eval_20241214",
        output_dir="training_data",
        status_filter="passed",
    )

Directory structure created:
    training_data/
    |-- episodes/
    |   |-- episode_001.json       # Episode schema format
    |   |-- episode_002.json
    |   |-- ...
    |-- screenshots/
    |   |-- episode_001/
    |   |   |-- step_000.png
    |   |   |-- step_001.png
    |   |-- episode_002/
    |-- manifest.json              # Index of all exported episodes
    |-- training_samples.jsonl     # JSONL format for training
"""

from __future__ import annotations

import json
import logging
import shutil
from dataclasses import dataclass, field
from datetime import datetime
from pathlib import Path
from typing import Any, Literal

from openadapt_evals.benchmarks.viewer import (
    load_benchmark_metadata,
    load_benchmark_summary,
    load_task_results,
)
from openadapt_ml.schema import (
    Action,
    ActionType,
    BenchmarkSource,
    Coordinates,
    Episode,
    Observation,
    Step,
    save_episode,
)

logger = logging.getLogger(__name__)


StatusFilter = Literal["passed", "failed", "all"]


@dataclass
class ExportStats:
    """Statistics from a trace export operation."""

    total_tasks: int = 0
    exported_tasks: int = 0
    skipped_tasks: int = 0
    total_steps: int = 0
    exported_screenshots: int = 0
    errors: list[str] = field(default_factory=list)


@dataclass
class TraceExporter:
    """Export WAA benchmark traces as training data.

    Filters and converts benchmark execution traces to Episode format,
    copies screenshots, and creates training-ready data files.

    Args:
        benchmark_dir: Path to benchmark results directory containing metadata.json,
            summary.json, and tasks/ subdirectory.
        output_dir: Output directory for exported training data.
        status_filter: Filter by task status ("passed", "failed", "all").
        copy_screenshots: Whether to copy screenshots to output directory.
        create_jsonl: Whether to create training_samples.jsonl file.
        viewport_size: Default viewport size (width, height) for normalizing coordinates.
    """

    benchmark_dir: Path
    output_dir: Path
    status_filter: StatusFilter = "passed"
    copy_screenshots: bool = True
    create_jsonl: bool = True
    viewport_size: tuple[int, int] = (1920, 1200)

    def __post_init__(self):
        self.benchmark_dir = Path(self.benchmark_dir)
        self.output_dir = Path(self.output_dir)

    def export(self) -> list[Episode]:
        """Export traces according to configuration.

        Returns:
            List of Episode objects created from the traces.
        """
        # Load benchmark data
        metadata = load_benchmark_metadata(self.benchmark_dir)
        load_benchmark_summary(self.benchmark_dir)
        tasks = load_task_results(self.benchmark_dir)

        logger.info(
            f"Loaded {len(tasks)} tasks from {self.benchmark_dir.name} "
            f"(model: {metadata.get('model_id', 'unknown')})"
        )

        # Filter tasks
        filtered_tasks = self._filter_tasks(tasks)
        logger.info(f"Filtered to {len(filtered_tasks)} tasks with status={self.status_filter}")

        if not filtered_tasks:
            logger.warning("No tasks match the filter criteria")
            return []

        # Create output directories
        self._setup_output_dirs()

        # Convert and export
        episodes = []
        stats = ExportStats(total_tasks=len(tasks))

        for i, task in enumerate(filtered_tasks):
            try:
                episode = self._convert_task_to_episode(task, i, metadata)
                episodes.append(episode)

                # Save episode JSON
                episode_path = self.output_dir / "episodes" / f"{episode.episode_id}.json"
                save_episode(episode, episode_path)

                # Copy screenshots if enabled
                if self.copy_screenshots:
                    self._copy_task_screenshots(task, episode.episode_id)
                    stats.exported_screenshots += len(task.get("screenshots", []))

                stats.exported_tasks += 1
                stats.total_steps += len(episode.steps)

                logger.debug(
                    f"Exported episode {episode.episode_id}: "
                    f"{len(episode.steps)} steps, success={episode.success}"
                )

            except Exception as e:
                error_msg = f"Failed to export task {task.get('task_id', 'unknown')}: {e}"
                logger.error(error_msg)
                stats.errors.append(error_msg)
                stats.skipped_tasks += 1

        # Create manifest
        self._create_manifest(episodes, metadata, stats)

        # Create JSONL training file
        if self.create_jsonl:
            self._create_training_jsonl(episodes)

        # Log summary
        logger.info(
            f"Export complete: {stats.exported_tasks}/{stats.total_tasks} tasks, "
            f"{stats.total_steps} steps, {stats.exported_screenshots} screenshots"
        )
        if stats.errors:
            logger.warning(f"{len(stats.errors)} errors during export")

        return episodes

    def _filter_tasks(self, tasks: list[dict[str, Any]]) -> list[dict[str, Any]]:
        """Filter tasks by status.

        Args:
            tasks: List of task dictionaries from load_task_results.

        Returns:
            Filtered list of tasks.
        """
        if self.status_filter == "all":
            return tasks

        filtered = []
        for task in tasks:
            execution = task.get("execution", {})
            success = execution.get("success", False)

            if self.status_filter == "passed" and success:
                filtered.append(task)
            elif self.status_filter == "failed" and not success:
                filtered.append(task)

        return filtered

    def _setup_output_dirs(self) -> None:
        """Create output directory structure."""
        self.output_dir.mkdir(parents=True, exist_ok=True)
        (self.output_dir / "episodes").mkdir(exist_ok=True)
        if self.copy_screenshots:
            (self.output_dir / "screenshots").mkdir(exist_ok=True)

    def _convert_task_to_episode(
        self,
        task: dict[str, Any],
        index: int,
        metadata: dict[str, Any],
    ) -> Episode:
        """Convert a benchmark task to Episode format.

        Args:
            task: Task dictionary from load_task_results.
            index: Task index for episode ID generation.
            metadata: Benchmark metadata.

        Returns:
            Episode instance.
        """
        definition = task.get("definition", {})
        execution = task.get("execution", {})
        screenshots = task.get("screenshots", [])
        execution_steps = execution.get("steps", [])

        task_id = task.get("task_id", f"task_{index:03d}")
        episode_id = f"waa_{task_id}"

        # Convert execution steps to Episode steps
        steps = []
        for step_idx, step_data in enumerate(execution_steps):
            step = self._convert_step(step_data, step_idx, screenshots)
            steps.append(step)

        return Episode(
            episode_id=episode_id,
            task_id=task_id,
            instruction=definition.get("instruction", ""),
            goal=definition.get("instruction", ""),
            steps=steps,
            success=execution.get("success", False),
            final_reward=execution.get("score", 0.0),
            source=BenchmarkSource.WAA,
            source_file=str(self.benchmark_dir / "tasks" / task_id),
            agent_model=metadata.get("model_id", "unknown"),
            environment="Windows 11",
            tags=[
                definition.get("domain", "unknown"),
                "waa",
                "benchmark",
            ],
            metadata={
                "benchmark_name": metadata.get("benchmark_name", "waa"),
                "run_name": metadata.get("run_name"),
                "domain": definition.get("domain"),
                "num_steps": execution.get("num_steps", len(steps)),
                "total_time_seconds": execution.get("total_time_seconds"),
                "error": execution.get("error"),
                "reason": execution.get("reason"),
                "evaluation_spec": definition.get("evaluation_spec"),
            },
        )

    def _convert_step(
        self,
        step_data: dict[str, Any],
        step_idx: int,
        screenshots: list[str],
    ) -> Step:
        """Convert a benchmark execution step to Episode Step format.

        Args:
            step_data: Step data from execution.json.
            step_idx: Step index.
            screenshots: List of screenshot paths.

        Returns:
            Step instance.
        """
        action_data = step_data.get("action", {})

        # Build observation
        screenshot_path = None
        if step_idx < len(screenshots):
            screenshot_path = screenshots[step_idx]
        elif step_data.get("screenshot_path"):
            screenshot_path = step_data["screenshot_path"]

        observation = Observation(
            screenshot_path=screenshot_path,
            screen_size=self.viewport_size,
        )

        # Convert action type
        action_type = self._map_action_type(action_data.get("type", "click"))

        # Build action with coordinates
        action_kwargs: dict[str, Any] = {
            "type": action_type,
            "raw": action_data,
        }

        # Handle coordinates - convert to normalized if pixel values
        x = action_data.get("x")
        y = action_data.get("y")
        if x is not None and y is not None:
            # Check if already normalized (0-1 range)
            if 0 <= x <= 1 and 0 <= y <= 1:
                action_kwargs["normalized_coordinates"] = (x, y)
            else:
                # Assume pixel coordinates, normalize
                norm_x = x / self.viewport_size[0]
                norm_y = y / self.viewport_size[1]
                action_kwargs["normalized_coordinates"] = (norm_x, norm_y)
                # Also store pixel coordinates
                action_kwargs["coordinates"] = Coordinates(x=int(x), y=int(y))

        # Handle text for type action
        if action_data.get("text"):
            action_kwargs["text"] = action_data["text"]

        # Handle key for key action
        if action_data.get("key"):
            action_kwargs["key"] = action_data["key"]

        # Handle modifiers
        if action_data.get("modifiers"):
            action_kwargs["modifiers"] = action_data["modifiers"]

        # Handle scroll
        if action_data.get("scroll_direction"):
            action_kwargs["scroll_direction"] = action_data["scroll_direction"]
        if action_data.get("scroll_amount"):
            action_kwargs["scroll_amount"] = int(action_data["scroll_amount"])

        # Handle drag end coordinates
        end_x = action_data.get("end_x")
        end_y = action_data.get("end_y")
        if end_x is not None and end_y is not None:
            if 0 <= end_x <= 1 and 0 <= end_y <= 1:
                action_kwargs["normalized_end"] = (end_x, end_y)
            else:
                norm_end_x = end_x / self.viewport_size[0]
                norm_end_y = end_y / self.viewport_size[1]
                action_kwargs["normalized_end"] = (norm_end_x, norm_end_y)
                action_kwargs["end_coordinates"] = Coordinates(x=int(end_x), y=int(end_y))

        # Handle element targeting
        if action_data.get("target_node_id"):
            from openadapt_ml.schema import UIElement

            action_kwargs["element"] = UIElement(
                element_id=action_data.get("target_node_id"),
                role=action_data.get("target_role"),
                name=action_data.get("target_name"),
            )

        action = Action(**action_kwargs)

        return Step(
            step_index=step_idx,
            observation=observation,
            action=action,
            reasoning=step_data.get("reasoning"),
            timestamp=step_data.get("timestamp"),
        )

    def _map_action_type(self, action_type_str: str) -> ActionType:
        """Map benchmark action type string to ActionType enum.

        Args:
            action_type_str: Action type string from benchmark.

        Returns:
            ActionType enum value.
        """
        mapping = {
            "click": ActionType.CLICK,
            "double_click": ActionType.DOUBLE_CLICK,
            "right_click": ActionType.RIGHT_CLICK,
            "type": ActionType.TYPE,
            "key": ActionType.KEY,
            "scroll": ActionType.SCROLL,
            "drag": ActionType.DRAG,
            "hover": ActionType.HOVER,
            "wait": ActionType.WAIT,
            "done": ActionType.DONE,
            "answer": ActionType.DONE,
            "failed": ActionType.FAIL,
            "fail": ActionType.FAIL,
        }
        return mapping.get(action_type_str.lower(), ActionType.CLICK)

    def _copy_task_screenshots(self, task: dict[str, Any], episode_id: str) -> None:
        """Copy task screenshots to output directory.

        Args:
            task: Task dictionary.
            episode_id: Episode ID for output subdirectory.
        """
        screenshots = task.get("screenshots", [])
        if not screenshots:
            return

        # Create episode screenshot directory
        episode_screenshots_dir = self.output_dir / "screenshots" / episode_id
        episode_screenshots_dir.mkdir(parents=True, exist_ok=True)

        for i, rel_path in enumerate(screenshots):
            src_path = self.benchmark_dir / rel_path
            if src_path.exists():
                dest_path = episode_screenshots_dir / f"step_{i:03d}.png"
                shutil.copy2(src_path, dest_path)

    def _create_manifest(
        self,
        episodes: list[Episode],
        metadata: dict[str, Any],
        stats: ExportStats,
    ) -> None:
        """Create manifest.json with export metadata.

        Args:
            episodes: List of exported episodes.
            metadata: Benchmark metadata.
            stats: Export statistics.
        """
        manifest = {
            "export_timestamp": datetime.utcnow().isoformat(),
            "source_benchmark": metadata.get("benchmark_name", "waa"),
            "source_run": metadata.get("run_name"),
            "source_model": metadata.get("model_id"),
            "status_filter": self.status_filter,
            "statistics": {
                "total_tasks": stats.total_tasks,
                "exported_tasks": stats.exported_tasks,
                "skipped_tasks": stats.skipped_tasks,
                "total_steps": stats.total_steps,
                "exported_screenshots": stats.exported_screenshots,
                "errors": len(stats.errors),
            },
            "episodes": [
                {
                    "episode_id": ep.episode_id,
                    "task_id": ep.task_id,
                    "instruction": ep.instruction,
                    "num_steps": len(ep.steps),
                    "success": ep.success,
                    "domain": ep.metadata.get("domain") if ep.metadata else None,
                }
                for ep in episodes
            ],
        }

        manifest_path = self.output_dir / "manifest.json"
        with open(manifest_path, "w") as f:
            json.dump(manifest, f, indent=2)

        logger.info(f"Created manifest: {manifest_path}")

    def _create_training_jsonl(self, episodes: list[Episode]) -> None:
        """Create JSONL file for training.

        Each line contains a training sample with:
        - instruction: Task instruction
        - screenshot_path: Path to screenshot
        - action: Action taken
        - reasoning: Optional reasoning

        Args:
            episodes: List of exported episodes.
        """
        jsonl_path = self.output_dir / "training_samples.jsonl"

        with open(jsonl_path, "w") as f:
            for episode in episodes:
                for step in episode.steps:
                    sample = {
                        "episode_id": episode.episode_id,
                        "task_id": episode.task_id,
                        "instruction": episode.instruction,
                        "step_index": step.step_index,
                        "screenshot_path": step.observation.screenshot_path,
                        "action_type": step.action.type.value,
                        "action": {
                            "type": step.action.type.value,
                            "coordinates": (
                                {
                                    "x": step.action.coordinates.x,
                                    "y": step.action.coordinates.y,
                                }
                                if step.action.coordinates
                                else None
                            ),
                            "normalized_coordinates": step.action.normalized_coordinates,
                            "text": step.action.text,
                            "key": step.action.key,
                            "modifiers": step.action.modifiers,
                            "scroll_direction": step.action.scroll_direction,
                            "scroll_amount": step.action.scroll_amount,
                        },
                        "reasoning": step.reasoning,
                        "domain": episode.metadata.get("domain") if episode.metadata else None,
                        "success": episode.success,
                    }
                    f.write(json.dumps(sample) + "\n")

        logger.info(f"Created training JSONL: {jsonl_path}")


def export_traces(
    benchmark_dir: str | Path,
    output_dir: str | Path,
    status_filter: StatusFilter = "passed",
    copy_screenshots: bool = True,
    create_jsonl: bool = True,
    viewport_size: tuple[int, int] = (1920, 1200),
) -> list[Episode]:
    """Convenience function to export benchmark traces.

    Args:
        benchmark_dir: Path to benchmark results directory.
        output_dir: Output directory for exported training data.
        status_filter: Filter by task status ("passed", "failed", "all").
        copy_screenshots: Whether to copy screenshots to output directory.
        create_jsonl: Whether to create training_samples.jsonl file.
        viewport_size: Default viewport size for normalizing coordinates.

    Returns:
        List of Episode objects created from the traces.

    Example:
        episodes = export_traces(
            benchmark_dir="benchmark_results/waa_eval_20241214",
            output_dir="training_data",
            status_filter="passed",
        )
        print(f"Exported {len(episodes)} episodes")
    """
    exporter = TraceExporter(
        benchmark_dir=Path(benchmark_dir),
        output_dir=Path(output_dir),
        status_filter=status_filter,
        copy_screenshots=copy_screenshots,
        create_jsonl=create_jsonl,
        viewport_size=viewport_size,
    )
    return exporter.export()


def list_available_runs(
    benchmark_results_dir: str | Path = "benchmark_results",
) -> list[dict[str, Any]]:
    """List available benchmark runs for export.

    Args:
        benchmark_results_dir: Base directory containing benchmark results.

    Returns:
        List of dictionaries with run information.
    """
    results_dir = Path(benchmark_results_dir)
    if not results_dir.exists():
        return []

    runs = []
    for run_dir in sorted(results_dir.iterdir()):
        if not run_dir.is_dir():
            continue

        metadata_path = run_dir / "metadata.json"
        summary_path = run_dir / "summary.json"

        run_info = {
            "run_name": run_dir.name,
            "path": str(run_dir),
        }

        if metadata_path.exists():
            with open(metadata_path) as f:
                metadata = json.load(f)
            run_info.update(
                {
                    "benchmark_name": metadata.get("benchmark_name"),
                    "model_id": metadata.get("model_id"),
                    "created_at": metadata.get("created_at"),
                }
            )

        if summary_path.exists():
            with open(summary_path) as f:
                summary = json.load(f)
            run_info.update(
                {
                    "num_tasks": summary.get("num_tasks", 0),
                    "num_success": summary.get("num_success", 0),
                    "success_rate": summary.get("success_rate", 0.0),
                }
            )

        runs.append(run_info)

    return runs
