from __future__ import annotations

import json
import os
from pathlib import Path
from typing import Any, Dict, Iterator, List, Optional

from harbor.agents.factory import AgentFactory
from harbor.models.agent.name import AgentName

"""
Trace extraction and conversion utilities to provide reusable helpers for
turning sandbox episode traces into HF Datasets‑ready artifacts.

Schema of exported rows (per episode):
    - conversations: list of {"role": str, "content": str}
        - Built from trajectory.json by extracting steps up to each agent turn.
        - Roles typically include "system", "user", and "assistant".
    - agent: str                # agent name (e.g., "terminus-2")
    - model: str                # underlying model name
    - model_provider: str       # model provider id
    - date: str                 # ISO start time of the run
    - task: Optional[str]       # task name from run metadata
    - episode: str              # episode identifier (e.g., "episode-0")
    - run_id: str               # job/run identifier
    - trial_name: Optional[str] # trial name associated with the run
    - result: Optional[str]     # reward bucket (e.g., "1.0") or exception name
    - tool_definitions: Optional[list] # tool definitions available to the agent
    - instruction: Optional[str] # task description (included when requested)
    - verifier_output: Optional[str] # concatenated verifier stdout/stderr (optional)

Notes and options:
    - Only trials that contain agent/trajectory.json are discovered.
    - Episodes are inferred from agent steps in the trajectory.
    - Success filtering can include/exclude trials based on reward in result.json.
    - If to_sharegpt=True, a "conversations_sharegpt" column is added with the
      ShareGPT-style [{"from": "human|gpt|system", "value": str}] messages.
    - rows_to_dataset() returns a datasets.Dataset built from the rows; optional
      push to Hub via push_dataset().
"""

try:
    from datasets import Dataset
except Exception:  # pragma: no cover - optional dep at import time
    Dataset = None  # type: ignore


class MultimodalExportError(Exception):
    """Raised when attempting to export multimodal trajectories for text-only training.

    This error is raised explicitly to prevent silent data loss when multimodal
    trajectories (containing images) are processed by text-only export pipelines.
    """

    pass


_RESULT_JSON_CACHE: dict[Path, Any] = {}


def _read_json_cached(path: Path) -> Any:
    """Read JSON from path with caching and graceful error handling."""
    resolved = path.resolve()
    if resolved in _RESULT_JSON_CACHE:
        return _RESULT_JSON_CACHE[resolved]
    try:
        data = json.loads(path.read_text())
    except (json.JSONDecodeError, OSError):
        data = None
    _RESULT_JSON_CACHE[resolved] = data
    return data


# --------------------
# Multimodal detection
# --------------------


def _content_has_images(content: Any) -> bool:
    """Check if content contains image parts."""
    if isinstance(content, list):
        for part in content:
            if isinstance(part, dict) and part.get("type") == "image":
                return True
    return False


def _step_has_multimodal_content(step: dict) -> bool:
    """Check if a step contains multimodal content."""
    message = step.get("message")
    if _content_has_images(message):
        return True

    observation = step.get("observation")
    if observation:
        for result in observation.get("results", []):
            if _content_has_images(result.get("content")):
                return True
    return False


def _trajectory_has_multimodal_content(trajectory_data: dict) -> bool:
    """Check if trajectory contains any multimodal content."""
    # Scan steps for multimodal content
    for step in trajectory_data.get("steps", []):
        if _step_has_multimodal_content(step):
            return True
    return False


# --------------------
# Message normalization
# --------------------


def normalize_message_content(content: Any) -> str:
    """Return a best-effort textual content from various message encodings.

    Handles:
    - str
    - list[dict] with {'text': '...'} entries (first element preferred)
    - arbitrary serializable objects (json-dumped)
    """
    if content is None:
        return ""
    if isinstance(content, str):
        return content
    if isinstance(content, list) and content:
        first = content[0]
        if isinstance(first, dict) and "text" in first:
            val = first.get("text")
            return val if isinstance(val, str) else json.dumps(val, ensure_ascii=False)
    try:
        return json.dumps(content, ensure_ascii=False)
    except TypeError:
        return str(content)


