"""Validate Executor proposals against runtime-observed trace events."""

from __future__ import annotations

from dataclasses import dataclass
from enum import StrEnum

from .memory import ObservationRecord, TaskLease
from .trace import EpisodeTrace, has_action_receipts, is_grounding_receipt


class ExecutionOutcome(StrEnum):
    DONE = "done"
    PROGRESS = "progress"
    BLOCKED = "blocked"
    FAILED = "failed"


@dataclass(frozen=True)
class ExecutionDelta:
    task_id: str
    outcome: ExecutionOutcome
    summary: str
    evidence_excerpt: str | None


@dataclass(frozen=True)
class ValidExecution:
    run_id: str
    task_id: str
    attempt_id: str
    lease_revision: int
    trace_episode_id: str
    outcome: ExecutionOutcome
    summary: str
    observation: str | None
    evidence_sequences: tuple[int, ...]
    recovered_transport_failure: bool = False
    evidence_span_widened: bool = False
    evidence_projected: bool = False
    evidence_fallback: bool = False
    evidence_truncated: bool = False
    evidence_unresolved: bool = False
    reused_observation_id: str | None = None


class ExecutionValidationError(ValueError):
    pass


def _newline_normalized_with_offsets(value: str) -> tuple[str, tuple[int, ...]]:
    characters: list[str] = []
    offsets: list[int] = []
    index = 0
    while index < len(value):
        offsets.append(index)
        if value.startswith("\r\n", index):
            characters.append("\n")
            index += 2
        else:
            characters.append(value[index])
            index += 1
    return "".join(characters), tuple(offsets)


def _exact_receipt_slice(output: str, excerpt: str) -> str | None:
    exact_start = output.find(excerpt)
    if exact_start >= 0:
        return output[exact_start : exact_start + len(excerpt)]

    normalized_output, offsets = _newline_normalized_with_offsets(output)
    normalized_excerpt, _ = _newline_normalized_with_offsets(excerpt)
    normalized_start = normalized_output.find(normalized_excerpt)
    if normalized_start < 0:
        return None
    normalized_end = normalized_start + len(normalized_excerpt)
    original_start = offsets[normalized_start]
    original_end = offsets[normalized_end] if normalized_end < len(offsets) else len(output)
    return output[original_start:original_end]


def _ordered_exact_line_span(output: str, excerpt: str) -> str | None:
    """Widen exact ordered lines to one bounded contiguous receipt span."""
    normalized_output, offsets = _newline_normalized_with_offsets(output)
    normalized_excerpt, _ = _newline_normalized_with_offsets(excerpt)
    wanted_lines = normalized_excerpt.splitlines()
    if len(wanted_lines) < 2 or any(not line for line in wanted_lines):
        return None

    output_lines: list[tuple[str, int, int]] = []
    start = 0
    for line_with_ending in normalized_output.splitlines(keepends=True):
        line = line_with_ending.removesuffix("\n")
        output_lines.append((line, start, start + len(line)))
        start += len(line_with_ending)
    if start < len(normalized_output):
        output_lines.append((normalized_output[start:], start, len(normalized_output)))

    matched: list[tuple[int, int]] = []
    search_from = 0
    for wanted in wanted_lines:
        for index in range(search_from, len(output_lines)):
            line, line_start, line_end = output_lines[index]
            if line != wanted:
                continue
            matched.append((line_start, line_end))
            search_from = index + 1
            break
        else:
            return None

    normalized_start = matched[0][0]
    normalized_end = matched[-1][1]
    original_start = offsets[normalized_start]
    original_end = offsets[normalized_end] if normalized_end < len(offsets) else len(output)
    span = output[original_start:original_end]
    return span if len(span) <= 4_000 else None


def _unique_line_envelope(output: str, wanted_lines: list[str]) -> str | None:
    normalized_output, offsets = _newline_normalized_with_offsets(output)
    if not wanted_lines or len(set(wanted_lines)) != len(wanted_lines):
        return None

    output_lines: list[tuple[str, int, int]] = []
    start = 0
    for line_with_ending in normalized_output.splitlines(keepends=True):
        line = line_with_ending.removesuffix("\n")
        output_lines.append((line, start, start + len(line)))
        start += len(line_with_ending)
    if start < len(normalized_output):
        output_lines.append((normalized_output[start:], start, len(normalized_output)))

    matched: list[tuple[int, int]] = []
    for wanted in wanted_lines:
        occurrences = [
            (line_start, line_end) for line, line_start, line_end in output_lines if line == wanted
        ]
        if len(occurrences) != 1:
            return None
        matched.append(occurrences[0])

    normalized_start = min(line_start for line_start, _ in matched)
    normalized_end = max(line_end for _, line_end in matched)
    original_start = offsets[normalized_start]
    original_end = offsets[normalized_end] if normalized_end < len(offsets) else len(output)
    span = output[original_start:original_end]
    return span if len(span) <= 4_000 else None


