import ast
import builtins
import json
import logging
import threading
import types
from collections.abc import Sequence
from pathlib import Path
from typing import Any, Callable, Generic, Mapping, TypeVar, cast

import pydantic

from stable_pydantic.model_graph import ModelNode

_log = logging.getLogger(__name__)

T = TypeVar("T", bound=pydantic.BaseModel)

_LOADED_SCHEMAS_LOCK = threading.RLock()  # For thread safety
LOADED_SCHEMAS: dict[Path, types.ModuleType] = {}
"""We cache the modules to avoid duplicating types."""

_LOADED_MIGRATIONS_LOCK = threading.RLock()  # For thread safety
LOADED_MIGRATIONS: dict[Path, types.ModuleType] = {}
"""We cache the modules to avoid duplicating types."""


class CurrentEntry(pydantic.BaseModel):
    "The current version is only present for convenience, to get a nice diff."

    path: Path
    source: str

    @staticmethod
    def create(path: Path, model: ModelNode) -> "CurrentEntry":
        _log.info(f"Creating current.py for {model.node.__name__}")

        source = model.clean_source_recursive()
        path.write_text(source)
        return CurrentEntry(
            path=path,
            source=source,
        )

    @staticmethod
    def open(path: Path) -> "CurrentEntry":
        return CurrentEntry(
            path=path,
            source=path.read_text(),
        )

    def assert_unchanged(self, model: ModelNode):
        new = model.clean_source_recursive()
        if self.source != new:
            raise ValueError(
                f"Current schema is not the same as live schema for model {model.node.__name__}.\n"
                + f"On disk:\n{self.source}\n\n"
                + f"Derived:\n{new}"
            )

    def update(self, model: ModelNode):
        new_source = model.clean_source_recursive()
        if self.source == new_source:
            _log.info(f"current.py for {model.node.__name__} is up to date")
            return
        else:
            _log.info(f"Updating current.py for {model.node.__name__}")
        self.source = new_source
        self.path.write_text(self.source)


class SchemaEntry(pydantic.BaseModel):
    version: int
    dir: Path
    source: str

    @staticmethod
    def create(dir: Path, version: int, model: ModelNode) -> "SchemaEntry":
        _log.info(f"Creating schema at version {version} for {model.node.__name__}")
        source = model.clean_source_recursive()
        (dir / _schema_file_name(version)).write_text(source)
        return SchemaEntry(
            version=version,
            dir=dir,
            source=source,
        )

    @staticmethod
    def open(dir: Path, version: int) -> "SchemaEntry":
        return SchemaEntry(
            version=version,
            dir=dir,
            source=(dir / _schema_file_name(version)).read_text(),
        )

    def assert_equal(self, model: ModelNode):
        new = model.clean_source_recursive()
        if self.source != new:
            raise ValueError(
                f"Schema at version {self.version} is not the same as live schema for model {model.node.__name__}.\n"
                + f"On disk:\n{self.source}\n\n"
                + f"Derived:\n{new}"
            )

    def isolated_module(self) -> types.ModuleType:
        """Create an isolated module containing the schema models."""
        file_path = self.dir / _schema_file_name(self.version)

        with _LOADED_SCHEMAS_LOCK:
            if file_path in LOADED_SCHEMAS:
                return LOADED_SCHEMAS[file_path]

            module_name = _schema_module_name(self.version)
            module = types.ModuleType(module_name)
            module.__file__ = str(file_path)

            # Execute the source code in the module's namespace
            try:
                exec(self.source, module.__dict__)
            except Exception as e:
                raise ValueError(
                    f"Error executing schema module {file_path}: {e}\n{self.source}"
                )

            LOADED_SCHEMAS[file_path] = module
            return module

    def isolated_model(self) -> ModelNode:
        module = self.isolated_module()

        models = [
            obj
            for obj in module.__dict__.values()  # dicts are ordered in Python 3.7+
            if isinstance(obj, type)
            and issubclass(obj, pydantic.BaseModel)
            and obj is not pydantic.BaseModel
        ]
        assert len(models) > 0, "No models found in schema module."
        # The last one is the root (dicts are ordered in Python 3.7+)
        model = models[-1]
        return ModelNode.new(model, module=module)

    def json_schema(self) -> dict[str, Any]:
        return self.isolated_model().node.model_json_schema()