# --------------------
# ShareGPT conversion
# --------------------


def openai_to_sharegpt(messages: list) -> list:
    role_map = {"user": "human", "assistant": "gpt", "system": "system"}
    out = []
    for m in messages:
        role = m.get("role")
        content = normalize_message_content(m.get("content"))
        if role not in role_map:
            continue
        out.append({"from": role_map[role], "value": content})
    return out


def convert_openai_to_sharegpt(
    dataset: "Dataset", conversations_column: str, output_column: str
) -> "Dataset":
    def f(row: dict):
        row[output_column] = openai_to_sharegpt(row.get(conversations_column) or [])
        return row

    return dataset.map(f)


# --------------------
# Discovery helpers
# --------------------


def is_trial_dir(p: Path) -> bool:
    """Heuristic: a trial directory has an agent directory with episodes."""
    agent = p / "agent"
    return agent.exists()


def iter_trial_dirs(root: Path, recursive: bool = True) -> Iterator[Path]:
    """Yield directories that look like trial outputs.

    If root itself is a trial dir, yield it once. Otherwise, recursively search
    for subdirectories containing an 'agent' subdirectory.
    """
    root = Path(root)
    if is_trial_dir(root):
        yield root
        return
    if not recursive:
        return
    for p in root.rglob("*"):
        if p.is_dir() and is_trial_dir(p):
            yield p


# --------------------
# Extraction logic
# --------------------


def _normalize_run_metadata(raw: Dict[str, Any]) -> Dict[str, Any]:
    """Extract the fields required by trace export from the sandbox result.json blob."""
    config = raw["config"]
    # Newer configs store agents under `agents`; take the first entry.
    agent_cfg = config.get("agent") or config["agents"][0]
    agent_info = raw.get("agent_info") or {}
    model_info = agent_info.get("model_info") or {
        "name": agent_cfg.get("model_name"),
        "provider": agent_cfg.get("provider") or config.get("engine"),
    }
    return {
        "agent_name": agent_cfg["name"],
        "model_name": model_info["name"],
        "model_provider": model_info["provider"],
        "start_time": raw["started_at"],
        "run_id": config.get("job_id")
        or config.get("job_name")
        or raw.get("trial_name")
        or "unknown",
        "task_name": raw["task_name"],
        "trial_name": raw["trial_name"],
        "raw_metadata": raw,
    }


def _load_result_data(trial_dir: Path) -> Dict[str, Any] | None:
    """Load and cache the per-trial result.json (metadata) file."""
    result_path = trial_dir / "result.json"
    if not result_path.exists():
        return None
    data = _read_json_cached(result_path)
    return data if isinstance(data, dict) else None


def _load_job_result_data(trial_dir: Path) -> Dict[str, Any] | None:
    """Search upwards for the job-level result.json that contains aggregate stats."""
    candidate = trial_dir.parent / "result.json"
    if not candidate.exists():
        return None
    data = _read_json_cached(candidate)
    if isinstance(data, dict) and "stats" in data:
        return data
    return None


def _extract_instruction(trial_dir: Path, agent_name: str) -> Optional[str]:
    """Return the initial instruction text extracted from the main trajectory."""
    agent_dir = trial_dir / "agent"
    if not agent_dir.exists():
        return None

    # Prefer the primary trajectory.json; fall back to the earliest continuation.
    trajectory_paths = sorted(
        path
        for path in agent_dir.glob("trajectory*.json")
        if ".summarization" not in path.name  # skip subagent summaries
    )
    if not trajectory_paths:
        return None

    try:
        trajectory = _read_json_cached(trajectory_paths[0])
    except OSError:
        return None

    steps = trajectory.get("steps")
    if not isinstance(steps, list):
        return None

    for step in steps:
        if not isinstance(step, dict):
            continue
        message = step.get("message")
        if isinstance(message, str) and message.strip():
            return message.strip()

    return None


