from __future__ import annotations

import hashlib
from dataclasses import dataclass, field
from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple

ROOT_TOKEN = "q"


def encode_path(indices: Sequence[int], root_token: str = ROOT_TOKEN) -> str:
    """
    Convert a tuple of branch indices into the canonical ``q/0/1`` path string.
    """

    if not indices:
        return root_token
    return f"{root_token}/" + "/".join(str(idx) for idx in indices)


def decode_path(path_str: str, root_token: str = ROOT_TOKEN) -> Tuple[int, ...]:
    """
    Parse a canonical path string back into indices.
    """

    if not path_str or path_str == root_token:
        return ()
    parts = path_str.split("/")
    filtered = [int(p) for p in parts if p and p != root_token]
    return tuple(filtered)


def ancestry_from_indices(indices: Sequence[int], root_token: str = ROOT_TOKEN) -> Tuple[str, ...]:
    """Return all prefix paths (including the root)."""

    ancestry = [root_token]
    for depth in range(len(indices)):
        ancestry.append(encode_path(indices[: depth + 1], root_token))
    return tuple(ancestry)


def path_is_prefix(prefix: str, path: str) -> bool:
    """
    Check whether ``prefix`` is equal to or a strict prefix of ``path``.
    """

    return path.startswith(prefix) and (path == prefix or path[len(prefix)] == "/")


@dataclass(frozen=True)
class TrajectoryKey:
    """
    Identifier for a trajectory within a LiTS search instance.

    Tree-search modules (e.g., :mod:`lits.agents.tree_search.mcts`) pass instances to
    :class:`~lits.memory.manager.LiTSMemoryManager` so that inherited memories can be
    resolved without leaking search-tree internals into the memory subsystem.
    """

    search_id: str
    indices: Tuple[int, ...] = field(default_factory=tuple)
    root_token: str = ROOT_TOKEN

    def __post_init__(self):
        for idx in self.indices:
            if not isinstance(idx, int):
                raise TypeError("Trajectory indices must be integers.")

    @property
    def depth(self) -> int:
        return len(self.indices)

    @property
    def path_str(self) -> str:
        return encode_path(self.indices, self.root_token)

    @property
    def ancestry_paths(self) -> Tuple[str, ...]:
        return ancestry_from_indices(self.indices, self.root_token)

    def parent(self) -> "TrajectoryKey":
        if not self.indices:
            return self
        return TrajectoryKey(self.search_id, self.indices[:-1], root_token=self.root_token)

    def child(self, branch: int) -> "TrajectoryKey":
        return TrajectoryKey(
            self.search_id,
            self.indices + (branch,),
            root_token=self.root_token,
        )

    def extends(self, prefix_path: str) -> bool:
        return path_is_prefix(prefix_path, self.path_str)

    @classmethod
    def from_path(cls, search_id: str, path: str, root_token: str = ROOT_TOKEN) -> "TrajectoryKey":
        return cls(search_id=search_id, indices=decode_path(path, root_token), root_token=root_token)

    def to_dict(self) -> Dict[str, Any]:
        """Serialize to a JSON-safe dict."""
        return {
            "search_id": self.search_id,
            "indices": list(self.indices),
            "root_token": self.root_token,
        }

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "TrajectoryKey":
        """Deserialize from a dict."""
        return cls(
            search_id=data["search_id"],
            indices=tuple(data.get("indices", ())),
            root_token=data.get("root_token", ROOT_TOKEN),
        )


@dataclass
class MemoryUnit:
    """
    Container for a stored memory fact.

    Attributes:
        id: Identifier returned by mem0/Qdrant.
        text: Natural language fact.
        search_id: Search instance identifier (maps to ``user_id`` in mem0).
        origin_path: Path string corresponding to the node that produced the memory.
        depth: Cached depth derived from ``origin_path``.
        ancestry_paths: Prefix paths stored in metadata.
        metadata: Original payload.
        created_at: Timestamp string generated by mem0.
        content_hash: Hash used for deduplication and similarity comparisons.
    """

    id: str
    text: str
    search_id: str
    origin_path: str
    depth: int
    ancestry_paths: Tuple[str, ...]
    metadata: dict = field(default_factory=dict)
    created_at: Optional[str] = None
    content_hash: Optional[str] = None

    def signature(self) -> str:
        """
        Unique key for set operations.  Falls back to a deterministic digest when mem0
        does not supply a hash (useful for ``LocalMemoryBackend`` and unit tests).
        """

        if self.content_hash:
            return self.content_hash
        digest = hashlib.md5(self.text.encode("utf-8")).hexdigest()
        return f"{digest}:{self.origin_path}"

    def inherited_by(self, trajectory_path: str) -> bool:
        """
        Return ``True`` if this memory belongs to an ancestor of ``trajectory_path``.
        """

        return path_is_prefix(self.origin_path, trajectory_path)


@dataclass
class TrajectorySimilarity:
    """
    Result returned by the trajectory search stage.

    Attributes:
        trajectory_path: Path string for the retrieved trajectory.
        score: Overlap score (``|norm(M(t)) ∩ norm(M(~t))| / |norm(M(t))|``).
        missing_units: Units present in the retrieved trajectory but absent from the
            current trajectory – the output of ``Sel`` in the specification.
        overlapping_units: Units responsible for the similarity hit; provided to help
            evaluation modules reason about provenance.
    """

    trajectory_path: str
    score: float
    missing_units: Tuple[MemoryUnit, ...]
    overlapping_units: Tuple[MemoryUnit, ...]

    def to_prompt_section(self) -> str:
        """
        Render retrieved memories as a compact prompt fragment suitable for the policy
        LLM.  The caller typically concatenates multiple sections when building the
        expansion prompt inside tree search policies.
        """

        header = f"Trajectory {self.trajectory_path} (score={self.score:.2f})"
        details = "\n".join(f"- {unit.text}" for unit in self.missing_units)
        return f"{header}\n{details}".strip()
