# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd.
# SPDX-License-Identifier: AGPL-3.0

from __future__ import annotations

import asyncio
import json
from types import SimpleNamespace

import pytest

from openviking.models.embedder.base import EmbedResult, embed_compat
from openviking.models.vlm.base import VLMBase
from openviking.observability.context import (
    bind_operation_observability_context,
    bind_root_observability_context,
    get_operation_observability_context,
    get_root_observability_context,
    reset_operation_observability_context,
    reset_root_observability_context,
)
from openviking.server.identity import RequestContext, Role
from openviking.service.resource_service import ResourceService
from openviking.storage.collection_schemas import TextEmbeddingHandler
from openviking.storage.queuefs.semantic_dag import DagStats
from openviking.storage.queuefs.semantic_msg import SemanticMsg
from openviking.storage.queuefs.semantic_processor import SemanticProcessor
from openviking.telemetry import (
    get_current_telemetry,
    get_telemetry_runtime,
    register_telemetry,
    tracer_module,
    unregister_telemetry,
)
from openviking.telemetry.backends.memory import MemoryOperationTelemetry
from openviking.telemetry.context import bind_telemetry, bind_telemetry_stage
from openviking.telemetry.snapshot import TelemetrySnapshot
from openviking.telemetry.span_models import OperationSpanAttributes, RootSpanAttributes
from openviking_cli.session.user_id import UserIdentifier
from openviking_cli.utils import logger as logger_module


def test_telemetry_module_exports_snapshot_and_runtime():
    snapshot = TelemetrySnapshot(
        telemetry_id="tm_demo",
        summary={"duration_ms": 1.2},
    )
    usage = snapshot.to_usage_dict()

    assert usage == {"duration_ms": 1.2, "token_total": 0}
    assert get_telemetry_runtime().meter() is not None


def test_root_observability_context_bind_and_reset():
    root = RootSpanAttributes(http_method="GET", http_route="/demo", request_id="req-1")
    token = bind_root_observability_context(root)
    try:
        assert get_root_observability_context() is root
    finally:
        reset_root_observability_context(token)
    assert get_root_observability_context() is None


def test_operation_observability_context_bind_and_reset():
    operation = OperationSpanAttributes(operation="search.find", telemetry_id="tm-demo")
    token = bind_operation_observability_context(operation)
    try:
        assert get_operation_observability_context() is operation
    finally:
        reset_operation_observability_context(token)
    assert get_operation_observability_context() is None


def test_telemetry_snapshot_to_dict_supports_summary_only():
    snapshot = TelemetrySnapshot(
        telemetry_id="tm_demo",
        summary={"duration_ms": 1.2, "tokens": {"total": 3}},
    )

    payload = snapshot.to_dict(include_summary=True)

    assert payload == {
        "id": "tm_demo",
        "summary": {"duration_ms": 1.2, "tokens": {"total": 3}},
    }


def test_telemetry_summary_breaks_down_llm_and_embedding_token_usage():
    telemetry = MemoryOperationTelemetry(operation="resources.add_resource", enabled=True)
    telemetry.record_token_usage("llm", 11, 7)
    telemetry.record_token_usage("embedding", 13, 0)

    summary = telemetry.finish().summary
    assert telemetry.telemetry_id
    assert telemetry.telemetry_id.startswith("tm_")
    assert summary["tokens"]["total"] == 31
    assert summary["duration_ms"] >= 0
    assert summary["tokens"]["llm"] == {
        "input": 11,
        "output": 7,
        "total": 18,
    }
    assert summary["tokens"]["embedding"] == {"total": 13}
    assert "queue" not in summary
    assert "vector" not in summary
    assert "semantic_nodes" not in summary
    assert "memory" not in summary
    assert "errors" not in summary