def _read_verifier_output(trial_dir: Path) -> Optional[str]:
    """Read verifier stdout/stderr files and concatenate them when present."""
    verifier_dir = trial_dir / "verifier"
    stdout_path = verifier_dir / "test-stdout.txt"
    stderr_path = verifier_dir / "test-stderr.txt"

    parts: List[str] = []
    for path in (stdout_path, stderr_path):
        if not path.exists():
            continue
        try:
            parts.append(path.read_text())
        except OSError:
            continue

    if not parts:
        return None

    return "".join(parts)


def _coerce_reward_value(value: Any) -> str | Any:
    """Attempt to parse a numeric reward and return as string, else original."""
    try:
        return str(float(value))
    except (TypeError, ValueError):
        return value


def _extract_trial_result_value(trial_dir: Path, trial_name: str) -> Optional[str]:
    """Return exception type or reward bucket for a given trial."""
    result_data = _load_job_result_data(trial_dir)
    if not result_data:
        return None

    stats = result_data.get("stats")
    if not isinstance(stats, dict):
        return None

    evals = stats.get("evals")
    if not isinstance(evals, dict):
        return None

    for eval_data in evals.values():
        if not isinstance(eval_data, dict):
            continue

        exception_stats = eval_data.get("exception_stats")
        if isinstance(exception_stats, dict):
            for exc_name, entries in exception_stats.items():
                if isinstance(entries, list) and trial_name in entries:
                    return exc_name

        reward_stats = eval_data.get("reward_stats")
        if isinstance(reward_stats, dict):
            reward_map = reward_stats.get("reward")
            if isinstance(reward_map, dict):
                for reward_value, entries in reward_map.items():
                    if isinstance(entries, list) and trial_name in entries:
                        result_value = _coerce_reward_value(reward_value)
                        return (
                            result_value
                            if isinstance(result_value, str)
                            else str(result_value)
                        )

    return None


def load_run_metadata(trial_dir: Path) -> Dict[str, Any]:
    """Locate result.json for a trial and extract the required metadata."""
    data = _load_result_data(trial_dir)
    if data is None:
        raise FileNotFoundError(f"No result.json found for trial {trial_dir}")
    return _normalize_run_metadata(data)


def extract_conversations_from_trajectory(
    trajectory_file: Path, run_metadata: Dict[str, Any]
) -> List[Dict[str, Any]]:
    """Extract all episode conversations from a trajectory file.

    Reads the trajectory once and generates one conversation per episode.
    Episodes are determined by agent steps - each agent step marks the end of an episode.

    Note: Steps marked with is_copied_context=True are excluded from training data
    as they represent interactions already present in previous trajectories.

    Args:
        trajectory_file: Path to the trajectory.json file
        run_metadata: Run metadata dictionary

    Returns:
        List of conversation dicts, one per episode

    Raises:
        MultimodalExportError: If the trajectory contains multimodal content (images).
            Text-only export does not support multimodal trajectories to prevent
            silent data loss.
    """
    try:
        trajectory_data = json.loads(trajectory_file.read_text())
    except (json.JSONDecodeError, OSError) as e:
        print(f"[traces] Skipping trajectory {trajectory_file}: invalid JSON ({e})")
        return []

    # FAIL FAST: Check for multimodal content
    if _trajectory_has_multimodal_content(trajectory_data):
        raise MultimodalExportError(
            f"Trajectory '{trajectory_file}' contains multimodal content (images). "
            f"Text-only export is not supported for multimodal trajectories. "
            f"Use a multimodal training pipeline or filter out multimodal trajectories."
        )

    steps = trajectory_data.get("steps", [])

    # Use agent name from trajectory if available (for subagent trajectories),
    # otherwise fall back to run_metadata agent name
    agent_info = trajectory_data.get("agent", {})
    trajectory_agent_name = agent_info.get("name") or run_metadata["agent_name"]
    trajectory_model_name = agent_info.get("model_name") or run_metadata["model_name"]

    # Extract tool definitions if available
    tool_definitions = agent_info.get("tool_definitions")

    # Create a modified run_metadata for this specific trajectory
    trajectory_run_metadata = {
        **run_metadata,
        "agent_name": trajectory_agent_name,
        "model_name": trajectory_model_name,
        "tool_definitions": tool_definitions,
    }

    # Find all agent steps (each marks the end of an episode)
    # Exclude steps marked as copied context
    agent_step_indices = []
    for i, step in enumerate(steps):
        if step.get("source") == "agent":
            # Skip steps marked as copied context (e.g., handoff steps)
            if step.get("is_copied_context"):
                continue
            agent_step_indices.append(i)

    if not agent_step_indices:
        return []

    # Generate one conversation per episode
    conversations = []
    for episode_num, agent_step_idx in enumerate(agent_step_indices):
        conv = _extract_single_episode_conversation(
            steps[
                : agent_step_idx + 1
            ],  # Include all steps up to and including this agent step
            episode_num,
            trajectory_run_metadata,
        )
        if conv and conv.get("conversations"):
            conversations.append(conv)

    return conversations


