"""
RBSP Index Migration.

Handles upgrading old index formats to the current version without
requiring a full re-index.  Each ``Migration`` transforms index data
from one version to the next; ``MigrationManager`` chains them together
to reach the current version.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Any, Dict, List, Optional


# ---------------------------------------------------------------------------
# Index version
# ---------------------------------------------------------------------------

@dataclass
class IndexVersion:
    """Semantic version identifying an index format."""

    major: int
    minor: int
    patch: int

    def __str__(self) -> str:
        return f"{self.major}.{self.minor}.{self.patch}"

    def __eq__(self, other: object) -> bool:
        if not isinstance(other, IndexVersion):
            return NotImplemented
        return (self.major, self.minor, self.patch) == (other.major, other.minor, other.patch)

    def __lt__(self, other: object) -> bool:
        if not isinstance(other, IndexVersion):
            return NotImplemented
        return (self.major, self.minor, self.patch) < (other.major, other.minor, other.patch)

    def __le__(self, other: object) -> bool:
        if not isinstance(other, IndexVersion):
            return NotImplemented
        return self == other or self < other

    def __gt__(self, other: object) -> bool:
        if not isinstance(other, IndexVersion):
            return NotImplemented
        return (self.major, self.minor, self.patch) > (other.major, other.minor, other.patch)

    def __ge__(self, other: object) -> bool:
        if not isinstance(other, IndexVersion):
            return NotImplemented
        return self == other or self > other

    def __hash__(self) -> int:
        return hash((self.major, self.minor, self.patch))

    @staticmethod
    def from_string(s: str) -> IndexVersion:
        """Parse a ``"major.minor.patch"`` string into an :class:`IndexVersion`."""
        parts = s.strip().split(".")
        if len(parts) != 3:
            raise ValueError(f"Invalid version string: {s!r}")
        try:
            return IndexVersion(int(parts[0]), int(parts[1]), int(parts[2]))
        except ValueError:
            raise ValueError(f"Invalid version string: {s!r}")


# Current index format version.
CURRENT_VERSION = IndexVersion(3, 0, 0)


# ---------------------------------------------------------------------------
# Migration base class
# ---------------------------------------------------------------------------

class Migration:
    """Base class for a single index-format migration step.

    Subclasses must set :attr:`from_version`, :attr:`to_version`,
    :attr:`description` and implement :meth:`migrate`.
    """

    from_version: IndexVersion
    to_version: IndexVersion
    description: str = ""

    def migrate(self, index_data: dict) -> dict:
        """Transform *index_data* from :attr:`from_version` to :attr:`to_version`.

        Returns the transformed data dict (may mutate in place).
        """
        raise NotImplementedError

    def validate(self, index_data: dict) -> list[str]:
        """Return a list of validation errors for *index_data*.

        An empty list means the data is valid for ``to_version``.
        """
        return []


# ---------------------------------------------------------------------------
# Concrete migrations
# ---------------------------------------------------------------------------

class V1ToV2Migration(Migration):
    """V1 -> V2: add ``token_set`` field to each entry."""

    from_version = IndexVersion(1, 0, 0)
    to_version = IndexVersion(2, 0, 0)
    description = "Add token_set field to index entries"

    def migrate(self, index_data: dict) -> dict:
        entries = index_data.get("entries", {})
        for key, entry in entries.items():
            if "token_set" not in entry:
                tokens = entry.get("tokens", [])
                entry["token_set"] = list(set(tokens))
        index_data["version"] = str(self.to_version)
        return index_data

    def validate(self, index_data: dict) -> list[str]:
        errors: list[str] = []
        entries = index_data.get("entries", {})
        for key, entry in entries.items():
            if "token_set" not in entry:
                errors.append(f"Entry {key!r} missing 'token_set'")
        return errors


class V2ToV3Migration(Migration):
    """V2 -> V3: add ``term_freqs`` and ``doc_len`` fields to each entry."""

    from_version = IndexVersion(2, 0, 0)
    to_version = IndexVersion(3, 0, 0)
    description = "Add term_freqs and doc_len fields to index entries"

    def migrate(self, index_data: dict) -> dict:
        entries = index_data.get("entries", {})
        for key, entry in entries.items():
            if "term_freqs" not in entry:
                tokens = entry.get("tokens", [])
                freqs: Dict[str, int] = {}
                for t in tokens:
                    freqs[t] = freqs.get(t, 0) + 1
                entry["term_freqs"] = freqs
            if "doc_len" not in entry:
                entry["doc_len"] = len(entry.get("tokens", []))
        index_data["version"] = str(self.to_version)
        return index_data

    def validate(self, index_data: dict) -> list[str]:
        errors: list[str] = []
        entries = index_data.get("entries", {})
        for key, entry in entries.items():
            if "term_freqs" not in entry:
                errors.append(f"Entry {key!r} missing 'term_freqs'")
            if "doc_len" not in entry:
                errors.append(f"Entry {key!r} missing 'doc_len'")
        return errors


# ---------------------------------------------------------------------------
# Migration manager
# ---------------------------------------------------------------------------

class MigrationManager:
    """Registry of migrations that can upgrade index data to the current version.

    Finds the shortest chain of migrations from an old version to
    :attr:`current_version` and applies them in order.
    """

    def __init__(self, current_version: IndexVersion) -> None:
        self.current_version = current_version
        self._migrations: List[Migration] = []

    def register(self, migration: Migration) -> None:
        """Register a migration step."""
        self._migrations.append(migration)

    def needs_migration(self, data_version: IndexVersion) -> bool:
        """Return ``True`` if *data_version* is older than the current version."""
        return data_version < self.current_version

    def get_migration_path(self, from_version: IndexVersion) -> list[Migration]:
        """Find an ordered chain of migrations from *from_version* to *current_version*.

        Raises :class:`ValueError` if no path exists.
        """
        if from_version == self.current_version:
            return []

        # Build adjacency: from_version -> list of migrations
        by_source: Dict[IndexVersion, List[Migration]] = {}
        for m in self._migrations:
            by_source.setdefault(m.from_version, []).append(m)

        # BFS to find shortest path
        visited: Dict[IndexVersion, Optional[Migration]] = {from_version: None}
        queue: list[IndexVersion] = [from_version]

        while queue:
            current = queue.pop(0)
            if current == self.current_version:
                break
            for m in by_source.get(current, []):
                if m.to_version not in visited:
                    visited[m.to_version] = m
                    queue.append(m.to_version)
        else:
            if self.current_version not in visited:
                raise ValueError(
                    f"No migration path from {from_version} to {self.current_version}"
                )

        if self.current_version not in visited:
            raise ValueError(
                f"No migration path from {from_version} to {self.current_version}"
            )

        # Reconstruct path
        path: list[Migration] = []
        version = self.current_version
        while version != from_version:
            m = visited[version]
            assert m is not None
            path.append(m)
            version = m.from_version
        path.reverse()
        return path

    def migrate(self, index_data: dict, from_version: IndexVersion) -> dict:
        """Apply all necessary migrations to bring *index_data* to the current version.

        Returns the migrated data dict.  Raises :class:`ValueError` if
        the migration path cannot be found.
        """
        if not self.needs_migration(from_version):
            return index_data

        path = self.get_migration_path(from_version)
        for migration in path:
            index_data = migration.migrate(index_data)
        return index_data


# ---------------------------------------------------------------------------
# Default manager with pre-registered migrations
# ---------------------------------------------------------------------------

def create_default_manager() -> MigrationManager:
    """Create a :class:`MigrationManager` pre-loaded with all known migrations."""
    manager = MigrationManager(CURRENT_VERSION)
    manager.register(V1ToV2Migration())
    manager.register(V2ToV3Migration())
    return manager
