utils/trajectory_parser.py

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

import zstandard as zstd


def list_eval_files(logs_dir: str) -> List[Dict[str, Any]]:
    """List all .eval files in logs directory, sorted by date (newest first)."""
    eval_files = []
    logs_path = Path(logs_dir)

    if not logs_path.exists():
        return eval_files

    for file in logs_path.glob("*.eval"):
        try:
            name_parts = file.stem.split("_")
            if len(name_parts) < 2:
                continue
            timestamp_str = name_parts[0]
            task_type = "_".join(name_parts[1:-1])

            # Parse timestamp format: 2026-06-08T17-54-56-00-00
            # Convert to ISO format by replacing hyphens in time part with colons
            # Format is: YYYY-MM-DDTHH-MM-SS-TZ-TZ -> YYYY-MM-DDTHH:MM:SS+TZ:TZ
            parts = timestamp_str.split("T")
            if len(parts) == 2:
                date_part = parts[0]  # 2026-06-08
                time_part = parts[1]  # 17-54-56-00-00
                time_components = time_part.split("-")
                if len(time_components) >= 3:
                    # Reconstruct as HH:MM:SS+00:00 or similar
                    iso_str = f"{date_part}T{time_components[0]}:{time_components[1]}:{time_components[2]}"
                    if len(time_components) > 3:
                        # Add timezone if present
                        iso_str += (
                            f"+{time_components[3]}:{time_components[4]}"
                            if len(time_components) > 4
                            else "+00:00"
                        )
                    dt = datetime.fromisoformat(iso_str)
                else:
                    continue
            else:
                continue

            eval_files.append(
                {
                    "timestamp": timestamp_str,
                    "filename": file.name,
                    "task_type": task_type,
                    "datetime_obj": dt,
                    "full_path": str(file),
                }
            )
        except Exception:
            continue

    eval_files.sort(key=lambda x: x["datetime_obj"], reverse=True)
    return eval_files


def _extract_zstd_from_zip(zip_ref, filename: str) -> Optional[bytes]:
    """Extract a single file from zip with Zstandard decompression."""
    try:
        info = zip_ref.getinfo(filename)
        if info.compress_type != 93:  # Not zstandard, use normal extraction
            return zip_ref.read(filename)

        # Read raw compressed data
        zip_ref.fp.seek(info.header_offset)
        local_header = zip_ref.fp.read(30)
        fn_size = int.from_bytes(local_header[26:28], "little")
        extra_size = int.from_bytes(local_header[28:30], "little")
        zip_ref.fp.seek(info.header_offset + 30 + fn_size + extra_size)
        compressed_data = zip_ref.fp.read(info.compress_size)

        # Decompress using zstandard
        dctx = zstd.ZstdDecompressor()
        return dctx.decompress(compressed_data, max_output_size=info.file_size)
    except (KeyError, zstd.ZstdError):
        return None


def extract_eval_zip(eval_path: str) -> Optional[Dict[str, Any]]:
    """Extract and parse .eval zip file (Zstd-compressed)."""
    try:
        with zipfile.ZipFile(eval_path, "r") as zip_ref:
            # Extract header
            header_data = None
            try:
                header_json = _extract_zstd_from_zip(zip_ref, "header.json")
                if header_json:
                    header_data = json.loads(header_json.decode("utf-8"))
            except (KeyError, json.JSONDecodeError):
                pass

            # Extract first sample (usually only one)
            samples_data = None
            try:
                sample_files = [
                    f for f in zip_ref.namelist() if f.startswith("samples/")
                ]
                if sample_files:
                    sample_json = _extract_zstd_from_zip(zip_ref, sample_files[0])
                    if sample_json:
                        samples_data = json.loads(sample_json.decode("utf-8"))
            except (KeyError, json.JSONDecodeError):
                pass

            # Extract journal start
            journal_start = None
            try:
                journal_json = _extract_zstd_from_zip(zip_ref, "_journal/start.json")
                if journal_json:
                    journal_start = json.loads(journal_json.decode("utf-8"))
            except (KeyError, json.JSONDecodeError):
                pass

            return {
                "header": header_data,
                "samples": samples_data,
                "journal_start": journal_start,
            }
    except Exception:
        return None


