src/yaml_edit/engine.py

"""Allowlisted scalar edits with exact snapshot binding and post-edit verification."""

import json
from collections.abc import Sequence
from difflib import unified_diff
from typing import Never, cast

from ruamel.yaml import YAML
from ruamel.yaml.nodes import ScalarNode

from .models import (
    MAX_BATCH_CANDIDATE_BYTES,
    MAX_BATCH_OPERATIONS,
    MAX_BATCH_SCALAR_BYTES,
    BatchEntry,
    BatchPreview,
    BatchProposal,
    BatchRequest,
    Candidate,
    EditPreview,
    EditProposal,
    ErrorCode,
    Inspection,
    Path,
    Scalar,
    Tree,
    YamlEditError,
    is_scalar,
    validate_batch_proposal,
)
from .parsing import Parsed, parse_source
from .validation import validate_value


def inspect(source: str) -> Inspection:
    """Discover editable fields and bind their identities to this exact source."""
    parsed = parse_source(source)
    return Inspection(parsed.source_version, parsed.candidates)


def _candidate(parsed: Parsed, candidate_id: str) -> Candidate:
    for candidate in parsed.candidates:
        if candidate.id == candidate_id:
            return candidate
    raise YamlEditError(
        ErrorCode.UNKNOWN_TARGET, f"Not an editable existing field: {candidate_id!r}."
    )


def _same(left: Tree, right: Tree) -> bool:
    """Compare semantic trees without conflating bool/int or int/float."""
    if type(left) is not type(right):
        return False
    if isinstance(left, dict) and isinstance(right, dict):
        return left.keys() == right.keys() and all(_same(left[key], right[key]) for key in left)
    if isinstance(left, list) and isinstance(right, list):
        return len(left) == len(right) and all(
            _same(a, b) for a, b in zip(left, right, strict=True)
        )
    return left == right


def _with_value(tree: Tree, path: Path, value: Scalar) -> Tree:
    """Copy only the selected ancestor chain, retaining the original semantic tree."""
    if not path:
        return value
    head, *tail = path
    if isinstance(tree, dict) and isinstance(head, str) and head in tree:
        return {**tree, head: _with_value(tree[head], tuple(tail), value)}
    if isinstance(tree, list) and type(head) is int and 0 <= head < len(tree):
        result = tree.copy()
        result[head] = _with_value(tree[head], tuple(tail), value)
        return result
    raise YamlEditError(
        ErrorCode.INVARIANT_VIOLATION, "Discovered path does not match the semantic tree."
    )


def _plain_string(value: str) -> bool:
    # Avoid reinterpreting strings as booleans/numbers in either YAML 1.2 or
    # common YAML 1.1 consumers. Parsing never runs arbitrary constructors here.
    for version in ((1, 2), (1, 1)):
        parser = YAML(typ="safe", pure=True)
        parser.version = version
        node = parser.compose(value)
        if not isinstance(node, ScalarNode) or node.tag != "tag:yaml.org,2002:str":
            return False
        if node.value != value:
            return False
    return True


def _replacement(value: Scalar, original: ScalarNode) -> str:
    if isinstance(value, str):
        if original.style == "'":
            return "'" + value.replace("'", "''") + "'"
        if original.style != '"' and _plain_string(value):
            return value
    return json.dumps(value, ensure_ascii=False, allow_nan=False)


def _unified_diff(before: str, after: str) -> str:
    if before == after:
        return ""
    # Review output uses LF regardless of the source's newline convention.
    # The returned source itself remains byte-for-byte preserved outside the edit.
    lines = unified_diff(
        before.replace("\r\n", "\n").replace("\r", "\n").splitlines(keepends=True),
        after.replace("\r\n", "\n").replace("\r", "\n").splitlines(keepends=True),
        fromfile="before.yaml",
        tofile="after.yaml",
    )
    # difflib doesn't supply the conventional missing-newline marker itself.
    return "".join(
        line if line.endswith("\n") else line + "\n\\ No newline at end of file\n" for line in lines
    )


def _propose_parsed(parsed: Parsed, source: str, candidate_id: str, value: object) -> EditProposal:
    """Create a proposal from a private, request-local parsed source snapshot."""
    candidate = _candidate(parsed, candidate_id)
    proposal = EditProposal(
        source_version=parsed.source_version,
        candidate_id=candidate.id,
        path=candidate.path,
        expected_old_value=candidate.current_value,
        value=validate_value(candidate.kind, value),
    )
    # Includes post-edit structure, request/limit, and invariant checks.
    _preview(parsed, source, proposal)
    return proposal