def test_telemetry_summary_breaks_down_stage_token_usage():
    telemetry = MemoryOperationTelemetry(operation="search.find", enabled=True)
    telemetry.record_token_usage("embedding", 11, 0, stage="embed_query")
    telemetry.record_token_usage("rerank", 7, 0, stage="rerank")
    telemetry.record_token_usage("llm", 5, 3, stage="vlm")

    summary = telemetry.finish().summary

    assert summary["tokens"]["total"] == 26
    assert summary["tokens"]["rerank"] == {"total": 7}
    assert summary["tokens"]["stages"]["embed_query"]["embedding"] == {"total": 11}
    assert summary["tokens"]["stages"]["rerank"]["rerank"] == {"total": 7}
    assert summary["tokens"]["stages"]["vlm"]["llm"] == {
        "input": 5,
        "output": 3,
        "total": 8,
    }


@pytest.mark.asyncio
async def test_bind_telemetry_stage_propagates_across_async_tasks():
    telemetry = MemoryOperationTelemetry(operation="resource.process", enabled=True)

    async def _worker() -> None:
        await asyncio.sleep(0)
        get_current_telemetry().add_token_usage(6, 4)

    with bind_telemetry(telemetry):
        with bind_telemetry_stage("resource_summarize"):
            await asyncio.create_task(_worker())

    summary = telemetry.finish().summary
    assert summary["tokens"]["stages"]["resource_summarize"]["llm"] == {
        "input": 6,
        "output": 4,
        "total": 10,
    }


@pytest.mark.asyncio
async def test_embed_compat_binds_query_stage_for_embedding_tokens():
    telemetry = MemoryOperationTelemetry(operation="search.find", enabled=True)

    class _TelemetryAwareAsyncEmbedder:
        async def embed_async(self, text: str, is_query: bool = False) -> EmbedResult:
            assert text == "hello"
            assert is_query is True
            get_current_telemetry().record_token_usage("embedding", 9, 0)
            return EmbedResult(dense_vector=[0.1, 0.2])

    with bind_telemetry(telemetry):
        await embed_compat(_TelemetryAwareAsyncEmbedder(), "hello", is_query=True)

    summary = telemetry.finish().summary
    assert summary["tokens"]["stages"]["embed_query"]["embedding"] == {"total": 9}


def test_vlm_base_defaults_operation_tokens_to_vlm_stage():
    class _DummyVLM(VLMBase):
        def get_completion(self, *args, **kwargs):
            raise NotImplementedError()

        async def get_completion_async(self, *args, **kwargs):
            raise NotImplementedError()

        def get_vision_completion(self, *args, **kwargs):
            raise NotImplementedError()

        async def get_vision_completion_async(self, *args, **kwargs):
            raise NotImplementedError()

    telemetry = MemoryOperationTelemetry(operation="session.commit", enabled=True)
    with bind_telemetry(telemetry):
        _DummyVLM({"provider": "openai", "model": "gpt-4o-mini"}).update_token_usage(
            model_name="gpt-4o-mini",
            provider="openai",
            prompt_tokens=7,
            completion_tokens=5,
        )

    summary = telemetry.finish().summary
    assert summary["tokens"]["stages"]["vlm"]["llm"] == {
        "input": 7,
        "output": 5,
        "total": 12,
    }


def test_disabled_telemetry_still_has_request_id():
    telemetry = MemoryOperationTelemetry(operation="resources.add_resource", enabled=False)

    assert telemetry.telemetry_id
    assert telemetry.telemetry_id.startswith("tm_")


def test_telemetry_summary_uses_simplified_internal_metric_keys():
    summary = MemoryOperationTelemetry(
        operation="search.find",
        enabled=True,
    )
    summary.count("vector.searches", 2)
    summary.count("vector.scored", 5)
    summary.count("vector.passed", 3)
    summary.set("vector.returned", 2)
    summary.count("vector.scanned", 5)
    summary.set("vector.scan_reason", "")
    summary.set("semantic_nodes.total", 4)
    summary.set("semantic_nodes.done", 3)
    summary.set("semantic_nodes.pending", 1)
    summary.set("semantic_nodes.running", 0)
    summary.set("memory.extracted", 6)

    result = summary.finish().summary

    assert result["vector"] == {
        "searches": 2,
        "scored": 5,
        "passed": 3,
        "returned": 2,
        "scanned": 5,
        "scan_reason": "",
    }
    assert result["semantic_nodes"] == {
        "total": 4,
        "done": 3,
        "pending": 1,
    }
    assert result["memory"] == {"extracted": 6}


