"""Safe JSON Patch (RFC 6902) operations for LLM-generated structured updates.

This module intentionally supports a focused subset of RFC 6902 operations:
``add``, ``replace``, and ``remove``.  Unsupported operations (``move``,
``copy``, ``test``) raise :class:`PatchError`.

For LLM-forgiving flows, ``add`` paths auto-create missing intermediate
containers:

* missing object segments create ``dict`` containers
* missing list segments can append a placeholder container when the traversal
  index equals ``len(list)``
* list traversal beyond ``len(list)`` is rejected

Key features:

* **Policy enforcement** -- configurable limits on allowed operations, patch
  count, and path depth.
* **RFC 6901 JSON Pointer** resolution with proper ``~0``/``~1`` escaping.
* **String-concat edge case** -- when a patch targets ``/-`` on a *string*
  field the value is concatenated rather than raising a type error.  This
  matches the behaviour of the *trustcall* library and is the most common
  failure mode when an LLM tries to "append" to a text field.
* **Deep-copy semantics** -- the input document is never mutated.
* **Pydantic integration** -- :func:`apply_patch_and_validate` applies patches
  and validates the result against a Pydantic model / ``TypeAdapter``.
"""

from __future__ import annotations

import copy
import json
from collections.abc import Sequence
from dataclasses import dataclass
from typing import Any, Literal

from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, model_validator

from .api import ParseResult
from .json_pointer import parse_json_pointer

# ---------------------------------------------------------------------------
# Exceptions
# ---------------------------------------------------------------------------


class PatchError(Exception):
    """Base exception for all patch-related errors."""


class PolicyViolationError(PatchError):
    """Raised when a patch set violates the active :class:`PatchPolicy`."""


# ---------------------------------------------------------------------------
# Core models
# ---------------------------------------------------------------------------


class JsonPatchOp(BaseModel):
    """A single RFC 6902 JSON Patch operation.

    Only ``add``, ``replace``, and ``remove`` are supported.  ``value`` is
    required for ``add`` and ``replace`` and ignored for ``remove``.
    """

    model_config = ConfigDict(extra="forbid")

    op: Literal["add", "replace", "remove"]
    path: str
    value: Any | None = None

    @model_validator(mode="before")
    @classmethod
    def _require_value_for_add_replace(cls, data: Any) -> Any:
        if isinstance(data, dict):
            op = data.get("op")
            if op in ("add", "replace") and "value" not in data:
                raise ValueError(f"'{op}' operation requires a 'value' field")
        return data


class PatchDoc(BaseModel):
    """LLM-facing tool schema for submitting a batch of patches.

    ``json_doc_id``
        Identifier of the document being patched (use ``"doc"`` for
        single-document workflows).
    ``planned_edits``
        Short bullet-point plan produced by the LLM for debugging purposes;
        the plan is *not* executed.
    ``patches``
        Ordered list of :class:`JsonPatchOp` operations.
    """

    model_config = ConfigDict(extra="forbid")

    json_doc_id: str = Field(description="ID of the document to patch (use 'doc' if single).")
    planned_edits: str = Field(description="Short plan of edits (for debugging; not executed).")
    patches: list[JsonPatchOp] = Field(description="RFC 6902 add/replace/remove operations.")


@dataclass(slots=True)
class PatchPolicy:
    """Safety rails for patch application.

    Attributes
    ----------
    allow_remove : bool
        Whether ``remove`` operations are permitted.  Default ``False``
        to prevent accidental data loss.
    max_ops : int
        Maximum number of operations in a single patch set.
    max_path_depth : int
        Maximum number of path segments (``/``-separated tokens) allowed
        in a single JSON Pointer.
    allow_append : bool
        Whether ``/-`` (array append) is permitted.
    """

    allow_remove: bool = False
    max_ops: int = 50
    max_path_depth: int = 32
    allow_append: bool = True


# ---------------------------------------------------------------------------
# JSON Pointer helpers (RFC 6901)
# ---------------------------------------------------------------------------


