"""Provider adapters, exercised with no AWS account and no network.

Every Bedrock call goes through a recording stub, so what is under test is the
part that can actually be wrong offline: the request body we send, and our
reading of the response bodies AWS documents. The one thing these tests cannot
establish is whether AWS accepts that body — that needs the live probe, and
the coordination handoff says so explicitly.
"""
import io
import json
import math

import pytest

from ingestion.embed import (
    BGE_M3,
    COHERE_V4,
    TITAN_V2,
    INPUT_DOCUMENT,
    INPUT_QUERY,
    Boto3BedrockInvoker,
    EmbeddingVector,
    build_provider,
    provider_names,
    text_digest,
)
from ingestion.embed import probe
from ingestion.embed.bedrock_cohere import CohereEmbedV4
from ingestion.embed.bedrock_titan import TitanTextEmbeddingsV2
from ingestion.embed.local_bge_m3 import BgeM3Local


class RecordingInvoker:
    """Stands in for Bedrock; remembers every request it was handed."""

    def __init__(self, responses):
        self._responses = list(responses)
        self.calls = []

    def invoke_json(self, model_id, payload, accept="application/json"):
        self.calls.append(
            {"model_id": model_id, "payload": payload, "accept": accept}
        )
        return self._responses.pop(0)


def _titan_response(dimensions=1024, token_count=12):
    return {
        "embedding": [0.01] * dimensions,
        "inputTextTokenCount": token_count,
        "embeddingsByType": {"float": [0.01] * dimensions},
    }


def _cohere_by_type_response(rows, dimensions=1024):
    return {
        "id": "stub-id",
        "response_type": "embeddings_by_type",
        "embeddings": {"float": [[0.02] * dimensions for _ in range(rows)]},
        "texts": ["stub"] * rows,
    }


def _cohere_floats_response(rows, dimensions=1024):
    return {
        "id": "stub-id",
        "response_type": "embeddings_floats",
        "embeddings": [[0.02] * dimensions for _ in range(rows)],
    }


def test_titan_request_body_matches_the_documented_v2_shape():
    invoker = RecordingInvoker([_titan_response()])
    provider = TitanTextEmbeddingsV2(invoker, dimensions=1024, normalize=True)

    provider.embed_documents(["paracetamol"])

    payload = invoker.calls[0]["payload"]
    assert invoker.calls[0]["model_id"] == "amazon.titan-embed-text-v2:0"
    assert payload == {
        "inputText": "paracetamol",
        "dimensions": 1024,
        "normalize": True,
    }


def test_titan_records_provenance_and_reported_token_count():
    invoker = RecordingInvoker([_titan_response(token_count=7)])
    provider = TitanTextEmbeddingsV2(invoker)

    vector = provider.embed_documents(["paracetamol"]).vectors[0]

    assert vector.model_id == "amazon.titan-embed-text-v2:0"
    assert vector.provider == TITAN_V2
    assert vector.dimensions == 1024
    assert vector.input_kind == INPUT_DOCUMENT
    assert vector.normalized is True
    assert vector.input_token_count == 7
    assert vector.text_sha256 == text_digest("paracetamol")


def test_titan_sends_one_request_per_text():
    invoker = RecordingInvoker([_titan_response(), _titan_response()])
    provider = TitanTextEmbeddingsV2(invoker)

    batch = provider.embed_documents(["a", "b"])

    assert batch.request_count == 2
    assert len(batch.vectors) == 2


def test_titan_rejects_a_dimension_the_model_does_not_offer():
    with pytest.raises(ValueError, match="supports"):
        TitanTextEmbeddingsV2(RecordingInvoker([]), dimensions=768)


def test_cohere_uses_search_document_for_corpus_and_search_query_for_queries():
    invoker = RecordingInvoker(
        [_cohere_by_type_response(1), _cohere_by_type_response(1)]
    )
    provider = CohereEmbedV4(invoker)

    provider.embed_documents(["metformin"])
    provider.embed_queries(["liều metformin"])

    assert invoker.calls[0]["payload"]["input_type"] == "search_document"
    assert invoker.calls[1]["payload"]["input_type"] == "search_query"


def test_cohere_request_body_pins_dimension_float_type_and_no_truncation():
    invoker = RecordingInvoker([_cohere_by_type_response(2)])
    provider = CohereEmbedV4(invoker, dimensions=1024)

    provider.embed_documents(["a", "b"])

    payload = invoker.calls[0]["payload"]
    assert invoker.calls[0]["model_id"] == "cohere.embed-v4:0"
    assert payload["texts"] == ["a", "b"]
    assert payload["embedding_types"] == ["float"]
    # Left unset the model would return 1536, which no 1024-wide collection
    # can accept.
    assert payload["output_dimension"] == 1024
    # An over-length input must fail, not arrive silently shortened.
    assert payload["truncate"] == "NONE"
    assert invoker.calls[0]["accept"] == "*/*"


def test_cohere_reads_the_embeddings_by_type_response():
    invoker = RecordingInvoker([_cohere_by_type_response(2)])
    provider = CohereEmbedV4(invoker)

    batch = provider.embed_documents(["a", "b"])

    assert len(batch.vectors) == 2
    assert all(len(v.values) == 1024 for v in batch.vectors)
    assert batch.request_count == 1


def test_cohere_also_reads_the_plain_embeddings_floats_response():
    invoker = RecordingInvoker([_cohere_floats_response(2)])
    provider = CohereEmbedV4(invoker)

    batch = provider.embed_documents(["a", "b"])

    assert len(batch.vectors) == 2
    assert all(len(v.values) == 1024 for v in batch.vectors)


def test_cohere_leaves_normalization_unknown_because_the_docs_do_not_say():
    invoker = RecordingInvoker([_cohere_by_type_response(1)])

    vector = CohereEmbedV4(invoker).embed_documents(["a"]).vectors[0]

    assert vector.normalized is None