class MigrationEntry(pydantic.BaseModel):
    from_version: int
    to_version: int
    path: Path
    source: str

    @staticmethod
    def create(
        from_version: int, to_version: int, path: Path, source: str
    ) -> "MigrationEntry":
        _validate_upgrade_module(source, path)
        path.write_text(source)
        return MigrationEntry(
            from_version=from_version,
            to_version=to_version,
            path=path,
            source=source,
        )

    @staticmethod
    def open(from_version: int, to_version: int, path: Path) -> "MigrationEntry":
        entry = MigrationEntry(
            from_version=from_version,
            to_version=to_version,
            path=path,
            source=path.read_text(),
        )
        _validate_upgrade_module(entry.source, entry.path)
        return entry

    def isolated_module(
        self,
        from_schema: SchemaEntry,
        to_schema: SchemaEntry,
    ) -> types.ModuleType:
        """We pass a callback because the module needs the sibling schema modules to be registered while it runs."""
        with _LOADED_MIGRATIONS_LOCK:
            if self.path in LOADED_MIGRATIONS:
                return LOADED_MIGRATIONS[self.path]

            from_module = from_schema.isolated_module()
            to_module = to_schema.isolated_module()

            # Virtual modules to intercept
            virtual_modules = {
                _schema_module_name(from_schema.version): from_module,
                _schema_module_name(to_schema.version): to_module,
            }

            original_import = builtins.__import__

            def custom_import(
                name: str,
                globals: Mapping[str, object] | None = None,
                locals: Mapping[str, object] | None = None,
                fromlist: Sequence[str] = (),
                level: int = 0,
            ):
                if level == 0 and name in virtual_modules:
                    return virtual_modules[name]
                return original_import(name, globals, locals, fromlist, level)

            # We patch the import because we want the user to be able to import the sibling schemas,
            # but we don't want to pollute the global namespace (where there are also threading issues.)
            module = types.ModuleType(self.path.name)
            module.__file__ = str(self.path)
            module.__dict__["__builtins__"] = {
                **builtins.__dict__,
                "__import__": custom_import,
            }

            exec(self.source, module.__dict__)

            LOADED_MIGRATIONS[self.path] = module

            return module

    def upgrade_function(
        self, from_schema: SchemaEntry, to_schema: SchemaEntry
    ) -> Callable[[pydantic.BaseModel], pydantic.BaseModel]:
        module = self.isolated_module(from_schema, to_schema)
        if "upgrade" not in module.__dict__:
            raise ValueError(
                f"Migration file {self.path} must contain an 'upgrade' function."
            )

        upgrade_inner = module.__dict__["upgrade"]

        def upgrade(old: pydantic.BaseModel) -> pydantic.BaseModel:
            try:
                return upgrade_inner(old)
            except Exception as e:
                raise ValueError(
                    f"Failed to run upgrade function at {self.path}"
                    + f" from version {from_schema.version} to {to_schema.version}: {e}"
                ) from e

        return upgrade

    def process_dict(
        self,
        data: dict[str, Any],
        from_schema: SchemaEntry,
        to_schema: SchemaEntry,
    ) -> dict[str, Any]:
        loaded = from_schema.isolated_model().node.model_validate(data)
        processed = self.upgrade_function(from_schema, to_schema)(loaded)
        return processed.model_dump(mode="json")