def _exact_unique_line_envelope(output: str, excerpt: str) -> str | None:
    """Recover the true receipt span from a rich set of exact, unambiguous lines."""
    normalized_excerpt, _ = _newline_normalized_with_offsets(excerpt)
    wanted_lines = [line for line in normalized_excerpt.splitlines() if line]
    if len(wanted_lines) < 4:
        return None
    return _unique_line_envelope(output, wanted_lines)


@dataclass(frozen=True)
class _ReceiptMatch:
    sequence: int
    observation: str
    widened: bool = False
    projected: bool = False
    truncated: bool = False


def _grounding_receipts(trace: EpisodeTrace) -> tuple[tuple[int, str], ...]:
    structured_output_calls = {
        event.get("call_id")
        for event in trace.events
        if event.get("type") == "tool_call" and event.get("name") == "StructuredOutput"
    }
    return tuple(
        (int(event["sequence"]), str(event.get("output", "")))
        for event in trace.events
        if is_grounding_receipt(event, structured_output_calls) and str(event.get("output", ""))
    )


def _supported_single_receipt_projection(
    receipts: tuple[tuple[int, str], ...], excerpt: str
) -> _ReceiptMatch | None:
    """Project an exact multi-receipt quote onto one rich source receipt."""
    normalized_excerpt, _ = _newline_normalized_with_offsets(excerpt)
    wanted_lines = [line for line in normalized_excerpt.splitlines() if line]
    if len(wanted_lines) < 2 or len(set(wanted_lines)) != len(wanted_lines):
        return None

    matches_by_receipt: list[tuple[int, str, list[str]]] = []
    supported_lines: set[str] = set()
    for sequence, output in receipts:
        normalized_output, _ = _newline_normalized_with_offsets(output)
        output_lines = normalized_output.splitlines()
        matched_lines = [line for line in wanted_lines if output_lines.count(line) == 1]
        supported_lines.update(matched_lines)
        matches_by_receipt.append((sequence, output, matched_lines))
    if supported_lines != set(wanted_lines):
        return None
    if any(set(matched_lines) == supported_lines for _, _, matched_lines in matches_by_receipt):
        return None

    candidates: list[tuple[int, int, int, str]] = []
    for sequence, output, matched_lines in matches_by_receipt:
        if len(matched_lines) < 2:
            continue
        span = _unique_line_envelope(output, matched_lines)
        if span is None:
            continue
        candidates.append((sum(map(len, matched_lines)), len(matched_lines), sequence, span))
    if not candidates:
        return None
    _, _, sequence, span = max(candidates)
    return _ReceiptMatch(sequence, span, widened=True, projected=True)


def _matching_receipt(trace: EpisodeTrace, excerpt: str) -> _ReceiptMatch | None:
    receipts = _grounding_receipts(trace)
    if not excerpt.strip():
        return None

    for sequence, output in receipts:
        exact_slice = _exact_receipt_slice(output, excerpt)
        if exact_slice is not None:
            if len(exact_slice) > 4_000:
                return _ReceiptMatch(
                    sequence,
                    exact_slice[-4_000:],
                    truncated=True,
                )
            return _ReceiptMatch(sequence, exact_slice)
        widened_span = _ordered_exact_line_span(output, excerpt)
        if widened_span is not None:
            return _ReceiptMatch(sequence, widened_span, widened=True)
        exact_line_envelope = _exact_unique_line_envelope(output, excerpt)
        if exact_line_envelope is not None:
            return _ReceiptMatch(sequence, exact_line_envelope, widened=True)
    return _supported_single_receipt_projection(receipts, excerpt)


def _last_bounded_grounding_receipt(trace: EpisodeTrace) -> _ReceiptMatch | None:
    for sequence, output in reversed(_grounding_receipts(trace)):
        if output.strip():
            return _ReceiptMatch(
                sequence,
                output[-4_000:],
                truncated=len(output) > 4_000,
            )
    return None


def _matching_prior_observation(
    observations: tuple[ObservationRecord, ...],
    task_id: str,
    excerpt: str,
) -> ObservationRecord | None:
    """Find the newest same-task canonical observation containing the exact excerpt."""
    for observation in sorted(
        observations,
        key=lambda candidate: candidate.created_revision,
        reverse=True,
    ):
        if observation.task_id != task_id:
            continue
        if _exact_receipt_slice(observation.statement, excerpt) is not None:
            return observation
    return None