def test_cohere_splits_at_the_documented_96_text_ceiling():
    invoker = RecordingInvoker(
        [_cohere_by_type_response(96), _cohere_by_type_response(4)]
    )
    provider = CohereEmbedV4(invoker)

    batch = provider.embed_documents([f"t{i}" for i in range(100)])

    assert batch.request_count == 2
    assert len(invoker.calls[0]["payload"]["texts"]) == 96
    assert len(invoker.calls[1]["payload"]["texts"]) == 4
    assert len(batch.vectors) == 100


def test_cohere_rejects_a_batch_size_above_the_documented_ceiling():
    with pytest.raises(ValueError, match="batch_size"):
        CohereEmbedV4(RecordingInvoker([]), batch_size=97)


def test_a_wrong_width_vector_fails_instead_of_entering_the_corpus():
    invoker = RecordingInvoker([_titan_response(dimensions=512)])
    provider = TitanTextEmbeddingsV2(invoker, dimensions=1024)

    with pytest.raises(ValueError, match="512 dimensions"):
        provider.embed_documents(["a"])


def test_a_response_missing_its_vectors_fails_loudly():
    invoker = RecordingInvoker([{"id": "stub", "response_type": "x"}])

    with pytest.raises(ValueError, match="no 'embeddings' field"):
        CohereEmbedV4(invoker).embed_documents(["a"])


def test_a_count_mismatch_between_texts_and_vectors_fails():
    invoker = RecordingInvoker([_cohere_by_type_response(1)])

    with pytest.raises(ValueError, match="1 vectors for 2 texts"):
        CohereEmbedV4(invoker).embed_documents(["a", "b"])


def test_an_unknown_input_kind_is_refused_before_any_request_is_made():
    invoker = RecordingInvoker([])

    with pytest.raises(ValueError, match="input_kind"):
        CohereEmbedV4(invoker).embed(["a"], "search_document")
    assert invoker.calls == []


def test_empty_text_is_refused_before_any_request_is_made():
    invoker = RecordingInvoker([])

    with pytest.raises(ValueError, match="empty"):
        CohereEmbedV4(invoker).embed_documents(["a", "   "])
    assert invoker.calls == []


def test_bge_m3_runs_through_an_injected_encoder_with_no_weights_loaded():
    seen = []

    def encoder(texts):
        seen.append(list(texts))
        unit = 1.0 / math.sqrt(1024)
        return [[unit] * 1024 for _ in texts]

    provider = BgeM3Local(encoder=encoder, batch_size=2)
    batch = provider.embed_queries(["a", "b", "c"])

    assert seen == [["a", "b"], ["c"]]
    assert batch.request_count == 2
    assert len(batch.vectors) == 3
    assert batch.vectors[0].input_kind == INPUT_QUERY
    assert batch.vectors[0].model_id == "BAAI/bge-m3"
    # Injected encoder: we did not set normalize_embeddings, so we do not claim it.
    assert batch.vectors[0].normalized is None


def test_registry_builds_every_provider_without_touching_an_sdk():
    assert set(provider_names()) == {TITAN_V2, COHERE_V4, BGE_M3}

    titan = build_provider(TITAN_V2, invoker=RecordingInvoker([]))
    cohere = build_provider(COHERE_V4, invoker=RecordingInvoker([]))
    local = build_provider(BGE_M3)

    assert (titan.dimensions, cohere.dimensions, local.dimensions) == (
        1024,
        1024,
        1024,
    )
    assert titan.max_batch_size == 1
    assert cohere.max_batch_size == 96


def test_registry_rejects_an_unknown_provider_name():
    with pytest.raises(ValueError, match="unknown embedding provider"):
        build_provider("text-embedding-3-small")


class FakeBotoClient:
    """The shape boto3's bedrock-runtime client returns: a streaming body."""

    def __init__(self, response_body):
        self._response_body = response_body
        self.kwargs = None

    def invoke_model(self, **kwargs):
        self.kwargs = kwargs
        return {"body": io.BytesIO(json.dumps(self._response_body).encode())}


def test_boto3_invoker_serialises_the_request_and_reads_the_streamed_body():
    client = FakeBotoClient({"embedding": [0.5]})
    invoker = Boto3BedrockInvoker(region="us-east-1", client=client)

    body = invoker.invoke_json("some.model", {"inputText": "à"}, accept="*/*")

    assert body == {"embedding": [0.5]}
    assert client.kwargs["modelId"] == "some.model"
    assert client.kwargs["contentType"] == "application/json"
    assert client.kwargs["accept"] == "*/*"
    # Vietnamese must survive the round trip as characters, not \\u escapes
    # the model would then embed literally.
    assert json.loads(client.kwargs["body"]) == {"inputText": "à"}


def test_probe_measures_the_l2_norm_rather_than_trusting_the_docs():
    unit = 1.0 / math.sqrt(4)
    assert probe._l2_norm([unit] * 4) == pytest.approx(1.0)
    assert probe._l2_norm([3.0, 4.0]) == pytest.approx(5.0)


def test_probe_reports_a_vector_without_raising(capsys):
    vector = EmbeddingVector(
        values=[0.5, 0.5, 0.5, 0.5],
        text_sha256=text_digest("x"),
        provider=TITAN_V2,
        model_id="amazon.titan-embed-text-v2:0",
        dimensions=4,
        input_kind=INPUT_DOCUMENT,
        normalized=True,
        input_token_count=3,
    )

    probe._report(vector, latency_ms=12.5, requests=1)

    out = capsys.readouterr().out
    assert "amazon.titan-embed-text-v2:0" in out
    assert "measured L2 norm: 1.000000" in out