from __future__ import annotations

import uuid
from pathlib import Path
from typing import Any, cast
from unittest.mock import MagicMock

import pytest
from pgvector.vector import Vector
from psycopg.types.json import Jsonb

from app.core.config import Settings
from app.persistence.indexing import (
    ACTIVATE_CURRENT_CHUNKS_SQL,
    BEGIN_INVOCATION_SQL,
    CACHE_LOOKUP_SQL,
    DEACTIVATE_OLD_CHUNKS_SQL,
    FINISH_INVOCATION_SQL,
    INSERT_CACHE_SQL,
    LEASE_FENCE_SQL,
    LOAD_PLAN_ITEMS_SQL,
    LOAD_PLAN_SQL,
    MARK_DOCUMENT_ACTIVE_SQL,
    MARK_VERSION_READY_SQL,
    PREPARE_STALE_CHUNK_SQL,
    PROGRESS_SQL,
    UPDATE_CHUNK_FROM_CACHE_SQL,
    UPSERT_ASSIGNMENT_SQL,
    VERIFY_CACHE_FOR_WRITE_SQL,
    VERIFY_READY_CHUNK_SQL,
    PostgresIndexingRepository,
)
from app.persistence.job_queue import JobLease, LeaseLostError
from app.ports.model_providers import ProviderUsage
from app.services.indexing import (
    CachedEmbedding,
    EmbeddingCacheLookup,
    EmbeddingWrite,
    embedding_cache_key,
)

JOB_ID = uuid.UUID("10000000-0000-0000-0000-000000000001")
DOCUMENT_VERSION_ID = uuid.UUID("20000000-0000-0000-0000-000000000002")
DOCUMENT_ID = uuid.UUID("30000000-0000-0000-0000-000000000003")
KNOWLEDGE_BASE_ID = uuid.UUID("40000000-0000-0000-0000-000000000004")
LEASE_TOKEN = uuid.UUID("50000000-0000-0000-0000-000000000005")
CHUNK_ID = uuid.UUID("60000000-0000-0000-0000-000000000006")
INVOCATION_ID = uuid.UUID("70000000-0000-0000-0000-000000000007")
LEASE = JobLease(JOB_ID, "embedding-worker-a", LEASE_TOKEN)
PROFILE_HASH = "a" * 64
TEXT = "已批准的斑岩铜矿地质证据"
TEXT_HASH = __import__("hashlib").sha256(TEXT.encode()).hexdigest()
CACHE_KEY = embedding_cache_key(TEXT_HASH, PROFILE_HASH)


class Cursor:
    def __init__(
        self,
        *,
        one: dict[str, object] | None = None,
        many: list[dict[str, object]] | None = None,
    ) -> None:
        self._one = one
        self._many = many or []

    def fetchone(self) -> dict[str, object] | None:
        return self._one

    def fetchall(self) -> list[dict[str, object]]:
        return self._many


def _settings(tmp_path: Path) -> Settings:
    password = tmp_path / "postgres-password"
    password.write_text("synthetic-test-password", encoding="utf-8")
    return Settings(postgres_password_file=password)


def _repository(
    tmp_path: Path,
    execute: Any,
) -> tuple[PostgresIndexingRepository, MagicMock, MagicMock, MagicMock]:
    transaction = MagicMock()
    connection = MagicMock()
    connection.__enter__.return_value = connection
    connection.transaction.return_value = transaction
    connection.execute.side_effect = execute
    factory = MagicMock(return_value=connection)
    registrar = MagicMock()
    repository = PostgresIndexingRepository(
        _settings(tmp_path),
        connection_factory=factory,
        vector_registrar=registrar,
    )
    return repository, connection, factory, registrar


def _fence_row() -> dict[str, object]:
    return {"resource_id": DOCUMENT_VERSION_ID}


