"""Embedding 类型/端口/transport 测试(M2 设计 §7;T8)。

蓝本审计: GovDoc retrieval/embedding.py(分批/index 排序/维度校验)与
VT adapters/embedding.py(归一化);库裁决见设计 §7.3 表。
"""

import contextlib
import dataclasses
import json

import httpx
import pytest
from loguru import logger

from polygateway.errors import (
    RequestRejectedError,
    ResultInvalidError,
    SourceDeadError,
    TransientError,
)
from polygateway.ports import EmbeddingTransport
from polygateway.transports.openai_compat import OpenAICompatTransport
from polygateway.types import EmbeddingResponse, EmbeddingTransportResult, SourceConfig


def _src(**overrides):
    base = {
        "name": "e1",
        "provider": "openai",
        "base_url": "https://gw.example/v1",
        "api_key": "sk",
        "model": "embed-1",
        "timeout_s": 10.0,
        "est_tokens": 7,
    }
    base.update(overrides)
    return SourceConfig(**base)


class TestTypes:
    def test_embedding_response_frozen_with_defaults(self):
        resp = EmbeddingResponse(
            vectors=[[0.1, 0.2]],
            dim=2,
            model="m",
            provider="p",
            prompt_tokens=3,
            usage_source="measured",
            latency_ms=10,
            call_id="c",
            source_name="e1",
        )
        assert resp.cost is None
        with pytest.raises(dataclasses.FrozenInstanceError):
            resp.dim = 3

    def test_transport_result_frozen(self):
        r = EmbeddingTransportResult(
            vectors=[[1.0]], dim=1, prompt_tokens=1, usage_source="measured", raw={}
        )
        with pytest.raises(dataclasses.FrozenInstanceError):
            r.dim = 2


class _DummyEmbedTransport:
    async def embed(self, *, texts, source, call_id):
        raise NotImplementedError


def test_embedding_transport_protocol_runtime_checkable():
    assert isinstance(_DummyEmbedTransport(), EmbeddingTransport)
    assert isinstance(OpenAICompatTransport(), EmbeddingTransport)


def _transport_with(handler):
    return OpenAICompatTransport(
        client_factory=lambda source: httpx.AsyncClient(transport=httpx.MockTransport(handler))
    )


def _ok_body(vectors, *, usage=None, shuffle=False):
    data = [{"index": i, "embedding": v} for i, v in enumerate(vectors)]
    if shuffle:
        data = list(reversed(data))
    body = {"data": data}
    if usage is not None:
        body["usage"] = usage
    return body


class TestEmbedTransport:
    async def test_sorts_by_index_and_measures_usage(self):
        def handler(request):
            assert request.url.path.endswith("/embeddings")
            payload = json.loads(request.content)
            assert payload == {"model": "embed-1", "input": ["a", "b"]}
            return httpx.Response(
                200,
                json=_ok_body([[1.0, 0.0], [0.0, 1.0]], usage={"prompt_tokens": 5}, shuffle=True),
            )

        result = await _transport_with(handler).embed(texts=["a", "b"], source=_src(), call_id="c")
        assert result.vectors == [[1.0, 0.0], [0.0, 1.0]]  # 乱序响应按 index 重排
        assert result.dim == 2
        assert result.prompt_tokens == 5 and result.usage_source == "measured"

    async def test_missing_usage_is_unavailable(self):
        """usage 缺失不再退到 `est_tokens`(夹具填 7),与 chat 同口径记 0 + unavailable。"""

        def handler(request):
            return httpx.Response(200, json=_ok_body([[1.0]]))

        result = await _transport_with(handler).embed(texts=["a"], source=_src(), call_id="c")
        assert result.prompt_tokens == 0 and result.usage_source == "unavailable"

    @pytest.mark.parametrize(
        ("status", "exc_type"),
        [(401, SourceDeadError), (400, RequestRejectedError), (500, TransientError)],
    )
    async def test_http_errors_translate(self, status, exc_type):
        def handler(request):
            return httpx.Response(status, text="boom")

        with pytest.raises(exc_type):
            await _transport_with(handler).embed(texts=["a"], source=_src(), call_id="c")

    async def test_network_error_is_transient(self):
        def handler(request):
            raise httpx.ConnectError("refused")

        with pytest.raises(TransientError):
            await _transport_with(handler).embed(texts=["a"], source=_src(), call_id="c")

    @pytest.mark.parametrize(
        "body",
        [
            {"data": []},  # 空 data
            {"data": [{"index": 0, "embedding": [1.0]}]},  # 数量与输入不符(输入 2 条)
            {
                "data": [
                    {"index": 0, "embedding": [1.0, 2.0]},
                    {"index": 1, "embedding": [1.0]},  # 维度不一致
                ]
            },
            {"nope": True},  # 缺 data
        ],
    )
    async def test_malformed_payload_is_result_invalid(self, body):
        def handler(request):
            return httpx.Response(200, json=body)

        with pytest.raises(ResultInvalidError):
            await _transport_with(handler).embed(texts=["a", "b"], source=_src(), call_id="c")

    async def test_empty_texts_rejected(self):
        with pytest.raises(ValueError):
            await _transport_with(lambda r: None).embed(texts=[], source=_src(), call_id="c")