def propose(source: str, candidate_id: str, value: object) -> EditProposal:
    """Create a validated proposal. No files are read or written."""
    return _propose_parsed(parse_source(source), source, candidate_id, value)


def _preview(parsed: Parsed, source: str, proposal: EditProposal) -> EditPreview:
    if parsed.source_version != proposal.source_version:
        raise YamlEditError(
            ErrorCode.STALE_PROPOSAL, "Source changed; generate and review a new proposal."
        )
    candidate = _candidate(parsed, proposal.candidate_id)
    if candidate.path != proposal.path or not _same(
        candidate.current_value, proposal.expected_old_value
    ):
        raise YamlEditError(
            ErrorCode.INVALID_PROPOSAL, "Proposal does not match its allowlisted target."
        )
    value = validate_value(candidate.kind, proposal.value)
    if _same(value, candidate.current_value):
        return EditPreview(source=source, diff="", changed=False, candidate=candidate)
    original = parsed.nodes[candidate.path]
    if not isinstance(original, ScalarNode):
        raise YamlEditError(ErrorCode.INVARIANT_VIOLATION, "Discovered target is not a scalar.")
    start, end = candidate.location.start, candidate.location.end
    result = source[:start] + _replacement(value, original) + source[end:]
    reparsed = parse_source(result)
    expected = _with_value(parsed.data, candidate.path, value)
    if not _same(expected, reparsed.data):
        raise YamlEditError(
            ErrorCode.INVARIANT_VIOLATION, "Edit changed more than the intended semantic value."
        )
    return EditPreview(
        source=result, diff=_unified_diff(source, result), changed=True, candidate=candidate
    )


def preview(source: str, proposal: object) -> EditPreview:
    """Revalidate an untrusted proposal and return the exact source and review diff."""
    checked = EditProposal.from_dict(
        proposal.to_dict() if isinstance(proposal, EditProposal) else proposal
    )
    parsed = parse_source(source)
    return _preview(parsed, source, checked)


def apply(source: str, proposal: object) -> str:
    """Apply to an immutable string snapshot; callers handle persistence/confirmation."""
    return preview(source, proposal).source


def _batch_invalid() -> Never:
    raise YamlEditError(ErrorCode.INVALID_PROPOSAL, "Malformed or unsupported batch request.")


def _bounded_request_text(value: object, maximum: int, *, pointer: bool = False) -> str:
    if type(value) is not str or len(value) > maximum or (pointer and not value.startswith("/")):
        _batch_invalid()
    text = value
    try:
        raw = text.encode("utf-8")
    except UnicodeEncodeError as error:
        raise YamlEditError(ErrorCode.INVALID_PROPOSAL, "Malformed batch request text.") from error
    if len(raw) > maximum:
        _batch_invalid()
    return text


def _batch_requests(value: object) -> tuple[BatchRequest, ...]:
    if type(value) not in (list, tuple):
        _batch_invalid()
    requests = cast("list[object] | tuple[object, ...]", value)
    if not 1 <= len(requests) <= MAX_BATCH_OPERATIONS:
        _batch_invalid()
    checked: list[BatchRequest] = []
    for item in requests:
        if type(item) is not BatchRequest:
            _batch_invalid()
        candidate_value = cast("object", getattr(item, "candidate_id", None))
        requested_value = cast("object", getattr(item, "value", None))
        candidate_id = _bounded_request_text(
            candidate_value, MAX_BATCH_CANDIDATE_BYTES, pointer=True
        )
        if not is_scalar(requested_value):
            _batch_invalid()
        if type(requested_value) is str:
            _bounded_request_text(requested_value, MAX_BATCH_SCALAR_BYTES)
        checked.append(BatchRequest(candidate_id, requested_value))
    return tuple(checked)


def _resolve_requests(
    parsed: Parsed, requests: tuple[BatchRequest, ...]
) -> tuple[tuple[Candidate, Scalar], ...]:
    seen: set[str] = set()
    resolved: list[tuple[Candidate, Scalar]] = []
    for request in requests:
        if request.candidate_id in seen:
            _batch_invalid()
        seen.add(request.candidate_id)
        candidate = _candidate(parsed, request.candidate_id)
        resolved.append((candidate, validate_value(candidate.kind, request.value)))
    return tuple(sorted(resolved, key=lambda item: (item[0].location.start, item[0].id)))


def _proposal_entries(
    resolved: tuple[tuple[Candidate, Scalar], ...],
) -> tuple[BatchEntry, ...]:
    return tuple(
        BatchEntry(candidate.id, candidate.path, candidate.current_value, value)
        for candidate, value in resolved
    )


