"""Tests for trajectory parsing and summarization.

Covers:
    - format_step() for various tool call types (Bash, Edit, Read, Write)
    - extract_trajectory_summary() with missing files, empty trajectories, truncation
    - parse_raw_trajectory_steps() for dashboard export
"""

from __future__ import annotations

import json
from pathlib import Path

from clawbench.trajectory import (
    format_step,
    extract_trajectory_summary,
    parse_raw_trajectory_steps,
)


class TestFormatStep:
    def test_bash_step(self):
        """Bash tool call formats correctly."""
        step = {
            "source": "agent",
            "tool_calls": [{
                "function_name": "Bash",
                "arguments": {"command": "ls -la", "description": "List files"},
            }],
        }
        lines = format_step(step)
        assert len(lines) == 1
        assert "[Bash (List files)]: ls -la" in lines[0]

    def test_bash_no_description(self):
        """Bash tool call without description."""
        step = {
            "source": "agent",
            "tool_calls": [{
                "function_name": "Bash",
                "arguments": {"command": "pwd"},
            }],
        }
        lines = format_step(step)
        assert "[Bash]: pwd" in lines[0]

    def test_edit_step_truncates(self):
        """Edit tool call truncates old/new strings."""
        step = {
            "source": "agent",
            "tool_calls": [{
                "function_name": "Edit",
                "arguments": {
                    "file_path": "/app/main.py",
                    "old_string": "x" * 200,
                    "new_string": "y" * 200,
                },
            }],
        }
        lines = format_step(step)
        assert len(lines) == 1
        assert "[Edit /app/main.py]" in lines[0]
        # Truncated to 100 chars
        assert "x" * 100 in lines[0]
        assert "x" * 101 not in lines[0]

    def test_read_step(self):
        """Read tool call formats correctly."""
        step = {
            "source": "agent",
            "tool_calls": [{
                "function_name": "Read",
                "arguments": {"file_path": "/app/config.json"},
            }],
        }
        lines = format_step(step)
        assert "[Read]: /app/config.json" in lines[0]

    def test_write_step(self):
        """Write tool call shows content preview."""
        step = {
            "source": "agent",
            "tool_calls": [{
                "function_name": "Write",
                "arguments": {
                    "file_path": "/app/output.txt",
                    "content": "Hello world",
                },
            }],
        }
        lines = format_step(step)
        assert "[Write /app/output.txt]" in lines[0]
        assert "Hello world" in lines[0]

    def test_agent_thought(self):
        """Agent thought message (no tool calls)."""
        step = {
            "source": "agent",
            "message": "I need to check the config file first.",
            "tool_calls": [],
        }
        lines = format_step(step)
        assert "[Agent thought]" in lines[0]
        assert "config file" in lines[0]

    def test_agent_thought_truncated(self):
        """Long agent thoughts are truncated."""
        step = {
            "source": "agent",
            "message": "x" * 500,
            "tool_calls": [],
        }
        lines = format_step(step)
        assert lines[0].endswith("...")
        assert len(lines[0]) < 350

    def test_system_prompt_skipped(self):
        """Step 1 (system prompt) is summarized."""
        step = {
            "source": "user",
            "step_id": 1,
            "message": "Very long system prompt...",
        }
        lines = format_step(step)
        assert lines == ["[Step 1: Task instruction provided to agent]"]

    def test_observation_output(self):
        """Observation results are included."""
        step = {
            "source": "agent",
            "tool_calls": [{
                "function_name": "Bash",
                "arguments": {"command": "echo hello"},
            }],
            "observation": {
                "results": [{"content": "hello\n"}],
            },
        }
        lines = format_step(step)
        assert len(lines) == 2
        assert "→ hello" in lines[1]

    def test_observation_error_gets_more_space(self):
        """Important observations (errors) get more characters."""
        error_content = "ERROR: " + "x" * 300
        step = {
            "source": "agent",
            "tool_calls": [{
                "function_name": "Bash",
                "arguments": {"command": "make build"},
            }],
            "observation": {
                "results": [{"content": error_content}],
            },
        }
        lines = format_step(step)
        obs_line = lines[1]
        # Error content should get up to 500 chars
        assert len(obs_line) > 200

    def test_unknown_tool(self):
        """Unknown tool names are handled gracefully."""
        step = {
            "source": "agent",
            "tool_calls": [{
                "function_name": "CustomTool",
                "arguments": {"foo": "bar"},
            }],
        }
        lines = format_step(step)
        assert "[CustomTool]" in lines[0]