def test_init_tracer_forwards_headers_to_grpc_exporter(monkeypatch):
    captured = {}

    class FakeExporter:
        def __init__(self, **kwargs):
            captured.update(kwargs)

    monkeypatch.setattr(tracer_module, "OTLPGrpcSpanExporter", FakeExporter)
    monkeypatch.setattr(tracer_module, "BatchSpanProcessor", lambda exporter, **kwargs: exporter)

    class FakeTracerProvider:
        def __init__(self, resource=None):
            self.resource = resource

        def add_span_processor(self, _processor):
            return None

    monkeypatch.setattr(tracer_module, "TracerProvider", FakeTracerProvider)
    monkeypatch.setattr(
        tracer_module,
        "Resource",
        SimpleNamespace(create=lambda attrs: attrs),
    )
    monkeypatch.setattr(
        tracer_module,
        "otel_trace",
        SimpleNamespace(
            set_tracer_provider=lambda _provider: None,
            get_tracer=lambda service_name: f"tracer:{service_name}",
        ),
    )
    monkeypatch.setattr(
        tracer_module,
        "TraceContextTextMapPropagator",
        lambda: "propagator",
    )
    monkeypatch.setattr(tracer_module, "_setup_logging", lambda: None)
    monkeypatch.setattr(tracer_module, "_init_asyncio_instrumentation", lambda: None)

    tracer_module.init_tracer(
        endpoint="apmplus-cn-beijing.ivolces.com:4317",
        service_name="memorydb",
        protocol="grpc",
        insecure=True,
        headers={"x-byteapm-appkey": "trace-appkey"},
        enabled=True,
    )

    assert captured["endpoint"] == "apmplus-cn-beijing.ivolces.com:4317"
    assert captured["insecure"] is True
    assert captured["headers"] == {"x-byteapm-appkey": "trace-appkey"}


def test_init_tracer_forwards_headers_to_http_exporter(monkeypatch):
    captured = {}

    class FakeExporter:
        def __init__(self, **kwargs):
            captured.update(kwargs)

    monkeypatch.setattr(tracer_module, "OTLPHttpSpanExporter", FakeExporter)
    monkeypatch.setattr(tracer_module, "BatchSpanProcessor", lambda exporter, **kwargs: exporter)

    class FakeTracerProvider:
        def __init__(self, resource=None):
            self.resource = resource

        def add_span_processor(self, _processor):
            return None

    monkeypatch.setattr(tracer_module, "TracerProvider", FakeTracerProvider)
    monkeypatch.setattr(
        tracer_module,
        "Resource",
        SimpleNamespace(create=lambda attrs: attrs),
    )
    monkeypatch.setattr(
        tracer_module,
        "otel_trace",
        SimpleNamespace(
            set_tracer_provider=lambda _provider: None,
            get_tracer=lambda service_name: f"tracer:{service_name}",
        ),
    )
    monkeypatch.setattr(
        tracer_module,
        "TraceContextTextMapPropagator",
        lambda: "propagator",
    )
    monkeypatch.setattr(tracer_module, "_setup_logging", lambda: None)
    monkeypatch.setattr(tracer_module, "_init_asyncio_instrumentation", lambda: None)

    tracer_module.init_tracer(
        endpoint="https://apmplus-cn-beijing.ivolces.com/api/otlp/v1/traces",
        service_name="memorydb",
        protocol="http",
        headers={"X-ByteAPM-AppKey": "trace-appkey"},
        enabled=True,
    )

    assert captured["endpoint"] == "https://apmplus-cn-beijing.ivolces.com/api/otlp/v1/traces"
    assert captured["headers"] == {"X-ByteAPM-AppKey": "trace-appkey"}


