"""On-demand multilingual FastEmbed quality and operational benchmark."""

from __future__ import annotations

import json
import os
import time
from pathlib import Path

import pytest
from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, async_sessionmaker

from basic_memory.config import DatabaseBackend
from basic_memory.repository.embedding_provider_factory import create_embedding_provider
from basic_memory.schemas.search import SearchItemType, SearchQuery

from semantic.conftest import SearchCombo, create_search_service, skip_if_needed
from semantic.multilingual_benchmark import (
    RuntimeMeasurements,
    RetrievalObservation,
    benchmark_context,
    benchmark_retrieval_mode,
    benchmark_storage_case,
    embedding_benchmark_config,
    embedding_model_case,
    model_cache_size_bytes,
    process_peak_rss_bytes,
    process_rss_bytes,
    result_permalinks,
    retrieval_summaries,
    vector_storage_size_bytes,
    write_multilingual_benchmark_artifact,
)
from semantic.multilingual_corpus import MULTILINGUAL_CORPUS, seed_multilingual_documents

type EngineFactoryResult = tuple[AsyncEngine, async_sessionmaker[AsyncSession]]


@pytest.fixture
def multilingual_engine_factory(
    request: pytest.FixtureRequest,
) -> EngineFactoryResult | None:
    """Resolve only the SQL fixture selected for this benchmark process."""
    storage_case = benchmark_storage_case(os.getenv("BASIC_MEMORY_MULTILINGUAL_BACKEND", "sqlite"))
    fixture_name = (
        "sqlite_engine_factory"
        if storage_case.database_backend is DatabaseBackend.SQLITE
        else "postgres_engine_factory"
    )
    return request.getfixturevalue(fixture_name)


@pytest.mark.asyncio
@pytest.mark.benchmark
async def test_multilingual_embedding_benchmark(
    multilingual_engine_factory: EngineFactoryResult | None,
    tmp_path: Path,
) -> None:
    """Measure one model/backend pair in an isolated pytest process."""
    model_case = embedding_model_case(os.getenv("BASIC_MEMORY_MULTILINGUAL_MODEL", "bge-small-en"))
    storage_case = benchmark_storage_case(os.getenv("BASIC_MEMORY_MULTILINGUAL_BACKEND", "sqlite"))
    backend = storage_case.database_backend
    retrieval_mode = benchmark_retrieval_mode(
        os.getenv("BASIC_MEMORY_MULTILINGUAL_RETRIEVAL_MODE", "vector")
    )
    production_threshold = float(os.getenv("BASIC_MEMORY_MULTILINGUAL_THRESHOLD", "0.55"))
    if not 0.0 <= production_threshold <= 1.0:
        raise ValueError("BASIC_MEMORY_MULTILINGUAL_THRESHOLD must be between 0.0 and 1.0")

    combo = SearchCombo(
        name=f"{storage_case.key}-{model_case.key}",
        backend=backend,
        provider_name="fastembed",
        dimensions=model_case.dimensions,
    )
    skip_if_needed(combo)

    if multilingual_engine_factory is None:
        pytest.skip("Selected benchmark engine is not available")
    engine_factory_result = multilingual_engine_factory
    engine, _ = engine_factory_result

    milvus_storage_directory = tmp_path / "milvus"
    milvus_uri = None
    if storage_case.vector_index_name == "milvus":
        milvus_storage_directory.mkdir()
        milvus_uri = str(milvus_storage_directory / "vectors.db")

    benchmark_config = embedding_benchmark_config(
        model_case,
        storage_case,
        milvus_uri=milvus_uri,
    )
    provider = create_embedding_provider(benchmark_config)

    # Each invocation benchmarks exactly one model, so this first embedding is a genuine
    # process-cold model load while later indexing and queries exercise the warmed provider.
    rss_before_load = process_rss_bytes()
    load_started = time.perf_counter()
    warmup_vector = await provider.embed_query("multilingual benchmark warmup")
    cold_load_seconds = time.perf_counter() - load_started
    rss_after_load = process_rss_bytes()
    assert len(warmup_vector) == model_case.dimensions

    search_service = await create_search_service(
        engine_factory_result,
        combo,
        tmp_path,
        embedding_provider=provider,
        benchmark_config=benchmark_config,
    )

    indexing_started = time.perf_counter()
    entities = await seed_multilingual_documents(
        search_service,
        MULTILINGUAL_CORPUS.documents,
    )
    indexing_seconds = time.perf_counter() - indexing_started
    assert len(entities) == len(MULTILINGUAL_CORPUS.documents)
    rss_after_index = process_rss_bytes()

    observations: list[RetrievalObservation] = []
    for query in MULTILINGUAL_CORPUS.queries:
        ranking_started = time.perf_counter()
        ranked_results = await search_service.search(
            SearchQuery(
                text=query.text,
                retrieval_mode=retrieval_mode,
                entity_types=[SearchItemType.ENTITY],
                min_similarity=0.0,
            ),
            limit=10,
        )
        ranking_latency = time.perf_counter() - ranking_started

        accepted_started = time.perf_counter()
        accepted_results = await search_service.search(
            SearchQuery(
                text=query.text,
                retrieval_mode=retrieval_mode,
                entity_types=[SearchItemType.ENTITY],
                min_similarity=production_threshold,
            ),
            limit=10,
        )
        accepted_latency = time.perf_counter() - accepted_started

        observations.append(
            RetrievalObservation(
                query=query,
                ranked_permalinks=result_permalinks(ranked_results, query),
                accepted_permalinks=result_permalinks(accepted_results, query),
                ranking_latency_seconds=ranking_latency,
                accepted_latency_seconds=accepted_latency,
            )
        )

    summaries = retrieval_summaries(observations)
    runtime = RuntimeMeasurements(
        document_count=len(entities),
        cold_load_seconds=cold_load_seconds,
        indexing_seconds=indexing_seconds,
        rss_before_load_bytes=rss_before_load,
        rss_after_load_bytes=rss_after_load,
        rss_after_index_bytes=rss_after_index,
        peak_rss_bytes=process_peak_rss_bytes(),
        model_cache_bytes=model_cache_size_bytes(model_case, benchmark_config),
        vector_storage_bytes=await vector_storage_size_bytes(
            engine,
            storage_case,
            milvus_storage_directory=milvus_storage_directory,
        ),
    )
    context = benchmark_context(
        model_case,
        storage_case,
        retrieval_mode,
        MULTILINGUAL_CORPUS.version,
        production_threshold,
    )

    print("\nMultilingual benchmark context")
    print(json.dumps(context, ensure_ascii=False, indent=2, sort_keys=True))
    print("\nMultilingual benchmark quality")
    print(
        json.dumps(
            {summary.name: summary.as_metrics() for summary in summaries},
            indent=2,
            sort_keys=True,
        )
    )
    print("\nMultilingual benchmark runtime")
    print(json.dumps(runtime.as_metrics(), indent=2, sort_keys=True))

    output_value = os.getenv("BASIC_MEMORY_BENCHMARK_OUTPUT")
    if output_value:
        write_multilingual_benchmark_artifact(
            Path(output_value).expanduser(),
            context,
            summaries,
            runtime,
            observations,
        )