class ModelEntry(pydantic.BaseModel, Generic[T]):
    live: ModelNode

    path: Path
    current: CurrentEntry | None
    versions: dict[int, SchemaEntry]
    migrations: dict[int, dict[int, MigrationEntry]]
    """The migration files. Note that a version can have multiple migrations to other versions."""

    @staticmethod
    def open(at: Path, model: type[T]) -> "ModelEntry[T]":
        at = at / model.__name__
        at.mkdir(parents=True, exist_ok=True)

        current: CurrentEntry | None = None
        versions: dict[int, SchemaEntry] = {}
        migrations: dict[int, dict[int, MigrationEntry]] = {}

        for file in at.glob("*.py"):
            number = _parse_schema_file_name(file.name)
            migration = _parse_schema_migration_file_name(file.name)
            if file.name == "__init__.py":
                pass
            elif file.name == "current.py":
                current = CurrentEntry.open(file)
            elif number is not None:
                versions[number] = SchemaEntry.open(at, number)
            elif migration is not None:
                [a, b] = migration
                entry = MigrationEntry.open(a, b, file)
                migrations[a] = migrations[a] if a in migrations else {}
                migrations[a][b] = entry
            else:
                raise ValueError(f"Invalid schema file name: {file.name}")

        for from_version, to_versions in migrations.items():
            for to_version in to_versions.keys():
                if from_version not in versions:
                    raise ValueError(
                        f"From version {from_version} not found in versions."
                    )
                if to_version not in versions:
                    raise ValueError(f"To version {to_version} not found in versions.")

        return ModelEntry(
            live=ModelNode.new(model),
            path=at,
            current=current,
            versions=versions,
            migrations=migrations,
        )

    def name(self) -> str:
        return self.live.node.__name__

    def latest_version(self) -> int | None:
        return max(self.versions.keys()) if self.versions else None

    def next_version(self) -> int:
        latest_version = self.latest_version()
        return latest_version + 1 if latest_version is not None else 0

    def has_migration(self, from_version: int, to_version: int) -> bool:
        return (
            from_version in self.migrations
            and to_version in self.migrations[from_version]
        )

    def update_current(self):
        if self.current is None:
            self.current = CurrentEntry.create(self.path / "current.py", self.live)
        else:
            self.current.update(self.live)

    def update_versioned_schemas(self):
        self.update_current()

        latest_version = self.latest_version()
        next_version = self.next_version()

        if (
            latest_version is not None
            and self.versions[latest_version].source
            == self.live.clean_source_recursive()
        ):
            _log.info(f"{self.name()} at version {latest_version} is up to date")
            return

        self.versions[next_version] = SchemaEntry.create(
            self.path, next_version, self.live
        )

    def assert_unchanged(self):
        if self.current is None:
            raise ValueError(f"No current version present for model {self.name()}.")
        self.current.assert_unchanged(self.live)
        latest_version = self.latest_version()
        if latest_version is not None:
            self.versions[latest_version].assert_equal(self.live)

    def assert_compatible(self, forward: bool, backward: bool):
        self.load_all_isolated_modules()  # To ensure sanity checks.

        if not self.versions:
            return

        from stable_pydantic import compatibility

        versions: list[SchemaEntry] = sorted(
            self.versions.values(),
            key=lambda x: x.version,
        )

        for i in range(len(versions) - 1):
            from_version = versions[i]
            to_version = versions[i + 1]

            # Skip compatibility check if a migration exists for this version pair
            if self.has_migration(from_version.version, to_version.version):
                _log.info(
                    f"Skipping compatibility check for {from_version.version} -> {to_version.version} (migration exists)"
                )
                continue

            compat = compatibility.check(from_version, to_version)
            if forward and not compat.forward.compatible:
                raise ValueError(
                    f"For model {self.name()}, version {from_version.version}"
                    + f" is not stricter than version {to_version.version}. Forward compatibility is not maintained."
                    + f"\nfrom_schema:\n{from_version.source}"
                    + f"\nto_schema:\n{to_version.source}"
                    + f"\ncompat:\n{compat.forward.tree()}"
                )

            if backward and not compat.backward.compatible:
                raise ValueError(
                    f"For model {self.name()}, version {to_version.version}"
                    + f" is stricter than version {from_version.version}. Backward compatibility is not maintained."
                    + f"\nfrom_schema:\n{from_version.source}"
                    + f"\nto_schema:\n{to_version.source}"
                    + f"\ncompat:\n{compat.backward.tree()}"
                )

    def load_all_isolated_modules(self):
        """Useful to ensure all modules are valid."""
        for version in self.versions.values():
            version.isolated_module()
        for from_version, migrations in self.migrations.items():
            for to_version, migration in migrations.items():
                migration.isolated_module(
                    self.versions[from_version], self.versions[to_version]
                )

    def to_json(self, value: pydantic.BaseModel) -> str:
        latest_version = self.latest_version()
        if latest_version is None:
            raise ValueError(f"No schema versions found for model {self.name()}.")
        versioned = VersionedData(version=latest_version, data=value)
        return versioned.model_dump_json()

    def from_json(self, json_str: str) -> T:
        return self.from_dict(json.loads(json_str))

    def from_json_with_version(self, version: int, json_str: str) -> T:
        if version == self.latest_version():
            # Fast path
            value = self.live.node.model_validate_json(json_str)
            assert isinstance(value, self.live.node)
            return cast(T, value)

        return self.from_dict_with_version(version, json.loads(json_str))

    def from_dict(self, data: dict[str, Any]) -> T:
        if "version" not in data:
            raise ValueError("JSON data missing 'version' field.")
        if "data" not in data:
            raise ValueError("JSON data missing 'data' field.")

        return self.from_dict_with_version(data["version"], data["data"])

    def from_dict_with_version(self, version: int, data: dict[str, Any]) -> T:
        if version not in self.versions:
            raise ValueError(
                f"Unknown schema version {version} for model {self.name()}."
            )

        latest_version = self.latest_version()
        if latest_version is None:
            raise ValueError(f"No schema versions found for model {self.name()}.")

        v = version
        while v < latest_version:
            # Check for migration from current version
            if v in self.migrations:
                # Note here that we could skip over multiple versions if the user provides a migration for it.
                to_version = max(self.migrations[v].keys())
                if to_version > v:
                    migration = self.migrations[v][to_version]
                    data = migration.process_dict(
                        data,
                        self.versions[v],
                        self.versions[to_version],
                    )
                    v = to_version
                    continue

            # No migration - find next version (compatible upgrade)
            for i in range(v + 1, latest_version + 1):
                if i in self.versions:
                    v = i
                    break
            else:
                raise ValueError(f"No next version found for version {v}.")

        if v == latest_version:
            value = self.live.node.model_validate(data)
            assert isinstance(value, self.live.node)
            return cast(T, value)
        else:
            raise ValueError(
                f"Unable to upgrade from version {version} to {latest_version} for model {self.name()}."
            )