def _parse_pointer(path: str) -> list[str]:
    """Split a JSON Pointer string into unescaped reference tokens.

    Returns an empty list for the root pointer ``""``.  Raises
    :class:`PatchError` if the pointer does not start with ``/``.
    """
    try:
        return parse_json_pointer(path)
    except ValueError as exc:
        raise PatchError(str(exc)) from exc


def _resolve_parent(
    doc: Any, tokens: list[str], *, max_depth: int, create_missing: bool = False
) -> tuple[Any, str]:
    """Walk *doc* following *tokens* and return ``(parent, last_token)``.

    Parameters
    ----------
    doc
        The mutable document (already deep-copied).
    tokens
        Reference tokens parsed from the JSON Pointer.
    max_depth
        Maximum allowed number of tokens (enforced by policy).
    create_missing
        If ``True``, create intermediate ``dict`` containers for missing
        keys.  Used by the ``add`` operation.

    Raises
    ------
    PatchError
        If a token cannot be resolved (missing key, bad array index, etc.).
    PolicyViolationError
        If the path exceeds *max_depth*.
    """
    if len(tokens) > max_depth:
        raise PolicyViolationError(
            f"Path depth {len(tokens)} exceeds policy max_path_depth={max_depth}"
        )
    if not tokens:
        raise PatchError("Cannot operate on the document root (empty path)")

    current = doc
    for i, tok in enumerate(tokens[:-1]):
        if isinstance(current, dict):
            if tok not in current:
                if create_missing:
                    # Look ahead: if the next token is a digit or "-",
                    # the path implies list indexing, so create a list.
                    next_tok = tokens[i + 1]
                    if next_tok.isdigit() or next_tok == "-":
                        current[tok] = []
                    else:
                        current[tok] = {}
                    current = current[tok]
                    continue
                raise PatchError(f"Key {tok!r} not found while resolving path")
            current = current[tok]
        elif isinstance(current, list):
            next_tok = tokens[i + 1]
            if tok == "-":
                if not create_missing:
                    raise PatchError("The '-' index is only valid for 'add' operations")
                idx = len(current)
            else:
                try:
                    idx = int(tok)
                except ValueError:
                    raise PatchError(f"Invalid array index: {tok!r}") from None
                if idx < 0:
                    raise PatchError(
                        f"Invalid array index: negative indices are not allowed ({idx})"
                    )

            if idx < len(current):
                current = current[idx]
                continue

            if create_missing and idx == len(current):
                current.append([] if (next_tok.isdigit() or next_tok == "-") else {})
                current = current[idx]
                continue

            raise PatchError(f"Array index {idx} out of bounds (length {len(current)})")
        else:
            raise PatchError(f"Cannot traverse into {type(current).__name__} with token {tok!r}")
    return current, tokens[-1]


def _list_index(token: str, lst: list[Any], *, allow_dash: bool, allow_append: bool = False) -> int:
    """Convert a JSON Pointer token to a list index.

    ``-`` is only allowed when *allow_dash* is ``True`` (used for ``add``).
    *allow_append* permits ``idx == len(lst)`` so that RFC 6902 ``add`` can
    insert at the end of a list using an explicit numeric index.
    """
    if token == "-":
        if allow_dash:
            return len(lst)
        raise PatchError("The '-' index is only valid for 'add' operations")
    try:
        idx = int(token)
    except ValueError:
        raise PatchError(f"Invalid array index: {token!r}") from None
    if idx < 0:
        raise PatchError(f"Invalid array index: negative indices are not allowed ({idx})")
    if idx > len(lst) or (idx == len(lst) and not allow_append):
        raise PatchError(f"Array index {idx} out of bounds (length {len(lst)})")
    return idx


# ---------------------------------------------------------------------------
# Individual operation handlers
# ---------------------------------------------------------------------------


