"""Audit a completed trial from canonical SQLite state and immutable episode traces."""

from __future__ import annotations

import argparse
import json
import re
import sqlite3
from collections import Counter
from pathlib import Path
from typing import Any

from .plan import TaskKind, TaskStatus
from .trace import is_grounding_receipt

_SUPERVISOR_EPISODE_ID = re.compile(r"supervisor-r(?P<revision>\d+)(?:-a\d+)?\Z")


def _rows(database: Path, query: str) -> list[dict[str, Any]]:
    connection = sqlite3.connect(f"file:{database}?mode=ro", uri=True)
    connection.row_factory = sqlite3.Row
    try:
        return [dict(row) for row in connection.execute(query).fetchall()]
    finally:
        connection.close()


def _table_exists(database: Path, table: str) -> bool:
    rows = _rows(
        database,
        f"SELECT name FROM sqlite_master WHERE type = 'table' AND name = '{table}'",
    )
    return bool(rows)


def _is_hidden_provider_memory_action(action: dict[str, Any]) -> bool:
    if action.get("type") != "tool_call" or action.get("name") not in {"Write", "Edit"}:
        return False
    tool_input = action.get("input")
    if not isinstance(tool_input, dict):
        return False
    path = str(tool_input.get("file_path") or tool_input.get("path") or "").replace("\\", "/")
    return "/.claude/projects/" in path and "/memory/" in path


def _load_trace_object(path: Path, integrity_errors: list[str]) -> dict[str, Any] | None:
    if not path.exists():
        integrity_errors.append(f"missing {path.name}")
        return None
    try:
        value = json.loads(path.read_text(encoding="utf-8"))
    except (json.JSONDecodeError, UnicodeError):
        integrity_errors.append(f"malformed {path.name}")
        return None
    except OSError:
        integrity_errors.append(f"unreadable {path.name}")
        return None
    if not isinstance(value, dict):
        integrity_errors.append(f"malformed {path.name}")
        return None
    return value


def _episode_records(trace_root: Path) -> list[dict[str, Any]]:
    episodes: list[dict[str, Any]] = []
    if not trace_root.exists():
        return episodes
    for episode_dir in trace_root.iterdir():
        if not episode_dir.is_dir():
            continue
        integrity_errors: list[str] = []
        input_data = _load_trace_object(episode_dir / "input.json", integrity_errors) or {}
        output_path = episode_dir / "output.json"
        output_data = _load_trace_object(output_path, integrity_errors)
        events: list[dict[str, Any]] = []
        truncated_tail = False
        events_path = episode_dir / "events.jsonl"
        if events_path.exists():
            try:
                event_lines = events_path.read_text(encoding="utf-8").splitlines()
            except (OSError, UnicodeError):
                integrity_errors.append("unreadable events.jsonl")
                event_lines = []
            for line_number, line in enumerate(event_lines, start=1):
                if not line.strip():
                    continue
                try:
                    event = json.loads(line)
                except json.JSONDecodeError:
                    truncated_tail = True
                    integrity_errors.append(f"malformed events.jsonl at line {line_number}")
                    break
                if not isinstance(event, dict):
                    truncated_tail = True
                    integrity_errors.append(f"malformed events.jsonl at line {line_number}")
                    break
                events.append(event)
        else:
            integrity_errors.append("missing events.jsonl")
        event_counts = Counter(str(event.get("type", "unknown")) for event in events)
        actions = [
            {
                key: event.get(key)
                for key in (
                    "sequence",
                    "type",
                    "name",
                    "command",
                    "input",
                    "output",
                    "exit_code",
                    "is_error",
                    "call_id",
                )
                if key in event
            }
            for event in events
            if event.get("type") in {"tool_call", "tool_result", "command_run"}
        ]
        episodes.append(
            {
                "episode_id": episode_dir.name,
                "run_id": input_data.get("run_id"),
                "input_episode_id": input_data.get("episode_id"),
                "role": input_data.get("role"),
                "state_revision": input_data.get("state_revision"),
                "task_id": input_data.get("task_id"),
                "attempt_id": input_data.get("attempt_id"),
                "opened_at": input_data.get("opened_at"),
                "closed_at": output_data.get("closed_at")
                if isinstance(output_data, dict)
                else None,
                "success": output_data.get("success") if isinstance(output_data, dict) else False,
                "duration_ms": (
                    output_data.get("duration_ms") if isinstance(output_data, dict) else None
                ),
                "cost_usd": (
                    float(output_data.get("cost_usd", 0) or 0)
                    if isinstance(output_data, dict)
                    else 0.0
                ),
                "usage": (output_data.get("usage", {}) if isinstance(output_data, dict) else {}),
                "structured_output": (
                    output_data.get("structured_output") if isinstance(output_data, dict) else None
                ),
                "event_counts": dict(sorted(event_counts.items())),
                "actions": actions,
                "events": events,
                "truncated_tail": truncated_tail,
                "complete": not integrity_errors and output_data is not None and not truncated_tail,
                "integrity_errors": integrity_errors,
            }
        )
    episodes.sort(key=lambda episode: float(episode.get("opened_at") or 0))
    return episodes