def propose_batch(source: str, requests: Sequence[BatchRequest]) -> BatchProposal:
    """Create one validated proposal for explicit existing scalar targets."""
    checked_requests = _batch_requests(requests)
    parsed = parse_source(source)
    resolved = _resolve_requests(parsed, checked_requests)
    proposal = BatchProposal(parsed.source_version, _proposal_entries(resolved))
    _preview_batch(parsed, source, proposal)
    proposal.to_dict()
    return proposal


def _entry_matches(candidate: Candidate, entry: BatchEntry) -> bool:
    return candidate.path == entry.path and _same(candidate.current_value, entry.expected_old_value)


def _resolve_entries(
    parsed: Parsed, proposal: BatchProposal, source_length: int
) -> tuple[tuple[Candidate, Scalar], ...]:
    if parsed.source_version != proposal.source_version:
        raise YamlEditError(
            ErrorCode.STALE_PROPOSAL, "Source changed; generate and review a new batch proposal."
        )
    ids: set[str] = set()
    paths: set[Path] = set()
    spans: set[tuple[int, int]] = set()
    resolved: list[tuple[Candidate, Scalar]] = []
    for entry in proposal.entries:
        candidate = _candidate(parsed, entry.candidate_id)
        if not _entry_matches(candidate, entry):
            raise YamlEditError(
                ErrorCode.INVALID_PROPOSAL, "Batch entry does not match its allowlisted target."
            )
        span = (candidate.location.start, candidate.location.end)
        if candidate.id in ids or candidate.path in paths or span in spans:
            raise YamlEditError(ErrorCode.INVALID_PROPOSAL, "Batch targets must be distinct.")
        ids.add(candidate.id)
        paths.add(candidate.path)
        spans.add(span)
        resolved.append((candidate, validate_value(candidate.kind, entry.value)))
    canonical = sorted(resolved, key=lambda item: (item[0].location.start, item[0].id))
    if resolved != canonical:
        raise YamlEditError(ErrorCode.INVALID_PROPOSAL, "Batch entries are not in source order.")
    _check_batch_spans(resolved, source_length)
    return tuple(resolved)


def _check_batch_spans(resolved: list[tuple[Candidate, Scalar]], source_length: int) -> None:
    previous_end = -1
    for candidate, _value in resolved:
        location = candidate.location
        if not 0 <= location.start <= location.end <= source_length:
            raise YamlEditError(
                ErrorCode.INVARIANT_VIOLATION, "Discovered batch target span is invalid."
            )
        if location.start < previous_end:
            raise YamlEditError(
                ErrorCode.INVARIANT_VIOLATION, "Discovered batch target spans overlap."
            )
        previous_end = location.end


def _batch_expected(parsed: Parsed, resolved: tuple[tuple[Candidate, Scalar], ...]) -> Tree:
    expected: Tree = parsed.data
    for candidate, value in resolved:
        expected = _with_value(expected, candidate.path, value)
    return expected


def _batch_source(
    parsed: Parsed, source: str, resolved: tuple[tuple[Candidate, Scalar], ...]
) -> str:
    result = source
    for candidate, value in reversed(resolved):
        if _same(candidate.current_value, value):
            continue
        original = parsed.nodes.get(candidate.path)
        if not isinstance(original, ScalarNode):
            raise YamlEditError(
                ErrorCode.INVARIANT_VIOLATION, "Discovered batch target is not a scalar."
            )
        start, end = candidate.location.start, candidate.location.end
        if not 0 <= start <= end <= len(source):
            raise YamlEditError(
                ErrorCode.INVARIANT_VIOLATION, "Discovered batch target span is invalid."
            )
        result = result[:start] + _replacement(value, original) + result[end:]
    return result


def _preview_batch(parsed: Parsed, source: str, proposal: BatchProposal) -> BatchPreview:
    resolved = _resolve_entries(parsed, proposal, len(source))
    expected = _batch_expected(parsed, resolved)
    result = _batch_source(parsed, source, resolved)
    reparsed = parse_source(result)
    if not _same(expected, reparsed.data):
        raise YamlEditError(
            ErrorCode.INVARIANT_VIOLATION,
            "Batch edit changed more than the intended semantic values.",
        )
    return BatchPreview(
        source=result,
        diff=_unified_diff(source, result),
        changed=result != source,
        candidates=tuple(candidate for candidate, _value in resolved),
    )


def preview_batch(source: str, proposal: object) -> BatchPreview:
    """Revalidate an untrusted batch artifact and return its complete review result."""
    checked = validate_batch_proposal(proposal)
    parsed = parse_source(source)
    return _preview_batch(parsed, source, checked)


def apply_batch(source: str, proposal: object) -> str:
    """Apply an atomic batch to an immutable string snapshot."""
    return preview_batch(source, proposal).source