"""Trajectory parsing, summarization, and LLM diagnosis.

Called by: executor.py, backend.py
Calls: litellm (for LLM summarization), _env (for model config)

Pure functions that read trajectory JSON files — no backend dependency.
Any backend that produces trial_dir/agent/trajectory.json gets summarization for free.
"""

from __future__ import annotations

import json
from pathlib import Path


def format_step(step: dict) -> list[str]:
    """Format a single trajectory step into human-readable lines."""
    lines: list[str] = []
    source = step.get("source", "")
    message = step.get("message", "")
    tool_calls = step.get("tool_calls", []) or []
    observation = step.get("observation") or {}
    obs_results = observation.get("results", []) or []

    # Skip the system prompt / skill injection (step 1, usually huge)
    if source == "user" and step.get("step_id") == 1:
        lines.append("[Step 1: Task instruction provided to agent]")
        return lines

    # Agent messages (reasoning about what to do)
    if source == "agent" and message and not tool_calls:
        msg_trunc = message[:300] + ("..." if len(message) > 300 else "")
        lines.append(f"[Agent thought]: {msg_trunc}")

    # Tool calls
    for tc in tool_calls:
        fn_name = tc.get("function_name", "unknown")
        args = tc.get("arguments", {})

        if fn_name == "Bash":
            cmd = args.get("command", "")
            desc = args.get("description", "")
            entry = f"[Bash]: {cmd}"
            if desc:
                entry = f"[Bash ({desc})]: {cmd}"
        elif fn_name == "Edit":
            fp = args.get("file_path", "?")
            old = (args.get("old_string", "") or "")[:100]
            new = (args.get("new_string", "") or "")[:100]
            entry = f"[Edit {fp}]: '{old}' → '{new}'"
        elif fn_name == "Write":
            fp = args.get("file_path", "?")
            content_preview = (args.get("content", "") or "")[:200]
            entry = f"[Write {fp}]: {content_preview}..."
        elif fn_name == "Read":
            fp = args.get("file_path", "?")
            entry = f"[Read]: {fp}"
        else:
            entry = f"[{fn_name}]: {json.dumps(args)[:200]}"
        lines.append(entry)

    # Observation/output from tool calls
    for obs in obs_results:
        content = obs.get("content", "")
        if not content or not isinstance(content, str):
            continue
        is_important = any(kw in content.lower() for kw in [
            "error", "exception", "failed", "404", "401", "500",
            "curl", "http", "gmail", "label", "message", "thread",
        ])
        max_obs = 500 if is_important else 150
        obs_trunc = content[:max_obs] + ("..." if len(content) > max_obs else "")
        lines.append(f"  → {obs_trunc}")

    return lines


def extract_trajectory_summary(trial_dir: Path, max_chars: int = 8000) -> str:
    """Extract a concise action summary from a trial's trajectory JSON.

    Keeps the first few steps (approach context) and last steps (where
    failures happen), dropping the middle when the trajectory is too long.

    Returns a human-readable summary the proposer can learn from.
    """
    traj_path = trial_dir / "agent" / "trajectory.json"
    if not traj_path.exists():
        return "(no trajectory file found)"

    try:
        traj = json.loads(traj_path.read_text())
    except (json.JSONDecodeError, OSError):
        return "(failed to read trajectory)"

    steps = traj.get("steps", [])
    if not steps:
        return "(empty trajectory)"

    # Format every step, then decide what to keep
    formatted: list[list[str]] = [format_step(s) for s in steps]
    step_texts = ["\n".join(lines) for lines in formatted if lines]
    total_chars = sum(len(t) for t in step_texts)

    # Metrics footer
    metrics = traj.get("final_metrics", {})
    footer = (
        f"\n[Summary: {len(steps)} steps, "
        f"{metrics.get('total_prompt_tokens', 0)} prompt tokens, "
        f"{metrics.get('total_completion_tokens', 0)} completion tokens]"
    )

    if total_chars + len(footer) <= max_chars:
        # Everything fits
        return "\n".join(step_texts) + footer

    # Head+tail strategy: keep first 3 steps + as many tail steps as fit
    head_count = min(3, len(step_texts))
    head = step_texts[:head_count]
    head_chars = sum(len(t) for t in head)

    remaining = max_chars - head_chars - len(footer) - 60  # 60 for the gap marker
    tail: list[str] = []
    for text in reversed(step_texts[head_count:]):
        if remaining - len(text) < 0:
            break
        tail.append(text)
        remaining -= len(text)
    tail.reverse()

    skipped = len(step_texts) - head_count - len(tail)
    gap = f"\n... ({skipped} steps omitted — showing first {head_count} + last {len(tail)} steps) ...\n"

    return "\n".join(head) + gap + "\n".join(tail) + footer


def llm_summarize_trajectory(
    trajectory_text: str,
    task_name: str,
    score: float,
) -> str | None:
    """Use a fast LLM to summarize a trajectory into actionable insights.

    Returns a focused ~1500 char summary, or None on failure.
    Uses litellm for provider-agnostic model routing.
    """
    import litellm
    from clawbench._env import DEFAULT_OPTIMIZER_MODEL

    prompt = (
        f"Summarize this agent trajectory for task '{task_name}' (score: {score:.2f}).\n\n"
        "Focus on:\n"
        "1. What approach/strategy did the agent take?\n"
        "2. What tool/API calls did it make and what responses did it get?\n"
        "3. Where did it go wrong (if it failed)? What was the root cause?\n"
        "4. What would have worked better?\n\n"
        "Be specific about commands, endpoints, query parameters, and error messages.\n"
        "Keep it under 1500 characters.\n\n"
        "<trajectory>\n"
        f"{trajectory_text}\n"
        "</trajectory>"
    )

    try:
        response = litellm.completion(
            model=DEFAULT_OPTIMIZER_MODEL,
            max_tokens=1024,
            messages=[{"role": "user", "content": prompt}],
        )
        text = response.choices[0].message.content
        return text.strip() if text else None
    except Exception:
        pass
    return None


def parse_raw_trajectory_steps(trial_dir: Path) -> list[dict]:
    """Parse trajectory JSON into structured step entries for dashboard export.

    Returns a list of step dicts with type/content/tool/command fields.
    """
    traj_path = trial_dir / "agent" / "trajectory.json"
    if not traj_path.exists():
        return []

    raw_steps = []
    try:
        raw_traj = json.loads(traj_path.read_text())
        for step in raw_traj.get("steps", []):
            step_entry: dict = {"type": "thought", "content": ""}
            source = step.get("source", "")
            message = step.get("message", "")
            tool_calls = step.get("tool_calls", []) or []
            observation = step.get("observation") or {}

            if source == "agent" and message and not tool_calls:
                step_entry = {"type": "thought", "content": message[:500]}
            elif tool_calls:
                tc = tool_calls[0]
                fn_name = tc.get("function_name", "unknown")
                args = tc.get("arguments", {})
                step_entry = {
                    "type": "tool",
                    "tool": fn_name,
                    "command": args.get("command", ""),
                    "description": args.get("description", ""),
                }
                if fn_name != "Bash":
                    step_entry["content"] = json.dumps(args)[:300]

            obs_results = observation.get("results", []) or []
            if obs_results:
                obs_text = obs_results[0].get("content", "")[:500] if obs_results else ""
                if obs_text:
                    raw_steps.append(step_entry)
                    raw_steps.append({"type": "observation", "content": obs_text})
                    continue

            if step_entry.get("content") or step_entry.get("command") or step_entry.get("tool"):
                raw_steps.append(step_entry)
    except (json.JSONDecodeError, OSError):
        pass

    return raw_steps