# ═══════════ T9: EmbeddingClient 治理循环 ═══════════

import asyncio  # noqa: E402

from polygateway.backends.memory.breaker import InMemoryGate  # noqa: E402
from polygateway.backends.memory.limiter import InMemoryLimiter  # noqa: E402
from polygateway.config import EmbeddingSettings  # noqa: E402
from polygateway.embedding import EmbeddingClient  # noqa: E402
from polygateway.pricing import ModelPrice, PricingTable  # noqa: E402
from polygateway.sources import RoundRobinSelector  # noqa: E402
from polygateway.types import (  # noqa: E402
    BackpressurePolicy,
    BreakerConfig,
    GlobalLimits,
    RetryPolicy,
)
from tests.contracts.conftest import FakeClock  # noqa: E402
from tests.unit.test_backpressure import BoundedSleep  # noqa: E402

_BREAKER = BreakerConfig(fail_threshold=3, cooldown_s=60.0, probe_ttl_s=120.0)
_NO_GLOBAL = GlobalLimits(max_concurrency=0, rpm=0, tpm=0)


def _vec_for(texts):
    """确定性向量: 每条 text 一个 [len(text)] 一维向量,便于断言保序。"""
    return EmbeddingTransportResult(
        vectors=[[float(len(t))] for t in texts],
        dim=1,
        prompt_tokens=len(texts),
        usage_source="measured",
        raw={},
    )


class ScriptedEmbedTransport:
    """按脚本响应: 条目为 Exception / "ok"(按输入生成) / EmbeddingTransportResult / "hang"。"""

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

    async def embed(self, *, texts, source, call_id):
        self.calls.append((source.name, list(texts), call_id))
        action = self.script.pop(0)
        if isinstance(action, Exception):
            raise action
        if action == "hang":
            await asyncio.Event().wait()
        if action == "ok":
            return _vec_for(texts)
        return action


class _ClockAdvancingEmbedTransport:
    """按脚本 [(推进秒数, 动作), ...] 执行: 在一次尝试内部推进时钟(issue #8)。

    动作语义同 `ScriptedEmbedTransport`。stall 口径要区分"时间花在哪",
    故必须能让时钟只在 transport 内前进。
    """

    def __init__(self, script, clock):
        self.script = list(script)
        self.clock = clock
        self.calls = []

    async def embed(self, *, texts, source, call_id):
        self.calls.append((source.name, list(texts), call_id))
        advance, action = self.script.pop(0)
        self.clock.advance(advance)
        if isinstance(action, Exception):
            raise action
        if action == "hang":
            await asyncio.Event().wait()
        if action == "ok":
            return _vec_for(texts)
        return action


class _MemoryRecorder:
    def __init__(self):
        self.rows = []

    async def record_llm_call(self, **fields):
        self.rows.append(fields)


def _embed_client(sources, script, *, batch_size=2, telemetry=None, **overrides):
    limiter = InMemoryLimiter(
        scope="embed",
        sources={s.name: s for s in sources},
        global_limits=_NO_GLOBAL,
        lease_ttl_s=100.0,
    )
    kwargs = {
        "scope": "embed",
        "sources": sources,
        "selector": RoundRobinSelector(),
        "limiter": limiter,
        "breaker": InMemoryGate(config=_BREAKER),
        "transport": ScriptedEmbedTransport(script),
        "retry": RetryPolicy(max_attempts=3, backoff_base_s=0.001, backoff_max_s=0.01),
        "backpressure": BackpressurePolicy(stall_window_s=300.0, poll_interval_s=0.001),
        "batch_size": batch_size,
        "telemetry": telemetry,
    }
    kwargs.update(overrides)
    client = EmbeddingClient(**kwargs)
    return client, limiter