def _apply_add(
    doc: Any, tokens: list[str], value: Any, *, max_depth: int, allow_append: bool = True
) -> None:
    """Apply an ``add`` operation (mutates *doc* in-place)."""
    parent, last = _resolve_parent(doc, tokens, max_depth=max_depth, create_missing=True)
    if isinstance(parent, dict):
        parent[last] = value
    elif isinstance(parent, list):
        if last == "-":
            parent.append(value)
        else:
            idx = _list_index(last, parent, allow_dash=True, allow_append=allow_append)
            if idx == len(parent):
                parent.append(value)
            else:
                parent.insert(idx, value)
    elif isinstance(parent, str):
        # Trustcall edge case: ``/-`` on a string field means concatenation.
        if last == "-":
            _string_concat_in_parent(doc, tokens, value)
        else:
            raise PatchError(f"Cannot add into a string value at {'/'.join(tokens[:-1])!r}")
    else:
        raise PatchError(f"Cannot add into {type(parent).__name__} at {'/'.join(tokens[:-1])!r}")


def _apply_replace(doc: Any, tokens: list[str], value: Any, *, max_depth: int) -> None:
    """Apply a ``replace`` operation (mutates *doc* in-place).

    The target path **must** already exist; ``replace`` never creates
    intermediate containers.
    """
    parent, last = _resolve_parent(doc, tokens, max_depth=max_depth, create_missing=False)

    # Handle string-concat edge case for ``/-`` on a string field.
    if isinstance(parent, str) and last == "-":
        _string_concat_in_parent(doc, tokens, value)
        return

    if isinstance(parent, dict):
        if last not in parent:
            raise PatchError(f"Path does not exist (cannot replace): key {last!r}")
        parent[last] = value
    elif isinstance(parent, list):
        idx = _list_index(last, parent, allow_dash=False)
        parent[idx] = value
    else:
        raise PatchError(f"Cannot replace in {type(parent).__name__}")


def _apply_remove(doc: Any, tokens: list[str], *, max_depth: int) -> None:
    """Apply a ``remove`` operation (mutates *doc* in-place)."""
    parent, last = _resolve_parent(doc, tokens, max_depth=max_depth, create_missing=False)
    if isinstance(parent, dict):
        if last not in parent:
            raise PatchError(f"Path does not exist (cannot remove): key {last!r}")
        del parent[last]
    elif isinstance(parent, list):
        idx = _list_index(last, parent, allow_dash=False)
        del parent[idx]
    else:
        raise PatchError(f"Cannot remove from {type(parent).__name__}")


# ---------------------------------------------------------------------------
# String-concat edge case  (trustcall compatibility)
# ---------------------------------------------------------------------------


def _string_concat_in_parent(doc: Any, tokens: list[str], value: Any) -> None:
    """Handle ``/-`` applied to a string field by concatenating *value*.

    Walks the document to the *grandparent* and replaces the string at the
    parent key with the concatenated result.

    This is a practical edge case documented by the *trustcall* library: LLMs
    sometimes emit ``{"op": "add", "path": "/bio/-", "value": " extra text"}``
    when ``bio`` is a ``str``, intending to append text.  Rather than raising
    a type error we interpret this as string concatenation.
    """
    # tokens looks like [..., parent_key, "-"]
    parent_tokens = tokens[:-1]
    if not parent_tokens:
        raise PatchError("Cannot apply string-concat on the document root")
    grandparent_tokens = parent_tokens[:-1]
    field_key = parent_tokens[-1]

    if grandparent_tokens:
        gp, _ = _resolve_parent(
            doc, parent_tokens, max_depth=len(parent_tokens), create_missing=False
        )
        container = gp
    else:
        container = doc

    if isinstance(container, dict):
        old = container.get(field_key, "")
        if not isinstance(old, str):
            raise PatchError(
                f"Expected string for concat at {'/'.join(parent_tokens)!r}, "
                f"got {type(old).__name__}"
            )
        container[field_key] = old + str(value)
    elif isinstance(container, list):
        try:
            idx = int(field_key)
        except ValueError:
            raise PatchError(f"Invalid array index for string concat: {field_key!r}") from None
        if idx < 0 or idx >= len(container):
            raise PatchError(
                f"Array index {idx} out of bounds for string concat (length {len(container)})"
            )
        old = container[idx]
        if not isinstance(old, str):
            raise PatchError(
                f"Expected string for concat at {'/'.join(parent_tokens)!r}, "
                f"got {type(old).__name__}"
            )
        container[idx] = old + str(value)
    else:
        raise PatchError(f"Cannot apply string-concat: grandparent is {type(container).__name__}")