def _extract_single_episode_conversation(
    steps: List[Dict[str, Any]], episode_num: int, run_metadata: Dict[str, Any]
) -> Optional[Dict[str, Any]]:
    """Extract conversation for a single episode from trajectory steps.

    Episodes end with the assistant's response. Observations from each agent step
    are added as user messages before the next agent turn, except for the last
    agent step in the episode (whose observation belongs to the next episode).

    Steps marked with is_copied_context=True are included in the conversation
    but don't trigger new episode boundaries.

    Args:
        steps: List of trajectory steps for this episode
        episode_num: Episode number (0-indexed)
        run_metadata: Run metadata dictionary

    Returns:
        Conversation dict for this episode
    """
    conv: Dict[str, Any] = {
        "conversations": [],
        "agent": run_metadata["agent_name"],
        "model": run_metadata["model_name"],
        "model_provider": run_metadata["model_provider"],
        "date": run_metadata["start_time"],
        "task": None,  # to be filled by caller
        "episode": f"episode-{episode_num}",
        "run_id": run_metadata["run_id"],
        "trial_name": None,  # to be filled by caller
    }

    # Add tool definitions if available
    if run_metadata.get("tool_definitions"):
        conv["tool_definitions"] = run_metadata["tool_definitions"]

    # Track agent steps to know when to add observations
    # Note: We use ALL steps (including copied context) for building the conversation
    agent_steps = []
    for i, step in enumerate(steps):
        if step.get("source") == "agent":
            agent_steps.append(i)

    for i, step in enumerate(steps):
        source = step.get("source")
        message = step.get("message", "")

        if source == "system":
            # System messages become user role (task instruction from user)
            conv["conversations"].append(
                {
                    "role": "user",
                    "content": message,
                }
            )
        elif source == "user":
            # User messages remain user role
            conv["conversations"].append(
                {
                    "role": "user",
                    "content": message,
                }
            )
        elif source == "agent":
            # Agent messages become assistant role
            # Format: <think>reasoning_content</think> + message + tool_calls JSON
            content_parts = []

            # Prepend reasoning_content wrapped in <think></think> tags if present
            reasoning_content = step.get("reasoning_content")
            if reasoning_content:
                content_parts.append(f"<think>{reasoning_content}</think>")

            # Add the main message content
            if message:
                content_parts.append(message)

            # Append tool_calls in the specified format if present
            # Format: <tool_call>\n{"name": <function-name>, "arguments": <args-json-object>}\n</tool_call>
            tool_calls = step.get("tool_calls")
            if tool_calls:
                for tool_call in tool_calls:
                    tool_call_obj = {
                        "name": tool_call.get("function_name"),
                        "arguments": tool_call.get("arguments", {}),
                    }
                    tool_call_json = json.dumps(tool_call_obj, ensure_ascii=False)
                    content_parts.append(f"<tool_call>\n{tool_call_json}\n</tool_call>")

            assistant_content = "\n".join(content_parts) if content_parts else ""

            conv["conversations"].append(
                {
                    "role": "assistant",
                    "content": assistant_content,
                }
            )

            # Add observation as user message, UNLESS this is the last agent step in the episode
            # (the last agent step's observation belongs to the next episode)
            is_last_agent_step = i == agent_steps[-1]
            if not is_last_agent_step:
                observation = step.get("observation")
                if observation and isinstance(observation, dict):
                    results = observation.get("results", [])
                    if results:
                        # Combine all observation results into a single user message
                        observation_contents = []
                        for result in results:
                            if isinstance(result, dict) and "content" in result:
                                # Use normalize_message_content to handle both string
                                # and list content formats
                                observation_contents.append(
                                    normalize_message_content(result["content"])
                                )

                        if observation_contents:
                            observation_text = "\n".join(observation_contents)
                            conv["conversations"].append(
                                {
                                    "role": "user",
                                    "content": observation_text,
                                }
                            )

    return conv