def _plan_row() -> dict[str, object]:
    return {
        "knowledge_base_id": KNOWLEDGE_BASE_ID,
        "document_version_id": DOCUMENT_VERSION_ID,
        "review_state": "CLOUD_APPROVED",
        "outbound_manifest_sha256": "b" * 64,
        "expected_chunk_count": 1,
        "profile_hash": PROFILE_HASH,
        "model": "text-embedding-v4",
        "dimension": 1024,
        "synthetic": False,
        "manifest_count": 1,
        "eligible_chunk_count": 1,
    }


def _item_row(*, status: str = "PENDING") -> dict[str, object]:
    return {
        "chunk_id": CHUNK_ID,
        "ordinal": 0,
        "embedding_text": TEXT,
        "embedding_text_sha256": TEXT_HASH,
        "assignment_status": status,
    }


def _progress(*, ready: int, assignments: int = 1) -> dict[str, object]:
    return {
        "expected_count": 1,
        "chunk_count": 1,
        "assignment_count": assignments,
        "ready_count": ready,
    }


def _provider_write() -> EmbeddingWrite:
    return EmbeddingWrite(
        chunk_id=CHUNK_ID,
        batch_index=0,
        cache_key=CACHE_KEY,
        profile_hash=PROFILE_HASH,
        embedding_text_sha256=TEXT_HASH,
        source="provider",
        embedding=(1.0,) + (0.0,) * 1023,
        resolved_model="text-embedding-v4",
        provider_request_id="embed-request-1",
        usage=ProviderUsage(input_tokens=8, total_tokens=8),
        elapsed_ms=12.4,
    )


def _cache_write() -> EmbeddingWrite:
    return EmbeddingWrite(
        chunk_id=CHUNK_ID,
        batch_index=0,
        cache_key=CACHE_KEY,
        profile_hash=PROFILE_HASH,
        embedding_text_sha256=TEXT_HASH,
        source="cache",
        embedding=None,
        resolved_model="text-embedding-v4",
        provider_request_id=None,
        usage=ProviderUsage(),
        elapsed_ms=0.0,
    )


def test_load_plan_is_fenced_and_requires_manifest_profile_and_complete_projection(
    tmp_path: Path,
) -> None:
    calls: list[tuple[str, object]] = []

    def execute(statement: str, parameters: object) -> Cursor:
        calls.append((statement, parameters))
        if statement == LEASE_FENCE_SQL:
            return Cursor(one=_fence_row())
        if statement == LOAD_PLAN_SQL:
            return Cursor(one=_plan_row())
        if statement == LOAD_PLAN_ITEMS_SQL:
            return Cursor(many=[_item_row(status="STALE")])
        raise AssertionError("unexpected SQL")

    repository, _, factory, registrar = _repository(tmp_path, execute)

    plan = repository.load_approved_plan(
        lease=LEASE,
        document_version_id=DOCUMENT_VERSION_ID,
    )

    assert plan.document_version_id == DOCUMENT_VERSION_ID
    assert plan.profile.model == "text-embedding-v4"
    assert plan.items[0].assignment_status == "STALE"
    assert calls[0] == (
        LEASE_FENCE_SQL,
        (JOB_ID, LEASE.worker_id, LEASE_TOKEN),
    )
    assert calls[1][1] == (DOCUMENT_VERSION_ID,)
    assert calls[2][1] == (DOCUMENT_VERSION_ID,)
    factory.assert_called_once()
    registrar.assert_called_once()

    normalized_header = " ".join(LOAD_PLAN_SQL.lower().split())
    normalized_items = " ".join(LOAD_PLAN_ITEMS_SQL.lower().split())
    assert "version.review_state = 'cloud_approved'" in normalized_header
    assert "profile.enabled is true" in normalized_header
    assert (
        "knowledge_base.active_embedding_profile_hash = profile.profile_hash" in normalized_header
    )
    assert "join rag.outbound_manifest_items as item" in normalized_items
    assert "when assignment.status = 'ready' then 'stale'" in normalized_items