def audit_run(run_dir: str | Path, *, expected_flag: str) -> dict[str, Any]:
    """Return a full, deterministic audit of one trial artifact directory."""
    root = Path(run_dir)
    database = root / "state.sqlite3"
    runs = _rows(database, "SELECT * FROM runs")
    if len(runs) != 1:
        raise ValueError(f"expected exactly one run row, found {len(runs)}")
    tasks = _rows(database, "SELECT * FROM tasks ORDER BY created_revision, task_id")
    attempts = _rows(database, "SELECT * FROM attempts ORDER BY started_revision, attempt_id")
    observations = _rows(
        database,
        "SELECT * FROM observations ORDER BY created_revision, observation_id",
    )
    transitions = (
        _rows(database, "SELECT * FROM transitions ORDER BY revision")
        if _table_exists(database, "transitions")
        else []
    )
    episodes = _episode_records(root / "traces")
    episodes_by_id = {str(episode["episode_id"]): episode for episode in episodes}
    attempts_by_id = {str(attempt["attempt_id"]): attempt for attempt in attempts}

    grounded = True
    direct_quotes = True
    observation_identities_match = True
    for observation in observations:
        episode = episodes_by_id.get(str(observation["trace_episode_id"]))
        attempt = attempts_by_id.get(str(observation["attempt_id"]))
        if episode is None or attempt is None:
            grounded = False
            observation_identities_match = False
            break
        if (
            episode.get("role") != "executor"
            or episode.get("run_id") != runs[0]["run_id"]
            or episode.get("input_episode_id") != observation["trace_episode_id"]
            or episode.get("task_id") != observation["task_id"]
            or episode.get("attempt_id") != observation["attempt_id"]
            or episode.get("state_revision") != attempt["started_revision"]
            or attempt["run_id"] != runs[0]["run_id"]
            or attempt["task_id"] != observation["task_id"]
            or attempt["trace_episode_id"] != observation["trace_episode_id"]
        ):
            observation_identities_match = False
        events_by_sequence = {
            int(event["sequence"]): event
            for event in episode["events"]
            if isinstance(event.get("sequence"), int)
        }
        sequences = json.loads(str(observation["evidence_sequences"]))
        if len(sequences) != 1:
            direct_quotes = False
        structured_output_calls = {
            event.get("call_id")
            for event in episode["events"]
            if event.get("type") == "tool_call" and event.get("name") == "StructuredOutput"
        }
        for sequence in sequences:
            event = events_by_sequence.get(int(sequence))
            if event is None or not is_grounding_receipt(event, structured_output_calls):
                grounded = False
                break
            if str(observation["statement"]) not in str(event.get("output", "")):
                direct_quotes = False
        if not grounded:
            break

    observations_by_id = {
        str(observation["observation_id"]): observation for observation in observations
    }
    tasks_by_id = {str(task["task_id"]): task for task in tasks}
    basis_ids_exist = True
    basis_producers_are_dependencies = True
    exploit_bases_current = True
    for task in tasks:
        basis_ids = tuple(json.loads(str(task["basis_ids"])))
        dependency_ids = set(json.loads(str(task["depends_on"])))
        resolved_basis = [observations_by_id.get(str(basis_id)) for basis_id in basis_ids]
        if any(observation is None for observation in resolved_basis):
            basis_ids_exist = False
        producers = {
            str(observation["task_id"]) for observation in resolved_basis if observation is not None
        }
        if task["kind"] != TaskKind.RECOVER.value and not producers.issubset(dependency_ids):
            basis_producers_are_dependencies = False
        if task["kind"] != TaskKind.EXPLOIT.value:
            continue
        candidates = [
            observation
            for observation in observations
            if (producer := tasks_by_id.get(str(observation["task_id"]))) is not None
            and producer["kind"] == TaskKind.TEST.value
            and producer["status"] == TaskStatus.DONE.value
            and producer["target"] == task["target"]
        ]
        if not candidates:
            exploit_bases_current = False
            continue
        latest = max(
            candidates,
            key=lambda observation: (
                int(observation["created_revision"]),
                str(observation["observation_id"]),
            ),
        )
        if (
            str(latest["observation_id"]) not in basis_ids
            or str(latest["task_id"]) not in dependency_ids
        ):
            exploit_bases_current = False

    settled_failure_ids = {
        str(attempt["trace_episode_id"])
        for attempt in attempts
        if attempt["status"] != "active" and attempt.get("trace_episode_id")
    }
    transitions_by_revision = {
        int(transition["revision"]): transition for transition in transitions
    }
    failed_episodes_settled = True
    for episode in episodes:
        if not episode["complete"] or episode["success"] is True:
            continue
        if episode["role"] == "executor":
            settled = str(episode["episode_id"]) in settled_failure_ids
        elif episode["role"] == "supervisor" and isinstance(episode["state_revision"], int):
            transition = transitions_by_revision.get(int(episode["state_revision"]) + 1)
            settled = transition is not None and transition["kind"] in {
                "plan_committed",
                "supervisor_failed",
                "decision_limit_reached",
            }
        else:
            settled = False
        failed_episodes_settled = failed_episodes_settled and settled

    accounted_incomplete_episode_ids = {
        str(attempt["trace_episode_id"])
        for attempt in attempts
        if attempt["status"] != "active"
        and attempt.get("trace_episode_id")
        and attempt.get("failure_kind") in {"interrupted", "trace_corrupt"}
    }

    def incomplete_supervisor_is_settled(episode: dict[str, Any]) -> bool:
        match = _SUPERVISOR_EPISODE_ID.fullmatch(str(episode["episode_id"]))
        if match is None:
            return False
        transition = transitions_by_revision.get(int(match.group("revision")) + 1)
        return transition is not None and transition["kind"] in {
            "plan_committed",
            "supervisor_failed",
            "decision_limit_reached",
        }

    episode_journals_accounted_for = bool(episodes) and all(
        episode["complete"]
        or str(episode["episode_id"]) in accounted_incomplete_episode_ids
        or incomplete_supervisor_is_settled(episode)
        for episode in episodes
    )
    expected_revisions = list(range(int(runs[0]["revision"]) + 1))
    actual_revisions = [int(transition["revision"]) for transition in transitions]
    hidden_provider_memory_actions = [
        action
        for episode in episodes
        for action in episode["actions"]
        if _is_hidden_provider_memory_action(action)
    ]
    supervisor_actions = sum(
        action.get("type") == "command_run"
        or (action.get("type") == "tool_call" and action.get("name") != "StructuredOutput")
        for episode in episodes
        if episode["role"] == "supervisor"
        for action in episode["actions"]
    )

    completion_basis_valid = False
    if runs[0]["status"] == "completed" and transitions:

        def completed_basis(observation_id: object) -> bool:
            basis_observation = observations_by_id.get(str(observation_id))
            if basis_observation is None:
                return False
            producer_task = tasks_by_id.get(str(basis_observation["task_id"]))
            return producer_task is not None and producer_task["status"] == TaskStatus.DONE.value

        final_transition = transitions[-1]
        try:
            final_detail = json.loads(str(final_transition["detail"]))
        except (json.JSONDecodeError, TypeError):
            final_detail = {}
        finish_basis_ids = final_detail.get("finish_basis_ids", [])
        completion_basis_valid = (
            final_transition["kind"] == "plan_committed"
            and final_detail.get("finish") is True
            and isinstance(finish_basis_ids, list)
            and 1 <= len(finish_basis_ids) <= 4
            and len(finish_basis_ids) == len(set(finish_basis_ids))
            and all(completed_basis(observation_id) for observation_id in finish_basis_ids)
        )

    canonical_text = "\n".join(str(observation["statement"]) for observation in observations)
    checks = {
        "run_completed": runs[0]["status"] == "completed",
        "oracle_in_canonical_observation": expected_flag.casefold() in canonical_text.casefold(),
        "no_active_tasks": all(task["status"] != "active" for task in tasks),
        "no_active_attempts": all(attempt["status"] != "active" for attempt in attempts),
        "all_episodes_complete": episode_journals_accounted_for,
        "all_observations_grounded": bool(observations) and grounded,
        "all_observations_are_direct_quotes": bool(observations) and direct_quotes,
        "all_observation_identities_match": bool(observations) and observation_identities_match,
        "all_basis_ids_exist": basis_ids_exist,
        "all_basis_producers_are_dependencies": basis_producers_are_dependencies,
        "all_exploit_bases_current": exploit_bases_current,
        "completion_basis_valid": completion_basis_valid,
        "transition_timeline_complete": actual_revisions == expected_revisions,
        "all_failed_episodes_settled": failed_episodes_settled,
    }
    usage_fields = (
        "input_tokens",
        "cached_input_tokens",
        "output_tokens",
        "reasoning_output_tokens",
    )
    totals = {
        "episodes": len(episodes),
        "supervisor_episodes": sum(episode["role"] == "supervisor" for episode in episodes),
        "executor_episodes": sum(episode["role"] == "executor" for episode in episodes),
        "tasks": len(tasks),
        "attempts": len(attempts),
        "observations": len(observations),
        "transitions": len(transitions),
        "supervisor_actions": supervisor_actions,
        "hidden_provider_memory_actions": len(hidden_provider_memory_actions),
        "cost_usd": sum(float(episode["cost_usd"]) for episode in episodes),
        "provider_duration_ms": sum(int(episode["duration_ms"] or 0) for episode in episodes),
        "usage": {
            field: sum(int(episode["usage"].get(field, 0) or 0) for episode in episodes)
            for field in usage_fields
        },
    }
    return {
        "schema_version": 2,
        "run": runs[0],
        "expected_flag": expected_flag,
        "passed": all(checks.values()),
        "checks": checks,
        "totals": totals,
        "tasks": tasks,
        "attempts": attempts,
        "observations": observations,
        "transitions": transitions,
        "episodes": episodes,
    }


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--run-dir", type=Path, required=True)
    parser.add_argument("--expected-flag", required=True)
    parser.add_argument("--output", type=Path)
    args = parser.parse_args(argv)
    result = audit_run(args.run_dir, expected_flag=args.expected_flag)
    rendered = json.dumps(result, indent=2, sort_keys=True) + "\n"
    if args.output:
        args.output.write_text(rendered, encoding="utf-8")
        print(
            "AUDIT_RESULT="
            + json.dumps(
                {
                    "run_id": result["run"]["run_id"],
                    "passed": result["passed"],
                    "checks": result["checks"],
                    "totals": result["totals"],
                    "output": str(args.output),
                },
                sort_keys=True,
            )
        )
    else:
        print(rendered, end="")
    return 0 if result["passed"] else 1


if __name__ == "__main__":  # pragma: no cover
    raise SystemExit(main())


__all__ = ["audit_run", "main"]