def compile_execution(
    delta: ExecutionDelta,
    lease: TaskLease,
    trace: EpisodeTrace,
    *,
    recover_structured_transport: bool = False,
    allow_evidence_fallback: bool = True,
    prior_observations: tuple[ObservationRecord, ...] = (),
) -> ValidExecution:
    if delta.task_id != lease.task_id:
        raise ExecutionValidationError("Executor returned a result for the wrong task")
    transport_recoverable = recover_structured_transport and any(
        event.get("type") == "turn_completed"
        and event.get("stop_reason") in {"error_max_structured_output_retries", "error_max_turns"}
        for event in trace.events
    )
    if trace.output is None or (
        trace.output.get("success") is not True and not transport_recoverable
    ):
        raise ExecutionValidationError("Executor episode did not complete successfully")
    if trace.input.get("role") != "executor":
        raise ExecutionValidationError("evidence trace does not belong to an Executor")
    if trace.input.get("run_id") != lease.run_id:
        raise ExecutionValidationError("evidence trace belongs to a different run")
    if trace.input.get("state_revision") != lease.revision:
        raise ExecutionValidationError("evidence trace belongs to a different state revision")
    if trace.input.get("task_id") != lease.task_id:
        raise ExecutionValidationError("evidence trace belongs to a different task")
    if trace.input.get("attempt_id") != lease.attempt_id:
        raise ExecutionValidationError("evidence trace belongs to a different attempt")
    if (
        lease.trace_episode_id is not None
        and trace.input.get("episode_id") != lease.trace_episode_id
    ):
        raise ExecutionValidationError("evidence trace belongs to a different episode")
    if not delta.summary.strip():
        raise ExecutionValidationError("execution summary must not be empty")
    if len(delta.summary) > 2_000:
        raise ExecutionValidationError("execution summary exceeds 2000 characters")
    receipt_match = (
        _matching_receipt(trace, delta.evidence_excerpt) if delta.evidence_excerpt else None
    )
    reused_observation = (
        _matching_prior_observation(
            prior_observations,
            lease.task_id,
            delta.evidence_excerpt,
        )
        if receipt_match is None and delta.evidence_excerpt
        else None
    )
    evidence_fallback = False
    if (
        allow_evidence_fallback
        and receipt_match is None
        and reused_observation is None
        and (delta.evidence_excerpt or delta.outcome is ExecutionOutcome.DONE)
    ):
        receipt_match = _last_bounded_grounding_receipt(trace)
        evidence_fallback = receipt_match is not None
    evidence_sequences = (receipt_match.sequence,) if receipt_match is not None else ()

    has_grounded_evidence = bool(evidence_sequences) or reused_observation is not None
    evidence_unresolved = (
        receipt_match is None
        and reused_observation is None
        and bool(delta.evidence_excerpt)
        and delta.outcome is ExecutionOutcome.DONE
        and not has_action_receipts(trace)
        and any(observation.task_id == lease.task_id for observation in prior_observations)
    )
    if delta.evidence_excerpt and not has_grounded_evidence and not evidence_unresolved:
        raise ExecutionValidationError("the proposed evidence was not captured")
    observation = receipt_match.observation if receipt_match is not None else None
    outcome = (
        ExecutionOutcome.PROGRESS
        if (
            evidence_fallback
            or bool(receipt_match and receipt_match.truncated)
            or evidence_unresolved
        )
        and delta.outcome is ExecutionOutcome.DONE
        else delta.outcome
    )

    if outcome is ExecutionOutcome.DONE:
        if not observation and reused_observation is None:
            raise ExecutionValidationError("a completed task must propose an observation")
        if not has_grounded_evidence:
            raise ExecutionValidationError("a completed task must cite captured evidence")

    episode_id = trace.input.get("episode_id")
    if not isinstance(episode_id, str):
        raise ExecutionValidationError("trace has no episode identity")
    return ValidExecution(
        run_id=lease.run_id,
        task_id=lease.task_id,
        attempt_id=lease.attempt_id,
        lease_revision=lease.revision,
        trace_episode_id=episode_id,
        outcome=outcome,
        summary=delta.summary.strip(),
        observation=observation,
        evidence_sequences=evidence_sequences,
        recovered_transport_failure=transport_recoverable,
        evidence_span_widened=bool(receipt_match and receipt_match.widened),
        evidence_projected=bool(receipt_match and receipt_match.projected),
        evidence_fallback=evidence_fallback,
        evidence_truncated=bool(receipt_match and receipt_match.truncated),
        evidence_unresolved=evidence_unresolved,
        reused_observation_id=(reused_observation.id if reused_observation is not None else None),
    )


__all__ = [
    "ExecutionDelta",
    "ExecutionOutcome",
    "ExecutionValidationError",
    "ValidExecution",
    "compile_execution",
]