class TestEmbedBatching:
    async def test_batches_sequential_and_order_preserved(self):
        texts = ["a", "bb", "ccc", "dddd", "eeeee"]
        client, _ = _embed_client([_src()], ["ok", "ok", "ok"], batch_size=2)
        resp = await client.embed(texts)
        transport = client._transport
        assert [len(batch) for _, batch, _ in transport.calls] == [2, 2, 1]
        assert resp.vectors == [[1.0], [2.0], [3.0], [4.0], [5.0]]  # 全批拼接保序
        assert resp.prompt_tokens == 5 and resp.dim == 1

    async def test_empty_input_short_circuits(self):
        client, _ = _embed_client([_src()], [])
        resp = await client.embed([])
        assert resp.vectors == [] and resp.prompt_tokens == 0
        assert client._transport.calls == []

    async def test_usage_source_aggregates_conservatively(self):
        estimated = EmbeddingTransportResult(
            vectors=[[1.0], [1.0]], dim=1, prompt_tokens=9, usage_source="estimated", raw={}
        )
        client, _ = _embed_client([_src()], ["ok", estimated], batch_size=2)
        resp = await client.embed(["a", "b", "c", "d"])
        assert resp.usage_source == "estimated"  # 任一批 estimated 则整体 estimated
        assert resp.prompt_tokens == 2 + 9

    async def test_unavailable_batch_dominates_and_voids_cost(self):
        """三态合并优先级(设计 §3.2 #10/#11): 任一批不可得 → 整体不可得且 cost NULL。

        改前二值合并只看 `estimated`,measured+unavailable 会误标 measured;
        `_total_cost` 逐批求和还会给出一个偏低却看似有效的金额。
        """
        estimated = EmbeddingTransportResult(
            vectors=[[1.0], [1.0]], dim=1, prompt_tokens=9, usage_source="estimated", raw={}
        )
        unavailable = EmbeddingTransportResult(
            vectors=[[1.0], [1.0]], dim=1, prompt_tokens=0, usage_source="unavailable", raw={}
        )
        client, _ = _embed_client(
            [_src()],
            ["ok", estimated, unavailable],
            batch_size=2,
            pricing=PricingTable({"embed-1": ModelPrice(input_per_1m=1.0, output_per_1m=0.0)}),
        )
        resp = await client.embed(["a", "b", "c", "d", "e", "f"])
        assert resp.usage_source == "unavailable"  # unavailable 压过 estimated 与 measured
        assert resp.cost is None


class TestEmbedPostProcess:
    async def test_normalize_l2(self):
        raw = EmbeddingTransportResult(
            vectors=[[3.0, 4.0]], dim=2, prompt_tokens=1, usage_source="measured", raw={}
        )
        client, _ = _embed_client([_src()], [raw], normalize=True)
        resp = await client.embed(["x"])
        assert resp.vectors[0] == pytest.approx([0.6, 0.8])

    async def test_zero_vector_normalize_no_nan(self):
        raw = EmbeddingTransportResult(
            vectors=[[0.0, 0.0]], dim=2, prompt_tokens=1, usage_source="measured", raw={}
        )
        client, _ = _embed_client([_src()], [raw], normalize=True)
        resp = await client.embed(["x"])
        assert resp.vectors[0] == [0.0, 0.0]  # max(norm, 1e-12) 防除零(VT 语义)

    async def test_expected_dim_violation_is_result_invalid(self):
        client, _ = _embed_client([_src()], ["ok"], expected_dim=768)
        with pytest.raises(ResultInvalidError):
            await client.embed(["x"])