def test_expired_or_wrong_lease_stops_before_any_business_read_or_write(tmp_path: Path) -> None:
    calls: list[str] = []

    def execute(statement: str, _parameters: object) -> Cursor:
        calls.append(statement)
        if statement == LEASE_FENCE_SQL:
            return Cursor(one=None)
        raise AssertionError("business SQL must not run after a failed fence")

    repository, _, _, _ = _repository(tmp_path, execute)

    with pytest.raises(LeaseLostError):
        repository.fenced_persist_batch(
            lease=LEASE,
            document_version_id=DOCUMENT_VERSION_ID,
            profile_hash=PROFILE_HASH,
            writes=(_provider_write(),),
        )

    assert calls == [LEASE_FENCE_SQL]
    normalized = " ".join(LEASE_FENCE_SQL.lower().split())
    for condition in (
        "job.status = 'running'",
        "job.lease_owner = %s",
        "job.lease_token = %s",
        "job.lease_until >= now()",
        "job.resource_type = 'document_version'",
        "for update",
    ):
        assert condition in normalized


def test_cache_lookup_uses_physical_profile_text_key_and_revalidates_active_profile(
    tmp_path: Path,
) -> None:
    calls: list[tuple[str, object]] = []

    def execute(statement: str, parameters: object) -> Cursor:
        calls.append((statement, parameters))
        if statement == LEASE_FENCE_SQL:
            return Cursor(one=_fence_row())
        if statement == CACHE_LOOKUP_SQL:
            return Cursor(
                one={
                    "profile_hash": PROFILE_HASH,
                    "embedding_text_sha256": TEXT_HASH,
                    "resolved_model": "text-embedding-v4",
                    "dimension": 1024,
                }
            )
        raise AssertionError("unexpected SQL")

    repository, _, _, _ = _repository(tmp_path, execute)
    lookup = EmbeddingCacheLookup(CACHE_KEY, PROFILE_HASH, TEXT_HASH)

    found = repository.lookup_cache(lease=LEASE, lookups=(lookup,))

    assert found == {
        CACHE_KEY: CachedEmbedding(
            cache_key=CACHE_KEY,
            profile_hash=PROFILE_HASH,
            embedding_text_sha256=TEXT_HASH,
            resolved_model="text-embedding-v4",
            dimension=1024,
        )
    }
    assert calls[-1][1] == (DOCUMENT_VERSION_ID, PROFILE_HASH, TEXT_HASH)
    assert "vector_norm(cache.embedding) > 0" in CACHE_LOOKUP_SQL


def test_invocation_writes_are_fenced_metadata_only_and_bound_to_job_trace(
    tmp_path: Path,
) -> None:
    calls: list[tuple[str, object]] = []

    def execute(statement: str, parameters: object) -> Cursor:
        calls.append((statement, parameters))
        if statement == LEASE_FENCE_SQL:
            return Cursor(one=_fence_row())
        if statement == BEGIN_INVOCATION_SQL:
            return Cursor(one={"id": INVOCATION_ID})
        if statement == FINISH_INVOCATION_SQL:
            return Cursor(one={"id": INVOCATION_ID})
        raise AssertionError("unexpected SQL")

    repository, _, factory, _ = _repository(tmp_path, execute)

    invocation_id = repository.begin_model_invocation(
        lease=LEASE,
        trace_id=JOB_ID,
        profile_hash=PROFILE_HASH,
        model="text-embedding-v4",
        item_count=3,
    )
    repository.finish_model_invocation(
        lease=LEASE,
        invocation_id=invocation_id,
        status="SUCCEEDED",
        provider_request_id="request-1",
        usage=ProviderUsage(input_tokens=9, total_tokens=9),
        elapsed_ms=4.6,
        error_code=None,
    )

    assert invocation_id == INVOCATION_ID
    begin_call = next(call for call in calls if call[0] == BEGIN_INVOCATION_SQL)
    finish_call = next(call for call in calls if call[0] == FINISH_INVOCATION_SQL)
    assert begin_call[1] == (
        JOB_ID,
        3,
        DOCUMENT_VERSION_ID,
        PROFILE_HASH,
        "text-embedding-v4",
    )
    assert finish_call[1] == (
        "SUCCEEDED",
        "request-1",
        9,
        0,
        9,
        5,
        None,
        INVOCATION_ID,
        JOB_ID,
        DOCUMENT_VERSION_ID,
    )
    assert "embedding_text" not in BEGIN_INVOCATION_SQL
    assert "cloud_text" not in FINISH_INVOCATION_SQL
    assert "vector" not in FINISH_INVOCATION_SQL
    assert factory.call_count == 2