class SchemaFilesystem(pydantic.BaseModel):
    models: dict[type[pydantic.BaseModel], ModelEntry]  # type: ignore[type-arg]

    @staticmethod
    def open(
        path: Path | str, models: Sequence[type[pydantic.BaseModel]]
    ) -> "SchemaFilesystem":
        path = Path(path)
        return SchemaFilesystem(
            models={model: ModelEntry.open(path, model) for model in models}
        )

    def update_current(self):
        for model in self.models.values():
            model.update_current()
        self.assert_unchanged_schemas()

    def update_versioned_schemas(self):
        for model in self.models.values():
            _log.info(f"Updating schemas for {model.live.node.__name__}")
            model.update_versioned_schemas()

        self.assert_unchanged_schemas()

    def assert_unchanged_schemas(self):
        for model in self.models.values():
            model.assert_unchanged()

    def assert_compatible_schemas(self, forward: bool = False, backward: bool = True):
        """
        Assert that the schemas are compatible in the given direction.

        Forward compatible: old clients keep working with new data.
        Backward compatible: new clients keep working with old data.
        """
        for model in self.models.values():
            model.load_all_isolated_modules()
            model.assert_compatible(forward, backward)

    def to_json(
        self, model: type[pydantic.BaseModel], value: pydantic.BaseModel
    ) -> str:
        return self.models[model].to_json(value)


def _schema_module_name(version: int) -> str:
    return f"v{version}_schema"


def _schema_file_name(version: int) -> str:
    return f"v{version}_schema.py"


def _parse_schema_file_name(name: str) -> int | None:
    if not name.startswith("v") or not name.endswith("_schema.py"):
        return None
    number_part = name[len("v") : -len("_schema.py")]
    try:
        return int(number_part)
    except ValueError:
        return None


def _parse_schema_migration_file_name(name: str) -> tuple[int, int] | None:
    if not name.startswith("v") or not name.endswith(".py") or "_to_" not in name:
        return None

    split = name[len("v") : -len(".py")].split("_to_")
    if len(split) != 2:
        return None

    try:
        return int(split[0]), int(split[1])
    except ValueError:
        return None


class VersionedData(pydantic.BaseModel, Generic[T]):
    """Wrapper model that stores version alongside the serialized data."""

    version: int
    data: T


def _validate_upgrade_module(source: str, path: Path) -> None:
    try:
        tree = ast.parse(source)
    except SyntaxError as e:
        raise ValueError(f"Migration file {path} has invalid Python syntax: {e}")

    for node in ast.walk(tree):
        if isinstance(node, ast.FunctionDef) and node.name == "upgrade":
            return

    raise ValueError(f"Migration file {path} must contain an 'upgrade' function.")
