from __future__ import annotations

from collections.abc import AsyncIterator, Sequence
from dataclasses import replace
from uuid import UUID, uuid4

import pytest

from ragcore.application.indexes.models import (
    IndexBuildError,
    MalformedIndexPoint,
    StoredIndexPoint,
    index_point_fingerprint,
)
from ragcore.application.indexes.shadow_migration import (
    ShadowEmbeddingMigrationService,
    ShadowMigrationCheckpoint,
    ShadowMigrationStatus,
    ShadowPreflight,
)
from ragcore.infrastructure.qdrant.models import (
    CollectionKind,
    CollectionProfile,
    IndexPayload,
    IndexPoint,
)


class MemoryCheckpoints:
    def __init__(self, *, cancel_after_batch: bool = False) -> None:
        self.item: ShadowMigrationCheckpoint | None = None
        self.cancel_after_batch = cancel_after_batch

    async def load(self, _migration_id: UUID) -> ShadowMigrationCheckpoint | None:
        return self.item

    async def save(self, checkpoint: ShadowMigrationCheckpoint) -> None:
        if self.cancel_after_batch and checkpoint.request_batches == 1:
            checkpoint = replace(checkpoint, cancel_requested=True)
        self.item = checkpoint


class Source:
    def __init__(self, points: Sequence[IndexPoint]) -> None:
        self.points = points

    async def iter_points(self, _profile: CollectionProfile) -> AsyncIterator[IndexPoint]:
        for point in self.points:
            yield point


class Store:
    def __init__(self, *, fail_once: bool = False, omit_last: bool = False) -> None:
        self.points: dict[UUID, StoredIndexPoint] = {}
        self.fail_once = fail_once
        self.omit_last = omit_last
        self.alias = "rag_dense_legacy_v1"
        self.alias_mutations = 0
        self.upsert_calls = 0

    async def create_shadow(self, _profile: CollectionProfile) -> None:
        return None

    async def delete_shadow(self, _profile: CollectionProfile) -> None:
        raise AssertionError("validated shadow is retained")

    async def upsert_points(
        self, _profile: CollectionProfile, points: Sequence[IndexPoint]
    ) -> None:
        self.upsert_calls += 1
        if self.fail_once:
            self.fail_once = False
            raise TimeoutError("transient")
        for point in points:
            self.points[point.point_id] = StoredIndexPoint(
                point.point_id, index_point_fingerprint(point)
            )

    async def delete_points(self, _profile: CollectionProfile, _point_ids: Sequence[UUID]) -> None:
        raise AssertionError("no destructive cleanup")

    async def delete_raw_points(self, _profile: CollectionProfile, _raw_ids: Sequence[str]) -> None:
        raise AssertionError("no destructive cleanup")

    async def count_points(self, _profile: CollectionProfile) -> int:
        return len(self.points)

    async def list_points(
        self, _profile: CollectionProfile
    ) -> tuple[tuple[StoredIndexPoint, ...], tuple[MalformedIndexPoint, ...]]:
        values = tuple(self.points.values())
        return (values[:-1] if self.omit_last and values else values), ()

    async def current_alias_target(self, _alias_name: str) -> str:
        return self.alias

    async def activate(self, _profile: CollectionProfile) -> None:
        self.alias_mutations += 1

    async def rollback(self, _alias_name: str, _previous_collection: str | None) -> None:
        self.alias_mutations += 1


def profile() -> CollectionProfile:
    return CollectionProfile(
        kind=CollectionKind.DENSE,
        profile="modelforge-shadow",
        version="v1",
        dimensions=3,
        model_digest="a" * 64,
        embedding_space_id=uuid4(),
    )


def points(item: CollectionProfile, count: int = 3) -> tuple[IndexPoint, ...]:
    workspace, space, document, version, source = (uuid4() for _ in range(5))
    return tuple(
        IndexPoint(
            point_id=(chunk := uuid4()),
            dense_vector=(1.0, 0.0, 0.0),
            payload=IndexPayload(
                workspace_id=workspace,
                space_id=space,
                document_id=document,
                document_version_id=version,
                chunk_id=chunk,
                parent_chunk_id=None,
                source_id=source,
                status="published",
                content_type="text/plain",
                language="en",
                trust_level="internal",
                location={"structure_path": f"/{index}"},
                parser_version="parser-v1",
                chunk_profile_version=item.version,
                embedding_model_digest=item.model_digest,
                text=f"chunk {index}",
                excerpt=f"chunk {index}",
                embedding_space_id=item.embedding_space_id,
            ),
        )
        for index in range(count)
    )