def test_init_otel_log_handler_forwards_headers_to_grpc_exporter(monkeypatch):
    captured = {}

    class FakeExporter:
        def __init__(self, **kwargs):
            captured.update(kwargs)

    class FakeLoggerProvider:
        def __init__(self, resource=None):
            self.resource = resource

        def add_log_record_processor(self, _processor):
            return None

    monkeypatch.setattr(logger_module, "_otel_log_handler_initialized", False)
    monkeypatch.setattr(logger_module, "_otel_log_handler", None)
    monkeypatch.setattr(logger_module, "OTLPGrpcLogExporter", FakeExporter)
    monkeypatch.setattr(
        logger_module,
        "BatchLogRecordProcessor",
        lambda exporter: exporter,
    )
    monkeypatch.setattr(logger_module, "LoggerProvider", FakeLoggerProvider)
    monkeypatch.setattr(
        logger_module,
        "LoggingHandler",
        lambda **kwargs: SimpleNamespace(**kwargs),
    )
    monkeypatch.setattr(
        logger_module,
        "Resource",
        SimpleNamespace(create=lambda attrs: attrs),
    )
    monkeypatch.setattr(logger_module, "set_logger_provider", lambda _provider: None)
    monkeypatch.setattr(
        logger_module,
        "get_logger",
        lambda _name: SimpleNamespace(
            info=lambda *args, **kwargs: None, warning=lambda *args, **kwargs: None
        ),
    )

    handler = logger_module.init_otel_log_handler(
        protocol="grpc",
        endpoint="apmplus-cn-beijing.ivolces.com:4317",
        service_name="memorydb",
        insecure=True,
        headers={"x-byteapm-appkey": "log-appkey"},
        enabled=True,
    )

    assert handler is not None
    assert captured["endpoint"] == "apmplus-cn-beijing.ivolces.com:4317"
    assert captured["insecure"] is True
    assert captured["headers"] == {"x-byteapm-appkey": "log-appkey"}


def test_init_otel_log_handler_forwards_headers_to_http_exporter(monkeypatch):
    captured = {}

    class FakeExporter:
        def __init__(self, **kwargs):
            captured.update(kwargs)

    class FakeLoggerProvider:
        def __init__(self, resource=None):
            self.resource = resource

        def add_log_record_processor(self, _processor):
            return None

    monkeypatch.setattr(logger_module, "_otel_log_handler_initialized", False)
    monkeypatch.setattr(logger_module, "_otel_log_handler", None)
    monkeypatch.setattr(logger_module, "OTLPHttpLogExporter", FakeExporter)
    monkeypatch.setattr(
        logger_module,
        "BatchLogRecordProcessor",
        lambda exporter: exporter,
    )
    monkeypatch.setattr(logger_module, "LoggerProvider", FakeLoggerProvider)
    monkeypatch.setattr(
        logger_module,
        "LoggingHandler",
        lambda **kwargs: SimpleNamespace(**kwargs),
    )
    monkeypatch.setattr(
        logger_module,
        "Resource",
        SimpleNamespace(create=lambda attrs: attrs),
    )
    monkeypatch.setattr(logger_module, "set_logger_provider", lambda _provider: None)
    monkeypatch.setattr(
        logger_module,
        "get_logger",
        lambda _name: SimpleNamespace(
            info=lambda *args, **kwargs: None, warning=lambda *args, **kwargs: None
        ),
    )

    handler = logger_module.init_otel_log_handler(
        protocol="http",
        endpoint="https://apmplus-cn-beijing.ivolces.com/api/otlp/v1/logs",
        service_name="memorydb",
        headers={"X-ByteAPM-AppKey": "log-appkey"},
        enabled=True,
    )

    assert handler is not None
    assert captured["endpoint"] == "https://apmplus-cn-beijing.ivolces.com/api/otlp/v1/logs"
    assert captured["headers"] == {"X-ByteAPM-AppKey": "log-appkey"}