def test_provider_persist_inserts_cache_then_assignment_and_canonical_chunk(
    tmp_path: Path,
) -> None:
    calls: list[tuple[str, object]] = []

    def execute(statement: str, parameters: object = ()) -> Cursor:
        calls.append((statement, parameters))
        if statement == LEASE_FENCE_SQL:
            return Cursor(one=_fence_row())
        if statement in {
            INSERT_CACHE_SQL,
            VERIFY_CACHE_FOR_WRITE_SQL,
            UPSERT_ASSIGNMENT_SQL,
            UPDATE_CHUNK_FROM_CACHE_SQL,
        }:
            return Cursor(one={"id": CHUNK_ID})
        if statement == PREPARE_STALE_CHUNK_SQL:
            return Cursor()
        if statement == PROGRESS_SQL:
            return Cursor(one=_progress(ready=1))
        raise AssertionError("unexpected SQL")

    repository, _, _, _ = _repository(tmp_path, execute)

    progress = repository.fenced_persist_batch(
        lease=LEASE,
        document_version_id=DOCUMENT_VERSION_ID,
        profile_hash=PROFILE_HASH,
        writes=(_provider_write(),),
    )

    assert progress.ready_count == 1
    statements = [statement for statement, _ in calls]
    assert statements == [
        LEASE_FENCE_SQL,
        INSERT_CACHE_SQL,
        VERIFY_CACHE_FOR_WRITE_SQL,
        UPSERT_ASSIGNMENT_SQL,
        PREPARE_STALE_CHUNK_SQL,
        UPDATE_CHUNK_FROM_CACHE_SQL,
        PROGRESS_SQL,
    ]
    insert_parameters = cast(
        tuple[object, ...],
        next(parameters for statement, parameters in calls if statement == INSERT_CACHE_SQL),
    )
    assert isinstance(insert_parameters[2], Vector)
    assert isinstance(insert_parameters[5], Jsonb)
    assert TEXT not in repr(insert_parameters)
    assert "on conflict (profile_hash, embedding_text_sha256) do nothing" in " ".join(
        INSERT_CACHE_SQL.lower().split()
    )


def test_cache_recovery_repairs_stale_ready_projection_without_reinserting_cache(
    tmp_path: Path,
) -> None:
    calls: list[str] = []

    def execute(statement: str, _parameters: object = ()) -> Cursor:
        calls.append(statement)
        if statement == LEASE_FENCE_SQL:
            return Cursor(one=_fence_row())
        if statement in {VERIFY_CACHE_FOR_WRITE_SQL, UPSERT_ASSIGNMENT_SQL}:
            return Cursor(one={"available": 1})
        if statement == PREPARE_STALE_CHUNK_SQL:
            return Cursor(one={"id": CHUNK_ID})
        if statement == UPDATE_CHUNK_FROM_CACHE_SQL:
            return Cursor(one={"id": CHUNK_ID})
        if statement == PROGRESS_SQL:
            return Cursor(one=_progress(ready=1))
        raise AssertionError("unexpected SQL")

    repository, _, _, _ = _repository(tmp_path, execute)

    progress = repository.fenced_persist_batch(
        lease=LEASE,
        document_version_id=DOCUMENT_VERSION_ID,
        profile_hash=PROFILE_HASH,
        writes=(_cache_write(),),
    )

    assert progress.ready_count == 1
    assert INSERT_CACHE_SQL not in calls
    assert calls.index(PREPARE_STALE_CHUNK_SQL) < calls.index(UPDATE_CHUNK_FROM_CACHE_SQL)
    assert "index_status = 'EMBEDDING'" in PREPARE_STALE_CHUNK_SQL