def _extract_subagent_refs_from_trajectory(
    trajectory_file: Path,
) -> List[str]:
    """Extract subagent trajectory references from a trajectory file.

    Returns:
        List of trajectory filenames referenced in the trajectory (e.g., ["trajectory.summarization-1-summary.json"])
    """
    try:
        trajectory_data = json.loads(trajectory_file.read_text())
    except (json.JSONDecodeError, OSError):
        return []

    refs = []
    steps = trajectory_data.get("steps", [])

    for step in steps:
        observation = step.get("observation", {})
        if not observation:
            continue

        results = observation.get("results", [])
        for result in results:
            if not isinstance(result, dict):
                continue

            subagent_refs = result.get("subagent_trajectory_ref", [])
            for ref in subagent_refs:
                trajectory_path = ref.get("trajectory_path")
                if trajectory_path:
                    refs.append(trajectory_path)

    return refs


def collect_conversations_from_trial(
    trial_dir: Path,
    run_meta: Dict[str, Any],
    episodes: str = "all",
    verbose: bool = False,
    include_instruction: bool = False,
    include_verifier_output: bool = False,
) -> List[Dict[str, Any]]:
    """Collect conversation traces from a trial.

    Supports:
    - Single trajectory files
    - Continuation trajectories (when linear_history mode is enabled)

    Note: This function only collects traces from the main agent trajectory
    and its continuations. For subagent traces, use collect_subagent_traces().

    Args:
        trial_dir: Path to trial directory
        run_meta: Run metadata dictionary
        episodes: "all" to collect all episodes, or "last" to collect only the last completed episode
        verbose: Whether to print verbose output
        include_instruction: Include instruction text metadata when available
        include_verifier_output: Include verifier stdout/stderr metadata when available

    Returns:
        List of conversation dicts
    """
    task_name = run_meta["task_name"]
    trial_name = run_meta["trial_name"]
    result_value = _extract_trial_result_value(trial_dir, trial_name)
    instruction_text = (
        _extract_instruction(trial_dir, run_meta["agent_name"])
        if include_instruction
        else None
    )
    verifier_output = (
        _read_verifier_output(trial_dir) if include_verifier_output else None
    )

    agent_dir = trial_dir / "agent"

    # Build trajectory processing order
    # Only include main trajectory and continuation trajectories
    # Subagent trajectories should be exported separately
    trajectory_order = []

    # Start with main trajectory
    main_traj = agent_dir / "trajectory.json"
    if not main_traj.exists():
        if verbose:
            print(f"[traces] Trial {trial_dir.name}: no trajectory.json found")
        return []

    trajectory_order.append(main_traj)

    # Follow continuation chain using continued_trajectory_ref from each trajectory
    # (Subagent trajectories are NOT added here - they go to separate trace files)
    current_traj_path = main_traj
    while True:
        try:
            trajectory_data = json.loads(current_traj_path.read_text())
        except (json.JSONDecodeError, OSError):
            break

        # Check if this trajectory has a continuation
        continued_ref = trajectory_data.get("continued_trajectory_ref")
        if not continued_ref:
            break

        # Resolve the continuation path relative to the agent directory
        next_traj_path = agent_dir / continued_ref
        if not next_traj_path.exists():
            if verbose:
                print(
                    f"[traces] Warning: continued_trajectory_ref '{continued_ref}' in {current_traj_path.name} points to non-existent file"
                )
            break

        trajectory_order.append(next_traj_path)
        current_traj_path = next_traj_path

    # Extract conversations from all trajectory files in order
    all_conversations = []
    episode_offset = 0

    for traj_file in trajectory_order:
        conversations = extract_conversations_from_trajectory(traj_file, run_meta)

        # Adjust episode numbers to be globally unique across all trajectory files
        for conv in conversations:
            episode_num = int(conv["episode"].split("-")[1])
            conv["episode"] = f"episode-{episode_offset + episode_num}"

        all_conversations.extend(conversations)
        episode_offset += len(conversations)

    if not all_conversations:
        if verbose:
            print(
                f"[traces] Trial {trial_dir.name}: no conversations found in trajectories"
            )
        return []

    # Handle "last" episode filter
    if episodes == "last":
        all_conversations = [all_conversations[-1]]
        if verbose:
            print(
                f"[traces] Trial {trial_dir.name}: selected last episode {all_conversations[0]['episode']}"
            )

    # Fill in task and trial_name for all conversations
    for conv in all_conversations:
        conv["task"] = task_name
        conv["trial_name"] = trial_name
        conv["result"] = result_value
        if include_instruction:
            conv["instruction"] = instruction_text
        if include_verifier_output:
            conv["verifier_output"] = verifier_output

    if verbose:
        traj_count = len(trajectory_order)
        traj_suffix = f" ({traj_count} trajectory files)" if traj_count > 1 else ""
        print(
            f"[traces] Collected {len(all_conversations)} rows from trial {trial_dir.name}{traj_suffix}"
        )

    return all_conversations