def test_telemetry_summary_detects_groups_by_prefix_without_static_key_lists():
    telemetry = MemoryOperationTelemetry(operation="search.find", enabled=True)
    telemetry.set("vector.debug_probe", 1)
    telemetry.set("queue.semantic.processed", 2)
    telemetry.set("memory.extracted", 1)

    result = telemetry.finish().summary

    assert "vector" in result
    assert "queue" in result
    assert "memory" in result


@pytest.mark.asyncio
async def test_semantic_processor_binds_registered_operation_telemetry(monkeypatch):
    telemetry = MemoryOperationTelemetry(operation="resources.add_resource", enabled=True)
    register_telemetry(telemetry)

    processor = SemanticProcessor()

    class FakeVikingFS:
        async def ls(self, uri, ctx=None):
            return []

    class _FakeDagExecutor:
        def __init__(self, **kwargs):
            pass

        async def run(self, root_uri):
            assert get_current_telemetry() is telemetry
            get_current_telemetry().record_token_usage("llm", 11, 7)

        def get_stats(self):
            return DagStats()

    monkeypatch.setattr(
        "openviking.storage.queuefs.semantic_processor.get_viking_fs",
        lambda: FakeVikingFS(),
    )
    monkeypatch.setattr(
        "openviking.storage.queuefs.semantic_processor.SemanticDagExecutor",
        lambda **kwargs: _FakeDagExecutor(**kwargs),
    )

    try:
        await processor.on_dequeue(
            SemanticMsg(
                uri="viking://resources/demo",
                context_type="resource",
                recursive=False,
                telemetry_id=telemetry.telemetry_id,
            ).to_dict()
        )
    finally:
        unregister_telemetry(telemetry.telemetry_id)

    result = telemetry.finish()
    summary = result.summary
    assert summary["tokens"]["total"] == 18
    assert summary["tokens"]["llm"]["total"] == 18
    assert "embedding" not in summary["tokens"]


@pytest.mark.asyncio
async def test_semantic_processor_binds_metric_account_context(monkeypatch):
    processor = SemanticProcessor()
    ran = {"value": False}

    class FakeVikingFS:
        async def ls(self, uri, ctx=None):
            return []

    class _FakeDagExecutor:
        def __init__(self, **kwargs):
            pass

        async def run(self, root_uri):
            ran["value"] = True
            root_context = get_root_observability_context()
            assert root_context is not None
            assert root_context.account_id == "acct-semantic"

        def get_stats(self):
            return DagStats()

    monkeypatch.setattr(
        "openviking.storage.queuefs.semantic_processor.get_viking_fs",
        lambda: FakeVikingFS(),
    )
    monkeypatch.setattr(
        "openviking.storage.queuefs.semantic_processor.SemanticDagExecutor",
        lambda **kwargs: _FakeDagExecutor(**kwargs),
    )

    await processor.on_dequeue(
        SemanticMsg(
            uri="viking://resources/demo",
            context_type="resource",
            recursive=False,
            account_id="acct-semantic",
        ).to_dict()
    )
    assert ran["value"] is True