class TestExtractTrajectorySummary:
    def test_missing_file(self):
        """Returns placeholder when trajectory file doesn't exist."""
        result = extract_trajectory_summary(Path("/nonexistent/trial"))
        assert result == "(no trajectory file found)"

    def test_empty_trajectory(self, tmp_path):
        """Returns placeholder for empty steps."""
        traj_dir = tmp_path / "agent"
        traj_dir.mkdir(parents=True)
        (traj_dir / "trajectory.json").write_text(json.dumps({"steps": []}))
        result = extract_trajectory_summary(tmp_path)
        assert result == "(empty trajectory)"

    def test_invalid_json(self, tmp_path):
        """Returns placeholder for invalid JSON."""
        traj_dir = tmp_path / "agent"
        traj_dir.mkdir(parents=True)
        (traj_dir / "trajectory.json").write_text("not json")
        result = extract_trajectory_summary(tmp_path)
        assert result == "(failed to read trajectory)"

    def test_short_trajectory_fits(self, tmp_path):
        """Short trajectories are included in full."""
        traj_dir = tmp_path / "agent"
        traj_dir.mkdir(parents=True)
        traj = {
            "steps": [
                {"source": "agent", "message": "thinking", "tool_calls": []},
                {"source": "agent", "tool_calls": [{"function_name": "Bash", "arguments": {"command": "ls"}}]},
            ],
            "final_metrics": {"total_prompt_tokens": 100, "total_completion_tokens": 50},
        }
        (traj_dir / "trajectory.json").write_text(json.dumps(traj))
        result = extract_trajectory_summary(tmp_path)
        assert "[Agent thought]" in result
        assert "[Bash]: ls" in result
        assert "[Summary: 2 steps" in result

    def test_truncation_head_tail(self, tmp_path):
        """Long trajectories use head+tail strategy."""
        traj_dir = tmp_path / "agent"
        traj_dir.mkdir(parents=True)
        # Create many steps that exceed max_chars
        steps = []
        for i in range(50):
            steps.append({
                "source": "agent",
                "tool_calls": [{
                    "function_name": "Bash",
                    "arguments": {"command": f"command_{i} " + "x" * 100},
                }],
            })
        traj = {"steps": steps, "final_metrics": {}}
        (traj_dir / "trajectory.json").write_text(json.dumps(traj))

        result = extract_trajectory_summary(tmp_path, max_chars=2000)
        assert "steps omitted" in result
        # First steps present
        assert "command_0" in result
        # Last steps present
        assert f"command_{49}" in result or "command_4" in result


class TestParseRawTrajectorySteps:
    def test_missing_file(self, tmp_path):
        """Returns empty list when file doesn't exist."""
        result = parse_raw_trajectory_steps(tmp_path)
        assert result == []

    def test_thought_step(self, tmp_path):
        """Agent thought becomes thought entry."""
        traj_dir = tmp_path / "agent"
        traj_dir.mkdir(parents=True)
        traj = {
            "steps": [{
                "source": "agent",
                "message": "I should check the config",
                "tool_calls": [],
                "observation": {},
            }],
        }
        (traj_dir / "trajectory.json").write_text(json.dumps(traj))
        result = parse_raw_trajectory_steps(tmp_path)
        assert len(result) == 1
        assert result[0]["type"] == "thought"
        assert "config" in result[0]["content"]

    def test_tool_step_with_observation(self, tmp_path):
        """Tool call + observation produces two entries."""
        traj_dir = tmp_path / "agent"
        traj_dir.mkdir(parents=True)
        traj = {
            "steps": [{
                "source": "agent",
                "tool_calls": [{
                    "function_name": "Bash",
                    "arguments": {"command": "ls", "description": "list files"},
                }],
                "observation": {
                    "results": [{"content": "file1.txt\nfile2.txt"}],
                },
            }],
        }
        (traj_dir / "trajectory.json").write_text(json.dumps(traj))
        result = parse_raw_trajectory_steps(tmp_path)
        assert len(result) == 2
        assert result[0]["type"] == "tool"
        assert result[0]["tool"] == "Bash"
        assert result[0]["command"] == "ls"
        assert result[1]["type"] == "observation"
        assert "file1.txt" in result[1]["content"]