# ---------------------------------------------------------------------------
# Policy validation
# ---------------------------------------------------------------------------


def _validate_policy(patches: Sequence[JsonPatchOp], policy: PatchPolicy) -> None:
    """Check that *patches* conform to *policy*.

    Raises :class:`PolicyViolationError` on the first violation found.
    """
    if len(patches) > policy.max_ops:
        raise PolicyViolationError(
            f"Patch set contains {len(patches)} operations, "
            f"but policy allows at most {policy.max_ops}"
        )
    for patch in patches:
        # Disallow ``remove`` unless policy permits it.
        if patch.op == "remove" and not policy.allow_remove:
            raise PolicyViolationError("Remove operations are not allowed by the current policy")
        # Disallow ``/-`` append unless policy permits it.
        # Check for "-" in any token position, not just the final segment,
        # to prevent bypass via intermediate "/-/" paths.
        tokens = _parse_pointer(patch.path)
        if not policy.allow_append and any(t == "-" for t in tokens):
            raise PolicyViolationError(
                "Append (/-) operations are not allowed by the current policy"
            )
        # Check path depth (tokens already parsed above).
        if len(tokens) > policy.max_path_depth:
            raise PolicyViolationError(
                f"Path {patch.path!r} has depth {len(tokens)}, "
                f"exceeding policy max_path_depth={policy.max_path_depth}"
            )


# ---------------------------------------------------------------------------
# Public API: apply_patch
# ---------------------------------------------------------------------------


def apply_patch(
    doc: Any,
    patches: Sequence[JsonPatchOp],
    *,
    policy: PatchPolicy | None = None,
) -> Any:
    """Apply a sequence of JSON Patch operations to *doc* and return the result.

    The input *doc* is **never mutated**; a deep copy is made before applying
    any operations.

    Parameters
    ----------
    doc
        The document to patch.
    patches
        Ordered list of :class:`JsonPatchOp` operations.
    policy
        Safety policy to enforce.  See :class:`PatchPolicy`.

    Returns
    -------
    Any
        The patched document.

    Raises
    ------
    PatchError
        If any operation fails (missing path, bad index, etc.).
    PolicyViolationError
        If the patch set violates the active policy.
    """
    policy = policy or PatchPolicy()
    _validate_policy(patches, policy)
    result = copy.deepcopy(doc)
    for patch in patches:
        tokens = _parse_pointer(patch.path)
        if patch.op == "add":
            _apply_add(
                result,
                tokens,
                patch.value,
                max_depth=policy.max_path_depth,
                allow_append=policy.allow_append,
            )
        elif patch.op == "replace":
            _apply_replace(result, tokens, patch.value, max_depth=policy.max_path_depth)
        elif patch.op == "remove":
            _apply_remove(result, tokens, max_depth=policy.max_path_depth)
        else:
            raise PatchError(f"Unsupported operation: {patch.op!r}")
    return result


# ---------------------------------------------------------------------------
# Public API: apply_patch_and_validate
# ---------------------------------------------------------------------------


def apply_patch_and_validate[T](
    doc: Any | BaseModel,
    patches: Sequence[JsonPatchOp],
    target: type[T] | TypeAdapter[T],
    *,
    policy: PatchPolicy | None = None,
) -> ParseResult[T]:
    """Apply patches then validate the result against a Pydantic type.

    If *doc* is a :class:`~pydantic.BaseModel` it is first converted to a
    ``dict`` via :meth:`~pydantic.BaseModel.model_dump`.

    Returns a :class:`~parsantic.api.ParseResult` with the validated
    value, empty flags, and a score of ``0``.

    Raises
    ------
    PatchError
        If any patch operation fails.
    PolicyViolationError
        If the patches violate the active policy.
    pydantic.ValidationError
        If the patched document does not conform to *target*.
    """
    if isinstance(doc, BaseModel):
        doc_dict: Any = doc.model_dump(mode="json")
    else:
        doc_dict = doc

    patched = apply_patch(doc_dict, patches, policy=policy)

    adapter: TypeAdapter[T] = target if isinstance(target, TypeAdapter) else TypeAdapter(target)
    validated = adapter.validate_python(patched)
    return ParseResult(value=validated, flags=(), score=0)