class TestEmbedGovernance:
    async def test_transient_retries_then_succeeds(self):
        client, _ = _embed_client(
            [_src()], [TransientError("boom", status_code=500), "ok"], batch_size=8
        )
        resp = await client.embed(["a", "b"])
        assert resp.vectors == [[1.0], [1.0]]
        assert len(client._transport.calls) == 2

    async def test_source_dead_switches_source(self):
        s1, s2 = _src(name="e1"), _src(name="e2")
        client, _ = _embed_client(
            [s1, s2], [SourceDeadError("401", status_code=401), "ok"], batch_size=8
        )
        await client.embed(["a"])
        assert [name for name, _, _ in client._transport.calls] == ["e1", "e2"]

    async def test_cancel_releases_permit(self):
        client, limiter = _embed_client([_src(max_concurrency=1)], ["hang"])
        task = asyncio.create_task(client.embed(["a"]))
        while not (await limiter.source_stats("e1")).inflight:
            await asyncio.sleep(0.01)
        task.cancel()
        with pytest.raises(asyncio.CancelledError):
            await task
        assert (await limiter.source_stats("e1")).inflight == 0

    async def test_single_timeout_does_not_exhaust_stall_budget(self):
        """issue #8: 一次耗满 timeout 的尝试不得吃掉 stall 预算而使重试失效。

        embedding 只有 `_on_no_runnable` 一处 stall 判定, 故失效链条是
        "先超时一次(墙钟耗尽) → 再遇到无可用源 → 判死"。此处正是这条路径。
        """
        clock = FakeClock()
        limiter = held = None  # 闭包延迟求值: client 建好后才有 limiter

        async def toggle_permit(n):
            """首次退避占满 permit, 迫使下一轮走 _on_no_runnable; 之后放行。"""
            nonlocal held
            if n == 1:
                held = await limiter.try_acquire("e1", 0)
            else:
                await held.release()

        # 第一次尝试耗满 300s 超时失败, 随后被迫走一轮 _on_no_runnable——
        # stall 判定就在那里, 检验它有没有把这 300s 生产性时间算进 stall 账
        transport = _ClockAdvancingEmbedTransport(
            [(300.1, TransientError("timeout", status_code=504)), (0.0, "ok")], clock
        )
        client, limiter = _embed_client(
            [_src(max_concurrency=1)],
            [],
            now=clock,
            transport=transport,
            sleep=BoundedSleep(toggle_permit),
        )
        resp = await client.embed(["a"])
        assert resp.vectors == [[1.0]]
        assert len(transport.calls) == 2  # 第二次尝试确实发出了


class TestEmbedTelemetry:
    async def test_per_batch_rows_with_digest(self):
        rec = _MemoryRecorder()
        client, _ = _embed_client([_src()], ["ok", "ok"], batch_size=1, telemetry=rec)
        await client.embed(["hello", "x" * 5000], session_id="sess", parent_call_id="pc")
        assert len(rec.rows) == 2  # 每批一行
        row = rec.rows[0]
        assert row["session_id"] == "sess" and row["parent_call_id"] == "pc"
        assert row["completion_tokens"] == 0
        assert row["response"] == "<vectors n=1 dim=1>"  # 向量绝不入库
        assert len(rec.rows[1]["messages"]) < 1000  # 长文本截断后入库


class TestEmbedCallerDimensions:
    """issue #11: 调用方自定义维度必须沿 embed 链四层透传到每一行遥测。"""

    async def test_single_batch_row_carries_dimensions(self):
        rec = _MemoryRecorder()
        client, _ = _embed_client([_src()], ["ok"], batch_size=2, telemetry=rec)
        await client.embed(["a"], tenant_id="t1", meta={"batch": "b-42"})
        assert rec.rows[0]["tenant_id"] == "t1"
        assert rec.rows[0]["meta"] == '{"batch": "b-42"}'

    async def test_every_batch_row_carries_the_same_dimensions(self):
        """维度属于本次 `embed()` 调用,不随批次变化。

        只断言首行会漏掉"只有第一批带维度"的实现——那正是逐层透传最容易漏的形态。
        """
        rec = _MemoryRecorder()
        client, _ = _embed_client([_src()], ["ok", "ok", "ok"], batch_size=1, telemetry=rec)
        await client.embed(["a", "b", "c"], tenant_id="t1", meta={"batch": "b-42"})
        assert len(rec.rows) == 3  # 切成三批,每批一行
        assert [r["tenant_id"] for r in rec.rows] == ["t1", "t1", "t1"]
        assert [r["meta"] for r in rec.rows] == ['{"batch": "b-42"}'] * 3

    async def test_invalid_meta_rejected_before_any_telemetry(self):
        """校验在切批之前: 遥测层的失败都被降级成 warning,放下游等于没有校验。"""
        rec = _MemoryRecorder()
        client, _ = _embed_client([_src()], ["ok"], batch_size=2, telemetry=rec)
        with pytest.raises(ValueError, match="meta"):
            await client.embed(["a"], meta={"Bad Key": 1})
        assert rec.rows == []
        assert client._transport.calls == []  # 连调用都没发出

    async def test_defaults_land_as_sentinels(self):
        rec = _MemoryRecorder()
        client, _ = _embed_client([_src()], ["ok"], batch_size=2, telemetry=rec)
        await client.embed(["a"])
        assert rec.rows[0]["tenant_id"] == ""  # 空串哨兵,不是 None
        assert rec.rows[0]["meta"] == "{}"