def test_ready_idempotent_replay_skips_vector_update_and_verifies_existing_projection(
    tmp_path: Path,
) -> None:
    calls: list[str] = []

    def execute(statement: str, _parameters: object = ()) -> Cursor:
        calls.append(statement)
        if statement == LEASE_FENCE_SQL:
            return Cursor(one=_fence_row())
        if statement in {VERIFY_CACHE_FOR_WRITE_SQL, UPSERT_ASSIGNMENT_SQL}:
            return Cursor(one={"available": 1})
        if statement in {PREPARE_STALE_CHUNK_SQL, UPDATE_CHUNK_FROM_CACHE_SQL}:
            return Cursor(one=None)
        if statement == VERIFY_READY_CHUNK_SQL:
            return Cursor(one={"id": CHUNK_ID})
        if statement == PROGRESS_SQL:
            return Cursor(one=_progress(ready=1))
        raise AssertionError("unexpected SQL")

    repository, _, _, _ = _repository(tmp_path, execute)

    repository.fenced_persist_batch(
        lease=LEASE,
        document_version_id=DOCUMENT_VERSION_ID,
        profile_hash=PROFILE_HASH,
        writes=(_cache_write(),),
    )

    assert VERIFY_READY_CHUNK_SQL in calls
    assert "chunk.index_status <> 'READY'" in UPDATE_CHUNK_FROM_CACHE_SQL


def test_activation_checks_complete_projection_then_uses_trigger_safe_order(
    tmp_path: Path,
) -> None:
    calls: list[str] = []

    def execute(statement: str, _parameters: object = ()) -> Cursor:
        calls.append(statement)
        if statement == LEASE_FENCE_SQL:
            return Cursor(one=_fence_row())
        if statement == PROGRESS_SQL:
            return Cursor(one=_progress(ready=1))
        if statement == MARK_VERSION_READY_SQL:
            return Cursor(one={"document_id": DOCUMENT_ID})
        if statement == MARK_DOCUMENT_ACTIVE_SQL:
            return Cursor(one={"id": DOCUMENT_ID})
        if statement == DEACTIVATE_OLD_CHUNKS_SQL:
            return Cursor()
        if statement == ACTIVATE_CURRENT_CHUNKS_SQL:
            return Cursor(many=[{"id": CHUNK_ID}])
        raise AssertionError("unexpected SQL")

    repository, _, _, _ = _repository(tmp_path, execute)

    activated = repository.fenced_activate(
        lease=LEASE,
        document_version_id=DOCUMENT_VERSION_ID,
        profile_hash=PROFILE_HASH,
        expected_count=1,
    )

    assert activated is True
    assert calls == [
        LEASE_FENCE_SQL,
        PROGRESS_SQL,
        MARK_VERSION_READY_SQL,
        MARK_DOCUMENT_ACTIVE_SQL,
        DEACTIVATE_OLD_CHUNKS_SQL,
        ACTIVATE_CURRENT_CHUNKS_SQL,
    ]


def test_incomplete_activation_has_no_version_document_or_searchability_write(
    tmp_path: Path,
) -> None:
    calls: list[str] = []

    def execute(statement: str, _parameters: object = ()) -> Cursor:
        calls.append(statement)
        if statement == LEASE_FENCE_SQL:
            return Cursor(one=_fence_row())
        if statement == PROGRESS_SQL:
            return Cursor(one=_progress(ready=0, assignments=0))
        raise AssertionError("activation writes must not run while incomplete")

    repository, _, _, _ = _repository(tmp_path, execute)

    activated = repository.fenced_activate(
        lease=LEASE,
        document_version_id=DOCUMENT_VERSION_ID,
        profile_hash=PROFILE_HASH,
        expected_count=1,
    )

    assert activated is False
    assert calls == [LEASE_FENCE_SQL, PROGRESS_SQL]