def parse_trajectory(sample_json: Dict[str, Any]) -> Dict[str, Any]:
    """Parse sample JSON preserving execution order of events.

    Returns a timeline of events in chronological order instead of grouped by type.
    """
    timeline = []  # Events in execution order
    final_result = None

    # Extract from events (Inspect AI structure) - preserve order
    events = sample_json.get("events", [])
    for event in events:
        event_type = event.get("event")

        # Collect tool calls (including "think" tool)
        if event_type == "tool":
            function = event.get("function", "unknown")
            args = event.get("arguments", {})
            result = event.get("result", None)
            timeline.append(
                {
                    "type": "tool",
                    "name": function,
                    "arguments": args,
                    "result": result,
                    "timestamp": event.get("timestamp"),
                }
            )

        # Collect logger messages as insights
        if event_type == "logger":
            msg = event.get("message", {})
            if isinstance(msg, dict) and "message" in msg:
                timeline.append(
                    {
                        "type": "log",
                        "content": msg.get("message", ""),
                        "timestamp": event.get("timestamp"),
                    }
                )

    # Extract from messages if present (fallback)
    if "messages" in sample_json and not timeline:
        for msg in sample_json["messages"]:
            if "content" in msg:
                content = msg.get("content", "")
                if isinstance(content, str) and content.strip():
                    timeline.append(
                        {
                            "type": "message",
                            "role": msg.get("role", "unknown"),
                            "content": content,
                        }
                    )

            if "tool_calls" in msg:
                for tool_call in msg["tool_calls"]:
                    timeline.append(
                        {
                            "type": "tool",
                            "name": tool_call.get("name", "unknown"),
                            "arguments": tool_call.get("arguments", {}),
                            "result": tool_call.get("result", None),
                        }
                    )

    # Extract final output
    output = sample_json.get("output", {})
    if output:
        completion = output.get("completion", "")
        if completion:
            final_result = completion
        else:
            choices = output.get("choices", [])
            if choices and len(choices) > 0:
                msg = choices[0].get("message", {})
                if isinstance(msg, dict):
                    final_result = msg

    # If no timeline yet, create summary
    if not timeline and (events or output):
        timeline.append(
            {
                "type": "log",
                "content": f"Evaluation run completed. Processed {len(events)} events.",
            }
        )

    return {
        "timeline": timeline,  # Events in execution order
        "thinking_steps": [
            e for e in timeline if e["type"] == "log"
        ],  # For backward compat
        "tool_calls": [
            e for e in timeline if e["type"] == "tool"
        ],  # For backward compat
        "final_result": final_result,
        "raw": sample_json,
    }


def format_trajectory_for_display(
    trajectory: Dict[str, Any], show_thinking: bool = True, show_tools: bool = True
) -> str:
    """Format trajectory data for human-readable display."""
    output_parts = []

    if show_thinking and trajectory["thinking_steps"]:
        output_parts.append("## THINKING STEPS\n")
        for i, step in enumerate(trajectory["thinking_steps"], 1):
            role = step.get("role", "unknown").upper()
            content = step.get("content", "")
            output_parts.append(f"**Step {i} ({role}):**\n{content}\n")

    if show_tools and trajectory["tool_calls"]:
        output_parts.append("\n## TOOL CALLS\n")
        for i, call in enumerate(trajectory["tool_calls"], 1):
            name = call.get("name", "unknown")
            args = call.get("arguments", {})
            result = call.get("result", None)
            output_parts.append(f"**Call {i}: {name}**\n")
            output_parts.append(f"Arguments: {json.dumps(args, indent=2)}\n")
            if result:
                output_parts.append(f"Result: {json.dumps(result, indent=2)}\n")

    if trajectory["final_result"]:
        output_parts.append("\n## FINAL RESULT\n")
        if isinstance(trajectory["final_result"], str):
            output_parts.append(trajectory["final_result"])
        else:
            output_parts.append(json.dumps(trajectory["final_result"], indent=2))

    return "".join(output_parts)