@pytest.mark.asyncio
async def test_embedding_handler_binds_registered_operation_telemetry(monkeypatch):
    telemetry = MemoryOperationTelemetry(operation="resources.add_resource", enabled=True)
    register_telemetry(telemetry)

    class _TelemetryAwareEmbedder:
        def embed(self, text: str, is_query: bool = False) -> EmbedResult:
            assert text == "hello"
            assert is_query is False
            get_current_telemetry().record_token_usage("embedding", 9, 0)
            return EmbedResult(dense_vector=[0.1, 0.2])

    class _DummyConfig:
        def __init__(self):
            self.storage = SimpleNamespace(vectordb=SimpleNamespace(name="context"))
            self.embedding = SimpleNamespace(
                dimension=2,
                get_embedder=lambda: _TelemetryAwareEmbedder(),
                circuit_breaker=SimpleNamespace(
                    failure_threshold=5,
                    reset_timeout=300.0,
                    max_reset_timeout=300.0,
                ),
            )

    class _DummyVikingDB:
        is_closing = False

        async def upsert(self, _data, *, ctx=None):
            return "rec-1"

    monkeypatch.setattr(
        "openviking_cli.utils.config.get_openviking_config",
        lambda: _DummyConfig(),
    )

    handler = TextEmbeddingHandler(_DummyVikingDB())
    payload = {
        "data": json.dumps(
            {
                "id": "msg-1",
                "message": "hello",
                "telemetry_id": telemetry.telemetry_id,
                "context_data": {
                    "id": "id-1",
                    "uri": "viking://resources/sample",
                    "account_id": "default",
                    "abstract": "sample",
                },
            }
        )
    }

    try:
        await handler.on_dequeue(payload)
    finally:
        unregister_telemetry(telemetry.telemetry_id)

    result = telemetry.finish()
    summary = result.summary
    assert summary["tokens"]["embedding"] == {"total": 9}


@pytest.mark.asyncio
async def test_resource_service_add_resource_reports_queue_summary(monkeypatch):
    telemetry = MemoryOperationTelemetry(operation="resources.add_resource", enabled=True)
    queue_status = {
        "Semantic": {
            "processed": 2,
            "requeue_count": 0,
            "error_count": 1,
            "errors": [],
        },
        "Embedding": {
            "processed": 5,
            "requeue_count": 0,
            "error_count": 0,
            "errors": [],
        },
    }

    class _DummyProcessor:
        async def process_resource(self, **kwargs):
            return {
                "status": "success",
                "root_uri": "viking://resources/demo",
            }

    class _DummyRequestWaitTracker:
        def register_request(self, telemetry_id: str) -> None:
            del telemetry_id

        async def wait_for_request(self, telemetry_id: str, timeout=None) -> None:
            del telemetry_id, timeout

        def build_queue_status(self, telemetry_id: str):
            del telemetry_id
            return queue_status

        def cleanup(self, telemetry_id: str) -> None:
            del telemetry_id

    monkeypatch.setattr(
        "openviking.service.resource_service.get_request_wait_tracker",
        lambda: _DummyRequestWaitTracker(),
        raising=False,
    )

    class _DagStats:
        total_nodes = 3
        done_nodes = 2
        pending_nodes = 1
        in_progress_nodes = 0

    monkeypatch.setattr(
        "openviking.storage.queuefs.semantic_processor.SemanticProcessor.consume_dag_stats",
        classmethod(lambda cls, telemetry_id="", uri=None: _DagStats()),
    )

    service = ResourceService(
        vikingdb=object(),
        viking_fs=object(),
        resource_processor=_DummyProcessor(),
        skill_processor=object(),
    )
    ctx = RequestContext(user=UserIdentifier.the_default_user(), role=Role.ROOT)

    with bind_telemetry(telemetry):
        result = await service.add_resource(path="/tmp/demo.md", ctx=ctx, wait=True)

    assert result["root_uri"] == "viking://resources/demo"
    telemetry_result = telemetry.finish()
    summary = telemetry_result.summary
    assert summary["queue"] == {
        "semantic": {"processed": 2, "error_count": 1},
        "embedding": {"processed": 5},
    }
    assert summary["semantic_nodes"] == {
        "total": 3,
        "done": 2,
        "pending": 1,
    }
    assert "memory" not in summary
    assert "errors" not in summary


def test_telemetry_summary_includes_only_memory_group_when_memory_metrics_exist():
    telemetry = MemoryOperationTelemetry(operation="session.commit", enabled=True)
    telemetry.record_token_usage("llm", 5, 3)
    telemetry.set("memory.extracted", 4)

    summary = telemetry.finish().summary

    assert summary["memory"] == {"extracted": 4}
    assert "queue" not in summary
    assert "vector" not in summary
    assert "semantic_nodes" not in summary
    assert "errors" not in summary