def checkpoint(item: CollectionProfile, total: int = 3) -> ShadowMigrationCheckpoint:
    assert item.embedding_space_id is not None
    return ShadowMigrationCheckpoint(
        migration_id=uuid4(),
        source_index="rag_dense_legacy_v1",
        target_index=item.collection_name,
        source_embedding_space="legacy-observed",
        target_embedding_space_id=item.embedding_space_id,
        corpus_revision="corpus-sha256",
        total_chunks=total,
    )


async def preflight(*, passed: bool = True) -> ShadowPreflight:
    return ShadowPreflight(
        production_healthy=passed,
        source_available=True,
        modelforge_healthy=True,
        credential_valid=True,
        embedding_space_matches=True,
        dimension_matches=True,
        distance_metric_matches=True,
        qdrant_reachable=True,
        storage_sufficient=True,
        target_conflict_free=True,
    )


@pytest.mark.asyncio
async def test_preflight_failure_writes_no_vectors() -> None:
    item = profile()
    store, state = Store(), MemoryCheckpoints()
    service = ShadowEmbeddingMigrationService(store, state)
    with pytest.raises(IndexBuildError, match="preflight"):
        await service.run(
            checkpoint(item), item, Source(points(item)), lambda: preflight(passed=False)
        )
    assert not store.points
    assert state.item is not None and state.item.status is ShadowMigrationStatus.FAILED


@pytest.mark.asyncio
async def test_backfill_is_bounded_validated_and_never_mutates_source_alias() -> None:
    item, store, state = profile(), Store(), MemoryCheckpoints()
    service = ShadowEmbeddingMigrationService(store, state, batch_size=2)
    result = await service.run(checkpoint(item), item, Source(points(item)), preflight)
    assert result.status is ShadowMigrationStatus.READY_FOR_EVALUATION
    assert result.completed_chunks == 3 and result.failed_chunks == 0
    assert result.request_batches == 2
    assert result.validation["passed"] is True
    assert result.validation["source_alias_unchanged"] is True
    assert store.alias_mutations == 0


@pytest.mark.asyncio
async def test_resume_skips_integrity_valid_existing_chunk_and_is_idempotent() -> None:
    item, store, states = profile(), Store(), MemoryCheckpoints()
    source_points = points(item)
    existing = source_points[0]
    store.points[existing.point_id] = StoredIndexPoint(
        existing.point_id, index_point_fingerprint(existing)
    )
    initial = replace(
        checkpoint(item), status=ShadowMigrationStatus.BACKFILLING, completed_chunks=1
    )
    states.item = initial
    service = ShadowEmbeddingMigrationService(store, states, batch_size=1)
    result = await service.run(initial, item, Source(source_points), preflight)
    assert result.completed_chunks == 3
    assert store.upsert_calls == 2
    repeated = await service.run(initial, item, Source(source_points), preflight)
    assert repeated.status is ShadowMigrationStatus.READY_FOR_EVALUATION
    assert store.upsert_calls == 2


@pytest.mark.asyncio
async def test_cancel_finishes_current_batch_and_remains_resumable() -> None:
    item, store, states = profile(), Store(), MemoryCheckpoints(cancel_after_batch=True)
    source_points = points(item)
    service = ShadowEmbeddingMigrationService(store, states, batch_size=1)
    cancelled = await service.run(checkpoint(item), item, Source(source_points), preflight)
    assert cancelled.status is ShadowMigrationStatus.CANCELLED
    assert cancelled.completed_chunks == 1
    states.cancel_after_batch = False
    resumed = await service.run(cancelled, item, Source(source_points), preflight)
    assert resumed.status is ShadowMigrationStatus.READY_FOR_EVALUATION
    assert resumed.completed_chunks == 3


@pytest.mark.asyncio
async def test_transient_retry_is_bounded_and_counted() -> None:
    item, store, states = profile(), Store(fail_once=True), MemoryCheckpoints()
    service = ShadowEmbeddingMigrationService(
        store,
        states,
        batch_size=3,
        max_retries=1,
        retryable=lambda error: isinstance(error, TimeoutError),
    )
    result = await service.run(checkpoint(item), item, Source(points(item)), preflight)
    assert result.retried_chunks == 3
    assert store.upsert_calls == 2


@pytest.mark.asyncio
async def test_missing_point_blocks_evaluation_readiness() -> None:
    item, store, states = profile(), Store(omit_last=True), MemoryCheckpoints()
    service = ShadowEmbeddingMigrationService(store, states, batch_size=3)
    with pytest.raises(IndexBuildError, match="validation"):
        await service.run(checkpoint(item), item, Source(points(item)), preflight)
    assert states.item is not None and states.item.status is ShadowMigrationStatus.FAILED