from __future__ import annotations

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

"""
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 episode input messages (debug.json) and the assistant
          output (response.json or response.txt). Content is normalized into
          text via best‑effort rules (see normalize_message_content).
        - Roles typically include "system", "user", and a final "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 directory name (e.g., "episode-0001")
    - run_id: str               # job/run identifier
    - trial_name: Optional[str] # trial name associated with the run

Notes and options:
    - Only trials that contain agent/episode-* directories are discovered.
    - 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


# --------------------
# 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


def find_episode_dirs(trial_dir: Path) -> List[Path]:
    episodes_root = trial_dir / "agent"
    if not episodes_root.exists():
        return []
    eps = [
        d
        for d in episodes_root.iterdir()
        if d.is_dir() and d.name.startswith("episode-")
    ]
    return sorted(eps, key=lambda x: x.name)


# --------------------
# 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"]
    agent_cfg = config["agent"]
    agent_info = raw["agent_info"]
    model_info = agent_info["model_info"]
    return {
        "agent_name": agent_cfg["name"],
        "model_name": model_info["name"],
        "model_provider": model_info["provider"],
        "start_time": raw["started_at"],
        "run_id": config["job_id"],
        "task_name": raw["task_name"],
        "trial_name": raw["trial_name"],
        "raw_metadata": raw,
    }


def load_run_metadata(trial_dir: Path) -> Dict[str, Any]:
    """Locate result.json for a trial and extract the required metadata."""
    cur = trial_dir
    for _ in range(4):
        meta = cur / "result.json"
        if meta.exists():
            data = json.loads(meta.read_text())
            return _normalize_run_metadata(data)
        if cur.parent == cur:
            break
        cur = cur.parent
    raise FileNotFoundError(f"No result.json found for trial {trial_dir}")


def extract_conversation_from_episode(
    episode_dir: Path, run_metadata: Dict[str, Any]
) -> Optional[Dict[str, Any]]:
    debug_file = episode_dir / "debug.json"
    response_json = episode_dir / "response.json"
    response_txt = episode_dir / "response.txt"
    if not debug_file.exists() or (
        not response_json.exists() and not response_txt.exists()
    ):
        return None

    try:
        debug_data = json.loads(debug_file.read_text())
        response_data = (
            json.loads(response_json.read_text()) if response_json.exists() else None
        )
    except (json.JSONDecodeError, OSError) as e:
        print(f"[traces] Skipping episode {episode_dir}: invalid JSON ({e})")
        return None

    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": episode_dir.name,
        "run_id": run_metadata["run_id"],
        "trial_name": None,  # to be filled by caller
    }

    # Inputs
    msgs = debug_data.get("input") if isinstance(debug_data, dict) else None
    if isinstance(msgs, list):
        for m in msgs:
            if isinstance(m, dict) and "role" in m:
                conv["conversations"].append(
                    {
                        "role": m.get("role"),
                        "content": normalize_message_content(m.get("content")),
                    }
                )

    # Assistant output: stringify response unless there's a top-level textual content field
    assistant_text = None
    if isinstance(response_data, dict):
        # common fields where model text may live
        for k in ("content", "text", "message", "output"):
            v = response_data.get(k)
            if isinstance(v, (str, list, dict)):
                assistant_text = normalize_message_content(v)
                if assistant_text:
                    break
    elif response_txt.exists():
        try:
            assistant_text = response_txt.read_text()
        except OSError as e:
            print(f"[traces] Failed to read response.txt in {episode_dir}: {e}")
            assistant_text = None
    if not assistant_text:
        try:
            assistant_text = json.dumps(response_data, ensure_ascii=False)
        except TypeError:
            assistant_text = str(response_data)

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


def collect_conversations_from_trial(
    trial_dir: Path,
    run_meta: Dict[str, Any],
    episodes: str = "all",
    verbose: bool = False,
) -> List[Dict[str, Any]]:
    task_name = run_meta["task_name"]
    trial_name = run_meta["trial_name"]
    eps = find_episode_dirs(trial_dir)
    if episodes == "last" and eps:
        eps = [eps[-1]]

    rows: List[Dict[str, Any]] = []
    for ep in eps:
        conv = extract_conversation_from_episode(ep, run_meta)
        if not conv or not conv.get("conversations"):
            if verbose:
                dbg = []
                if not (ep / "debug.json").exists():
                    dbg.append("missing debug.json")
                if not (ep / "response.json").exists():
                    dbg.append("missing response.json")
                extra = f" ({', '.join(dbg)})" if dbg else ""
                print(f"[traces] Skipping {ep} — no conversations{extra}")
            continue
        conv["task"] = task_name
        conv["trial_name"] = trial_name
        rows.append(conv)
    if verbose:
        print(f"[traces] Collected {len(rows)} rows from trial {trial_dir.name}")
    return rows


# --------------------
# 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,
) -> "Dataset":
    """Export traces under root into a HF Dataset. If push=True and repo_id is set, upload."""
    root = Path(root)
    rows: 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:
        # TODO(harbor#traces): support additional agent types for export
        run_meta = load_run_metadata(trial_dir)
        agent_name = run_meta["agent_name"]
        if agent_name != "terminus-2":
            raise NotImplementedError(
                "trace export currently supports only the default terminus-2 agent"
            )
        # 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
        ep_dirs = find_episode_dirs(trial_dir)
        if verbose:
            print(f"[traces] Trial {trial_dir.name}: {len(ep_dirs)} episode dirs")
        rows.extend(
            collect_conversations_from_trial(
                trial_dir,
                episodes=episodes,
                verbose=verbose,
                run_meta=run_meta,
            )
        )

    ds = rows_to_dataset(rows)
    if to_sharegpt:
        ds = convert_openai_to_sharegpt(ds, "conversations", "conversations_sharegpt")
    if verbose:
        print(
            f"[traces] Prepared {len(rows)} rows; to_sharegpt={to_sharegpt}, push={push}, repo_id={repo_id}"
        )
    if push and repo_id:
        push_dataset(ds, repo_id)
    return ds


def _trial_is_success(
    trial_dir: Path, run_meta: Dict[str, Any] | None = None
) -> Optional[bool]:
    """Read trial_dir/result.json and return True if reward > 0, False if present and not, None if missing."""
    result_path = trial_dir / "result.json"
    if not result_path.exists():
        # Some layouts keep result.json at the job root; we strictly require per-trial for filtering
        return None
    try:
        data = json.loads(result_path.read_text())
        vr = data.get("verifier_result") if isinstance(data, dict) else None
        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
    except (json.JSONDecodeError, OSError) as e:
        print(f"[traces] Invalid or unreadable result.json in {trial_dir}: {e}")
        return None
