"""Document — a dict-like wrapper around a json-five model with comment preservation.

A ``Document`` represents a loaded JSON5/JSON/JSONC file. It provides a
dict-like interface for reading and modifying data, while internally
maintaining a json-five model that preserves comments and whitespace.

Changes are written back with ``save()``, ``save_as()``, or the
functional ``dump()`` / ``dumps()`` API.
"""

from __future__ import annotations

from pathlib import Path
from typing import Any, Iterator

from json5.dumper import ModelDumper, dumps
from json5.loader import DefaultLoader, ModelLoader
from json5.model import (
    DoubleQuotedString,
    Identifier,
    JSONObject,
    JSONText,
    SingleQuotedString,
)

from ._modelizer import modelize_json


class Document:
    """A comment-preserving JSON5 document with a dict-like interface.

    Created by :func:`json5rt.load`, :func:`json5rt.loads`, or
    :func:`json5rt.Document()` (empty). Acts like a dict for reading and
    modifying data, but preserves comments and whitespace when saved.

    Example::

        import json5rt

        doc = json5rt.load("config.json5")
        doc["port"] = 9090
        del doc["features"]
        doc.save()  # comments preserved in the file
    """

    def __init__(
        self,
        model: JSONText | None = None,
        source_path: Path | None = None,
    ) -> None:
        """Initialize a Document.

        Args:
            model: A json-five JSONText model (from ModelLoader). If None,
                an empty model is created.
            source_path: The file path this document was loaded from
                (used by ``save()`` to write back to the same file).
        """
        if model is None:
            model = JSONText(value=JSONObject(trailing_comma=None))
        self._model = model
        self._source_path = source_path

    # ── Dict-like read interface ─────────────────────────────────────────

    def __getitem__(self, key: str) -> Any:
        """Get a value by key. Returns plain Python objects (not model nodes)."""
        json_obj = self._model.value
        for i, k in enumerate(json_obj.keys):
            if self._key_matches(k, key):
                # Load the value node through DefaultLoader to get plain Python
                value_node = json_obj.values[i]
                loader = DefaultLoader()
                return loader.load(value_node)
        raise KeyError(key)

    def __contains__(self, key: str) -> bool:
        """Check if a key exists in the document."""
        json_obj = self._model.value
        for k in json_obj.keys:
            if self._key_matches(k, key):
                return True
        return False

    def __iter__(self) -> Iterator[str]:
        """Iterate over keys."""
        return iter(self.keys())

    def __len__(self) -> int:
        """Number of top-level keys."""
        return len(self._model.value.keys)

    def keys(self) -> list[str]:
        """Return all keys as strings."""
        result: list[str] = []
        for k in self._model.value.keys:
            result.append(self._key_name(k))
        return result

    def values(self) -> list[Any]:
        """Return all values as plain Python objects."""
        loader = DefaultLoader()
        return [loader.load(v) for v in self._model.value.values]

    def items(self) -> list[tuple[str, Any]]:
        """Return all (key, value) pairs."""
        loader = DefaultLoader()
        return [
            (self._key_name(k), loader.load(v))
            for k, v in zip(self._model.value.keys, self._model.value.values)
        ]

    def get(self, key: str, default: Any = None) -> Any:
        """Get a value by key, returning *default* if not found."""
        try:
            return self[key]
        except KeyError:
            return default

    # ── Dict-like write interface ────────────────────────────────────────

    def __setitem__(self, key: str, value: Any) -> None:
        """Set or update a key's value. Uses double-quoted strings for output."""
        json_obj = self._model.value
        new_value_model = modelize_json(value)

        for i, k in enumerate(json_obj.keys):
            if self._key_matches(k, key):
                # Replace only the value — keep the existing key node so
                # comments and whitespace before the key survive.
                json_obj.values[i] = new_value_model
                return

        # Key doesn't exist — append it
        new_key = DoubleQuotedString(key, f'"{key}"')
        new_key.wsc_before = ["\n  "]
        json_obj.keys.append(new_key)
        json_obj.values.append(new_value_model)

    def __delitem__(self, key: str) -> None:
        """Delete a key. Comments on other keys are preserved."""
        json_obj = self._model.value
        for i, k in enumerate(json_obj.keys):
            if self._key_matches(k, key):
                del json_obj.keys[i]
                del json_obj.values[i]
                return
        raise KeyError(key)

    def update(self, other: dict[str, Any]) -> None:
        """Update multiple keys at once."""
        for key, value in other.items():
            self[key] = value

    # ── Serialization ────────────────────────────────────────────────────

    def to_dict(self) -> dict[str, Any]:
        """Convert to a plain Python dict (comments are lost)."""
        loader = DefaultLoader()
        result = loader.load(self._model)
        return result if isinstance(result, dict) else {}

    def to_json(self, indent: int = 2) -> str:
        """Serialize to a strict JSON string (comments are lost, double-quoted)."""
        import json as stdlib_json
        return stdlib_json.dumps(self.to_dict(), indent=indent)

    def dumps(self) -> str:
        """Serialize to a JSON5/JSONC string with comments preserved."""
        return dumps(self._model, dumper=ModelDumper())

    # ── File operations ──────────────────────────────────────────────────

    def save(self, path: str | Path | None = None) -> Path:
        """Save the document to a file with comments preserved.

        If *path* is None, saves back to the file this document was loaded
        from. If the document was not loaded from a file, a ValueError is
        raised.

        Args:
            path: File path to save to. If None, uses the source path.

        Returns:
            The path the document was saved to.
        """
        target = Path(path) if path is not None else self._source_path
        if target is None:
            raise ValueError(
                "No source path — provide a path argument or use save_as()"
            )
        target.parent.mkdir(parents=True, exist_ok=True)
        target.write_text(self.dumps(), encoding="utf-8")
        return target

    def save_as(self, path: str | Path) -> Path:
        """Save the document to a new file with comments preserved.

        Unlike ``save()``, this always requires a path and updates the
        internal source path so subsequent ``save()`` calls write to the
        new location.

        Args:
            path: File path to save to.

        Returns:
            The path the document was saved to.
        """
        target = Path(path)
        result = self.save(target)
        self._source_path = target
        return result

    # ── Properties ───────────────────────────────────────────────────────

    @property
    def source_path(self) -> Path | None:
        """The file path this document was loaded from, or None."""
        return self._source_path

    @property
    def has_comments(self) -> bool:
        """Check if the document contains any comments."""
        text = self.dumps()
        return "//" in text or "/*" in text

    # ── Internal helpers ─────────────────────────────────────────────────

    @staticmethod
    def _key_matches(key_node: Any, key_name: str) -> bool:
        """Check if a model key node matches a string key name."""
        if isinstance(key_node, Identifier):
            return key_node.name == key_name
        if isinstance(key_node, DoubleQuotedString):
            return key_node.characters == key_name
        if isinstance(key_node, SingleQuotedString):
            return key_node.characters == key_name
        return False

    @staticmethod
    def _key_name(key_node: Any) -> str:
        """Extract the string name from a model key node."""
        if isinstance(key_node, Identifier):
            return key_node.name
        if isinstance(key_node, (DoubleQuotedString, SingleQuotedString)):
            return key_node.characters
        return str(key_node)

    # ── Dunder methods ───────────────────────────────────────────────────

    def __repr__(self) -> str:
        if self._source_path:
            return f"Document(source={self._source_path!s}, keys={list(self.keys())})"
        return f"Document(keys={list(self.keys())})"

    def __eq__(self, other: object) -> bool:
        if isinstance(other, Document):
            return self.to_dict() == other.to_dict()
        if isinstance(other, dict):
            return self.to_dict() == other
        return NotImplemented