def collect_subagent_traces(
    trial_dir: Path,
    run_meta: Dict[str, Any],
    episodes: str = "all",
    verbose: bool = False,
    include_instruction: bool = False,
    include_verifier_output: bool = False,
) -> Dict[str, List[Dict[str, Any]]]:
    """Collect traces from subagent trajectories (e.g., context summarization agents).

    Returns a dictionary mapping subagent trajectory types to their trace lists.
    The traces use the main agent name (not the subagent name) for consistency.

    For example:
    {
        "summarization-1-summary": [...],  # All traces have agent="terminus-2"
        "summarization-1-questions": [...],  # All traces have agent="terminus-2"
        "summarization-1-answers": [...]  # All traces have agent="terminus-2"
    }

    Args:
        trial_dir: Path to trial directory
        run_meta: Run metadata dictionary
        episodes: "all" to collect all episodes, or "last" to collect only the last completed episode
        verbose: Whether to print verbose output
        include_instruction: Include instruction text metadata when available
        include_verifier_output: Include verifier stdout/stderr metadata when available

    Returns:
        Dictionary mapping subagent trajectory types to lists of conversation dicts
    """
    agent_dir = trial_dir / "agent"
    subagent_traces = {}
    result_value = _extract_trial_result_value(trial_dir, run_meta["trial_name"])
    instruction_text = (
        _extract_instruction(trial_dir, run_meta["agent_name"])
        if include_instruction
        else None
    )
    verifier_output = (
        _read_verifier_output(trial_dir) if include_verifier_output else None
    )

    # Find all subagent trajectory files
    # They follow the pattern: trajectory.*.json (but not trajectory.cont-*.json)
    subagent_files = []
    for traj_file in sorted(agent_dir.glob("trajectory.*.json")):
        # Skip continuation files
        if ".cont-" in traj_file.name:
            continue
        subagent_files.append(traj_file)

    if not subagent_files:
        return subagent_traces

    # Extract conversations from each subagent trajectory
    for traj_file in subagent_files:
        conversations = extract_conversations_from_trajectory(traj_file, run_meta)

        if not conversations:
            continue

        # Extract the subagent type from the filename (e.g., "summarization-1-summary")
        # Filename format: trajectory.summarization-1-summary.json
        filename = traj_file.name
        if filename.startswith("trajectory.") and filename.endswith(".json"):
            subagent_type = filename[len("trajectory.") : -len(".json")]
        else:
            continue  # Skip files that don't match the expected pattern

        # Override agent name to use the main agent name (not the subagent name)
        # This is important because subagent names are implementation details
        main_agent_name = run_meta["agent_name"]
        for conv in conversations:
            conv["agent"] = main_agent_name

        # Handle "last" episode filter
        if episodes == "last" and conversations:
            conversations = [conversations[-1]]

        # Fill in task and trial_name for all conversations
        task_name = run_meta["task_name"]
        trial_name = run_meta["trial_name"]
        for conv in conversations:
            conv["task"] = task_name
            conv["trial_name"] = trial_name
            conv["result"] = result_value
            if include_instruction:
                conv["instruction"] = instruction_text
            if include_verifier_output:
                conv["verifier_output"] = verifier_output

        subagent_traces[subagent_type] = conversations

        if verbose:
            print(
                f"[traces] Collected {len(conversations)} rows from subagent trajectory {subagent_type} in trial {trial_dir.name}"
            )

    return subagent_traces