@contextlib.contextmanager
def _captured_warnings():
    """捕获库发出的 WARNING;loguru 不经标准 logging,pytest 的 caplog 抓不到。"""
    messages: list[str] = []
    sink_id = logger.add(messages.append, level="WARNING")
    try:
        yield messages
    finally:
        logger.remove(sink_id)


class TestExtraBodyStripped:
    """issue #4 决策 G: embedding 路径不消费 extra_body,剥离并 warning。"""

    async def test_stripped_with_warning_but_assembly_succeeds(self):
        """报错会让下游整个装配起不来,而这条路径本无采样语义(人类拍板)。"""
        with _captured_warnings() as warnings:
            client, _ = _embed_client([_src(extra_body={"temperature": 0})], ["ok"])
        assert client._sources[0].extra_body == {}
        assert any("extra_body" in m for m in warnings)
        assert any("dimensions" in m for m in warnings)  # 文案须指路,不能只说不支持
        await client.embed(["hi"])  # 装配后可正常工作

    async def test_telemetry_never_records_a_parameter_that_was_not_sent(self):
        """剥离的真正理由: embed payload 硬编码 {model, input},不剥离则审计表

        会显示这次调用带了 temperature=0——那是数据造假,比参数失效更坏。
        """
        rec = _MemoryRecorder()
        client, _ = _embed_client([_src(extra_body={"temperature": 0})], ["ok"], telemetry=rec)
        await client.embed(["hi"])
        assert rec.rows[0]["sampling"] is None

    async def test_no_warning_without_extra_body(self):
        with _captured_warnings() as warnings:
            client, _ = _embed_client([_src()], ["ok"])
        assert client._sources[0].extra_body == {}
        assert not [m for m in warnings if "extra_body" in m]


class TestEmbeddingSettings:
    _ENV = {
        "EMBED__QWEN__1__BASE_URL": "https://gw.example/v1",
        "EMBED__QWEN__1__API_KEY": "sk-a",
        "EMBED__QWEN__1__MODEL": "text-embedding-v3",
        "EMBED__QWEN__1__TIMEOUT_S": "60",
        "EMBED__RETRY__MAX_ATTEMPTS": "3",
        "EMBED__RETRY__BACKOFF_BASE_S": "1.0",
        "EMBED__RETRY__BACKOFF_MAX_S": "10.0",
        "EMBED__BREAKER__FAIL_THRESHOLD": "5",
        "EMBED__BREAKER__COOLDOWN_S": "60",
        "PGW_CACHE_BACKEND": "none",
        "PGW_TELEMETRY_BACKEND": "none",
        "EMBED__BATCH_SIZE": "64",
    }

    def test_loads_scope_and_batch(self):
        s = EmbeddingSettings.from_env("EMBED", env=self._ENV)
        assert s.gateway.sources[0].model == "text-embedding-v3"
        assert s.batch_size == 64 and s.normalize is False and s.expected_dim is None

    def test_batch_size_required_and_positive(self):
        env = {k: v for k, v in self._ENV.items() if k != "EMBED__BATCH_SIZE"}
        with pytest.raises(ValueError, match="BATCH_SIZE"):
            EmbeddingSettings.from_env("EMBED", env=env)
        with pytest.raises(ValueError, match="BATCH_SIZE"):
            EmbeddingSettings.from_env("EMBED", env={**self._ENV, "EMBED__BATCH_SIZE": "0"})

    def test_optional_normalize_and_dim(self):
        env = {**self._ENV, "EMBED__NORMALIZE": "true", "EMBED__EXPECTED_DIM": "768"}
        s = EmbeddingSettings.from_env("EMBED", env=env)
        assert s.normalize is True and s.expected_dim == 768

    def test_expected_dim_must_be_positive(self):
        """env 层的检查保留是为了报错能点出键名(构造期那道点的是字段名)。"""
        with pytest.raises(ValueError, match="EXPECTED_DIM"):
            EmbeddingSettings.from_env("EMBED", env={**self._ENV, "EMBED__EXPECTED_DIM": "0"})

    def test_from_settings_assembles_client(self):
        s = EmbeddingSettings.from_env("EMBED", env=self._ENV)
        client = EmbeddingClient.from_settings(s)
        assert isinstance(client, EmbeddingClient)