# ---------------------------------------------------------------------------
# Public API: normalize_patches
# ---------------------------------------------------------------------------


def normalize_patches(patches: Any) -> list[JsonPatchOp]:
    """Coerce various patch representations into a list of :class:`JsonPatchOp`.

    Handles common LLM failure modes:

    * A ``list[dict]`` of raw patch dicts.
    * A JSON **string** containing a list of patches.
    * A ``dict`` with a ``"patches"`` key wrapping the actual list.
    * A JSON string wrapping a dict with a ``"patches"`` key.
    * A single ``dict`` (treated as a one-element list).
    * Already-validated :class:`JsonPatchOp` instances (pass-through).

    Raises
    ------
    PatchError
        If the input cannot be interpreted as a list of patches.
    """
    # 1) If it's a string, try to parse it as JSON first.
    if isinstance(patches, str):
        patches = _parse_json_string(patches)

    # 2) If it's a dict with a "patches" key, unwrap.
    if isinstance(patches, dict):
        if "patches" in patches:
            patches = patches["patches"]
            # Handle string-wrapped patches: {"patches": "[{...}, {...}]"}
            if isinstance(patches, str):
                patches = _parse_json_string(patches)
        else:
            # Treat a single dict as a one-element list.
            patches = [patches]

    # 3) Must be iterable at this point.
    if not isinstance(patches, (list, tuple)):
        raise PatchError(
            f"Cannot normalize patches: expected list or dict, got {type(patches).__name__}"
        )

    result: list[JsonPatchOp] = []
    for item in patches:
        if isinstance(item, JsonPatchOp):
            result.append(item)
        elif isinstance(item, dict):
            try:
                result.append(JsonPatchOp(**item))
            except Exception as exc:
                raise PatchError(f"Cannot parse patch dict {item!r}: {exc}") from exc
        else:
            raise PatchError(
                f"Cannot normalize patch item: expected dict or JsonPatchOp, "
                f"got {type(item).__name__}"
            )
    return result


def _parse_json_string(text: str) -> Any:
    """Attempt to parse *text* as JSON, with SAP fallback.

    Tries ``json.loads`` first.  If that fails, attempts to import
    and use the SAP jsonish parser as a fallback for malformed JSON.
    """
    try:
        return json.loads(text)
    except (json.JSONDecodeError, ValueError):
        pass

    # Attempt SAP parse as a fallback for "JSON-ish" strings.
    try:
        from .jsonish import ParseOptions, parse_jsonish

        jv = parse_jsonish(
            text,
            options=ParseOptions(
                allow_markdown_json=True,
                allow_find_all_json_objects=True,
                allow_fixes=True,
                allow_as_string=False,
            ),
            is_done=True,
        )
        # Prefer candidates that validate as RFC6902 patch arrays.
        if jv.candidates:
            patch_adapter = TypeAdapter(list[JsonPatchOp])
            best_candidate: Any = None
            best_size = -1
            for cand in jv.candidates:
                candidate_value: Any = cand.value
                if isinstance(candidate_value, dict) and "patches" in candidate_value:
                    candidate_value = candidate_value["patches"]
                if isinstance(candidate_value, dict):
                    candidate_value = [candidate_value]
                if isinstance(candidate_value, list):
                    try:
                        patch_adapter.validate_python(candidate_value)
                        return candidate_value
                    except Exception:
                        pass
                try:
                    size = len(json.dumps(candidate_value, ensure_ascii=False, default=str))
                except Exception:
                    size = 0
                if size > best_size:
                    best_size = size
                    best_candidate = candidate_value
            if best_candidate is not None:
                return best_candidate
        return jv.value
    except Exception:
        pass

    raise PatchError(f"Cannot parse patch string as JSON: {text!r}")


# ---------------------------------------------------------------------------
# Convenience re-exports
# ---------------------------------------------------------------------------

__all__ = [
    "JsonPatchOp",
    "PatchDoc",
    "PatchError",
    "PatchPolicy",
    "PolicyViolationError",
    "apply_patch",
    "apply_patch_and_validate",
    "normalize_patches",
]