# --------------------
# Dataset helpers
# --------------------


def rows_to_dataset(rows: List[Dict[str, Any]]) -> "Dataset":
    if Dataset is None:  # pragma: no cover - import-time optionality
        raise RuntimeError("datasets is not installed")
    return Dataset.from_list(rows)


def push_dataset(dataset: "Dataset", repo_id: str, token: Optional[str] = None) -> None:
    token = token or os.getenv("HUGGINGFACE_TOKEN") or os.getenv("HF_TOKEN")
    if not token:
        raise RuntimeError("No HF token found in env (HUGGINGFACE_TOKEN/HF_TOKEN)")
    dataset.push_to_hub(repo_id, token=token)


# --------------------
# High-level entry
# --------------------


def export_traces(
    root: Path | str,
    recursive: bool = True,
    episodes: str = "all",
    to_sharegpt: bool = False,
    repo_id: Optional[str] = None,
    push: bool = False,
    verbose: bool = False,
    success_filter: Optional[str] = None,
    export_subagents: bool = True,
    include_instruction: bool = False,
    include_verifier_output: bool = False,
) -> "Dataset | Dict[str, Dataset]":
    """Export traces under root into a HF Dataset. If push=True and repo_id is set, upload.

    Args:
        root: Root directory containing trial directories
        recursive: Whether to search recursively for trial directories
        episodes: "all" to collect all episodes, or "last" to collect only the last completed episode
        to_sharegpt: Whether to convert to ShareGPT format
        repo_id: HuggingFace repo ID for pushing datasets
        push: Whether to push to HuggingFace Hub
        verbose: Whether to print verbose output
        success_filter: Optional filter for successful/failed trials ("success", "failure", or None)
        export_subagents: Whether to export subagent traces (default: True)
        include_instruction: Include instruction text metadata column
        include_verifier_output: Include verifier stdout/stderr metadata column

    Returns:
        If export_subagents=False: A single Dataset with main agent traces
        If export_subagents=True: A dictionary with:
            - "main": Dataset with main agent traces
            - "<subagent-type>": Dataset for each subagent trajectory type (e.g., "summarization-1-summary")
              Note: All traces use the main agent name, not subagent-specific names
    """
    root = Path(root)
    rows: List[Dict[str, Any]] = []
    subagent_rows: Dict[str, List[Dict[str, Any]]] = {}

    trial_dirs = list(iter_trial_dirs(root, recursive=recursive))
    print(f"[traces] Found {len(trial_dirs)} trial directories under {root}")

    for trial_dir in trial_dirs:
        run_meta = load_run_metadata(trial_dir)
        agent_name = run_meta["agent_name"]

        # Check if agent supports ATIF trajectory format
        agent_enum = AgentName(agent_name)
        agent_class = AgentFactory._AGENT_MAP.get(agent_enum)
        if agent_class is None or not agent_class.SUPPORTS_ATIF:
            raise NotImplementedError(
                f"{agent_name} does not support Harbor's trajectory format (ATIF), cannot export traces"
            )

        # Optional trial-level success/failure filter based on result.json
        if success_filter in ("success", "failure"):
            succ = _trial_is_success(trial_dir)
            if succ is None:
                if verbose:
                    print(
                        f"[traces] Trial {trial_dir.name}: missing result.json; skipping due to filter"
                    )
                continue
            if success_filter == "success" and not succ:
                continue
            if success_filter == "failure" and succ:
                continue

        # Collect main agent traces
        rows.extend(
            collect_conversations_from_trial(
                trial_dir,
                episodes=episodes,
                verbose=verbose,
                run_meta=run_meta,
                include_instruction=include_instruction,
                include_verifier_output=include_verifier_output,
            )
        )

        # Collect subagent traces if requested
        if export_subagents:
            trial_subagent_traces = collect_subagent_traces(
                trial_dir,
                run_meta=run_meta,
                episodes=episodes,
                verbose=verbose,
                include_instruction=include_instruction,
                include_verifier_output=include_verifier_output,
            )

            # Merge subagent traces from this trial with accumulated subagent traces
            for subagent_name, subagent_convs in trial_subagent_traces.items():
                if subagent_name not in subagent_rows:
                    subagent_rows[subagent_name] = []
                subagent_rows[subagent_name].extend(subagent_convs)

    # Create main dataset
    main_ds = rows_to_dataset(rows)
    if to_sharegpt:
        main_ds = convert_openai_to_sharegpt(
            main_ds, "conversations", "conversations_sharegpt"
        )

    if verbose:
        print(
            f"[traces] Prepared {len(rows)} main agent rows; to_sharegpt={to_sharegpt}, push={push}, repo_id={repo_id}"
        )

    # If no subagents or export_subagents=False, return just the main dataset
    if not export_subagents or not subagent_rows:
        if push and repo_id:
            push_dataset(main_ds, repo_id)
        return main_ds

    # Create subagent datasets
    result = {"main": main_ds}
    for subagent_type, subagent_trace_list in subagent_rows.items():
        subagent_ds = rows_to_dataset(subagent_trace_list)
        if to_sharegpt:
            subagent_ds = convert_openai_to_sharegpt(
                subagent_ds, "conversations", "conversations_sharegpt"
            )
        result[subagent_type] = subagent_ds

        if verbose:
            print(
                f"[traces] Prepared {len(subagent_trace_list)} rows for subagent trajectory {subagent_type}"
            )

        # Push subagent dataset if requested
        if push and repo_id:
            subagent_repo_id = f"{repo_id}-{subagent_type}"
            push_dataset(subagent_ds, subagent_repo_id)

    # Push main dataset if requested
    if push and repo_id:
        push_dataset(main_ds, repo_id)

    return result


def _trial_is_success(
    trial_dir: Path, run_meta: Dict[str, Any] | None = None
) -> Optional[bool]:
    """Determine success using job-level reward stats when available."""
    trial_name = run_meta["trial_name"] if run_meta else trial_dir.name
    result_value = _extract_trial_result_value(trial_dir, trial_name)
    if result_value is not None:
        try:
            return float(result_value) > 0.0
        except (ValueError, TypeError):
            return False

    data = _load_result_data(trial_dir)
    if not isinstance(data, dict):
        return None

    vr = data.get("verifier_result")
    reward = vr.get("reward") if isinstance(vr, dict) else None
    if reward is None:
        return False
    try:
        return float(reward) > 0.0
    except (ValueError, TypeError):
        return False
