import contextlib
import json
import queue
from datetime import datetime, timedelta
from types import SimpleNamespace
from unittest.mock import MagicMock

import pytest

from core.ops.ops_trace_manager import (
    OpsTraceManager,
    TraceQueueManager,
    TraceTask,
    TraceTaskName,
)


class DummyConfig:
    def __init__(self, **kwargs):
        self._data = kwargs

    def model_dump(self):
        return dict(self._data)


class DummyTraceInstance:
    instances: list["DummyTraceInstance"] = []

    def __init__(self, config):
        self.config = config
        DummyTraceInstance.instances.append(self)

    def api_check(self):
        return True

    def get_project_key(self):
        return "fake-key"

    def get_project_url(self):
        return "https://project.fake"


FAKE_PROVIDER_ENTRY = {
    "config_class": DummyConfig,
    "secret_keys": ["secret_value"],
    "other_keys": ["other_value"],
    "trace_instance": DummyTraceInstance,
}


class FakeProviderMap:
    def __init__(self, data):
        self._data = data

    def __getitem__(self, key):
        if key in self._data:
            return self._data[key]
        raise KeyError(f"Unsupported tracing provider: {key}")


class DummyTimer:
    def __init__(self, interval, function):
        self.interval = interval
        self.function = function
        self.name = ""
        self.daemon = False
        self.started = False

    def start(self):
        self.started = True

    def is_alive(self):
        return False


class FakeMessageFile:
    def __init__(self):
        self.url = "path/to/file"
        self.id = "file-id"
        self.type = "document"
        self.created_by_role = "role"
        self.created_by = "user"


def make_message_data(**overrides):
    created_at = datetime(2025, 2, 20, 12, 0, 0)
    base = {
        "id": "msg-id",
        "app_id": "app-id",
        "conversation_id": "conv-id",
        "created_at": created_at,
        "updated_at": created_at + timedelta(seconds=3),
        "message": "hello",
        "provider_response_latency": 1,
        "message_tokens": 5,
        "answer_tokens": 7,
        "answer": "world",
        "error": "",
        "status": "complete",
        "model_provider": "provider",
        "model_id": "model",
        "from_end_user_id": "end-user",
        "from_account_id": "account",
        "agent_based": False,
        "workflow_run_id": "workflow-run",
        "from_source": "source",
        "message_metadata": json.dumps({"usage": {"time_to_first_token": 1, "time_to_generate": 2}}),
        "agent_thoughts": [],
        "query": "sample-query",
        "inputs": "sample-input",
    }
    base.update(overrides)

    class MessageData:
        def __init__(self, data):
            self.__dict__.update(data)

        def to_dict(self):
            return dict(self.__dict__)

    return MessageData(base)


def make_agent_thought(tool_name, created_at):
    return SimpleNamespace(
        tools=[tool_name],
        created_at=created_at,
        tool_meta={
            tool_name: {
                "tool_config": {"foo": "bar"},
                "time_cost": 5,
                "error": "",
                "tool_parameters": {"x": 1},
            }
        },
    )


def make_workflow_run():
    return SimpleNamespace(
        workflow_id="wf-1",
        tenant_id="tenant",
        id="run-id",
        elapsed_time=10,
        status="finished",
        inputs_dict={"sys.file": ["f1"], "query": "search"},
        outputs_dict={"out": "value"},
        version="3",
        error=None,
        total_tokens=12,
        workflow_run_id="run-id",
        created_at=datetime(2025, 2, 20, 10, 0, 0),
        finished_at=datetime(2025, 2, 20, 10, 0, 5),
        triggered_from="user",
        app_id="app-id",
        to_dict=lambda self=None: {"run": "value"},
    )


def configure_db_scalar(session, *, message_file=None, workflow_app_log=None):
    """Configure session.scalar to return appropriate values for MessageFile/WorkflowAppLog lookups."""
    original_scalar = session.scalar

    def _side_effect(stmt):
        stmt_str = str(stmt)
        if "message_file" in stmt_str.lower():
            return message_file
        if "workflow_app_log" in stmt_str.lower():
            return workflow_app_log
        return original_scalar(stmt)

    session.scalar.side_effect = _side_effect


class DummySessionContext:
    scalar_values = []

    def __init__(self, engine):
        self._values = list(self.scalar_values)
        self._index = 0

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        return False

    def execute(self, *args, **kwargs):
        return self

    def scalar(self, *args, **kwargs):
        if self._index >= len(self._values):
            return None
        value = self._values[self._index]
        self._index += 1
        return value

    def scalars(self, *args, **kwargs):
        return self

    def all(self):
        return []


@pytest.fixture(autouse=True)
def patch_provider_map(monkeypatch: pytest.MonkeyPatch):
    monkeypatch.setattr(
        "core.ops.ops_trace_manager.provider_config_map", FakeProviderMap({"dummy": FAKE_PROVIDER_ENTRY})
    )
    OpsTraceManager.ops_trace_instances_cache.clear()
    OpsTraceManager.decrypted_configs_cache.clear()


@pytest.fixture(autouse=True)
def patch_timer_and_current_app(monkeypatch: pytest.MonkeyPatch):
    monkeypatch.setattr("core.ops.ops_trace_manager.threading.Timer", DummyTimer)
    monkeypatch.setattr("core.ops.ops_trace_manager.trace_manager_queue", queue.Queue())
    monkeypatch.setattr("core.ops.ops_trace_manager.trace_manager_timer", None)

    class FakeApp:
        def app_context(self):
            return contextlib.nullcontext()

    fake_current = MagicMock()
    fake_current._get_current_object.return_value = FakeApp()
    monkeypatch.setattr("core.ops.ops_trace_manager.current_app", fake_current)


@pytest.fixture(autouse=True)
def patch_sqlalchemy_session(monkeypatch: pytest.MonkeyPatch):
    monkeypatch.setattr("core.ops.ops_trace_manager.Session", DummySessionContext)


@pytest.fixture
def encryption_mocks(monkeypatch: pytest.MonkeyPatch):
    encrypt_mock = MagicMock(side_effect=lambda tenant, value: f"enc-{value}")
    batch_decrypt_mock = MagicMock(side_effect=lambda tenant, values: [f"dec-{value}" for value in values])
    obfuscate_mock = MagicMock(side_effect=lambda value: f"ob-{value}")
    monkeypatch.setattr("core.ops.ops_trace_manager.encrypt_token", encrypt_mock)
    monkeypatch.setattr("core.ops.ops_trace_manager.batch_decrypt_token", batch_decrypt_mock)
    monkeypatch.setattr("core.ops.ops_trace_manager.obfuscated_token", obfuscate_mock)
    return encrypt_mock, batch_decrypt_mock, obfuscate_mock


@pytest.fixture
def mock_db(monkeypatch: pytest.MonkeyPatch):
    session = MagicMock()
    session.scalars.return_value.all.return_value = ["chat"]
    db_mock = MagicMock()
    db_mock.session = session
    db_mock.engine = MagicMock()
    monkeypatch.setattr("core.ops.ops_trace_manager.db", db_mock)
    return session


@pytest.fixture
def workflow_repo_fixture(monkeypatch: pytest.MonkeyPatch):
    repo = MagicMock()
    repo.get_workflow_run_by_id_without_tenant.return_value = make_workflow_run()
    monkeypatch.setattr(TraceTask, "_get_workflow_run_repo", classmethod(lambda cls: repo))
    return repo


@pytest.fixture
def trace_task_message(monkeypatch: pytest.MonkeyPatch, mock_db):
    message_data = make_message_data()
    monkeypatch.setattr("core.ops.ops_trace_manager.get_message_data", lambda msg_id: message_data)
    configure_db_scalar(mock_db, message_file=FakeMessageFile(), workflow_app_log=SimpleNamespace(id="log-id"))
    return message_data


def test_encrypt_tracing_config_handles_star_and_encrypt(encryption_mocks):
    encrypted = OpsTraceManager.encrypt_tracing_config(
        "tenant",
        "dummy",
        {"secret_value": "value", "other_value": "info"},
        current_trace_config={"secret_value": "keep"},
    )
    assert encrypted["secret_value"] == "enc-value"
    assert encrypted["other_value"] == "info"


def test_encrypt_tracing_config_preserves_star(encryption_mocks):
    encrypted = OpsTraceManager.encrypt_tracing_config(
        "tenant",
        "dummy",
        {"secret_value": "*", "other_value": "info"},
        current_trace_config={"secret_value": "keep"},
    )
    assert encrypted["secret_value"] == "keep"


def test_decrypt_tracing_config_caches(encryption_mocks):
    _, decrypt_mock, _ = encryption_mocks
    payload = {"secret_value": "enc", "other_value": "info"}
    first = OpsTraceManager.decrypt_tracing_config("tenant", "dummy", payload)
    second = OpsTraceManager.decrypt_tracing_config("tenant", "dummy", payload)
    assert first == second
    assert decrypt_mock.call_count == 1


def test_obfuscated_decrypt_token(encryption_mocks):
    _, _, obfuscate_mock = encryption_mocks
    result = OpsTraceManager.obfuscated_decrypt_token("dummy", {"secret_value": "value", "other_value": "info"})
    assert "secret_value" in result
    assert result["secret_value"] == "ob-value"
    obfuscate_mock.assert_called_once()


def test_get_decrypted_tracing_config_returns_config(encryption_mocks, mock_db):
    trace_config_data = SimpleNamespace(tracing_config={"secret_value": "enc", "other_value": "info"})
    app = SimpleNamespace(id="app-id", tenant_id="tenant")
    mock_db.scalar.side_effect = [trace_config_data, app]

    decrypted = OpsTraceManager.get_decrypted_tracing_config("app-id", "dummy")
    assert decrypted["other_value"] == "info"


def test_get_decrypted_tracing_config_missing_trace_config(mock_db):
    mock_db.scalar.return_value = None
    assert OpsTraceManager.get_decrypted_tracing_config("app-id", "dummy") is None


def test_get_decrypted_tracing_config_raises_for_missing_app(mock_db):
    trace_config_data = SimpleNamespace(tracing_config={"secret_value": "enc"})
    mock_db.scalar.side_effect = [trace_config_data, None]
    with pytest.raises(ValueError, match="App not found"):
        OpsTraceManager.get_decrypted_tracing_config("app-id", "dummy")


def test_get_decrypted_tracing_config_raises_for_none_config(mock_db):
    trace_config_data = SimpleNamespace(tracing_config=None)
    mock_db.scalar.side_effect = [trace_config_data, SimpleNamespace(tenant_id="tenant")]
    with pytest.raises(ValueError, match="Tracing config cannot be None"):
        OpsTraceManager.get_decrypted_tracing_config("app-id", "dummy")


def test_get_ops_trace_instance_handles_none_app(mock_db):
    mock_db.get.return_value = None
    assert OpsTraceManager.get_ops_trace_instance("app-id") is None


def test_get_ops_trace_instance_returns_none_when_disabled(mock_db, monkeypatch: pytest.MonkeyPatch):
    app = SimpleNamespace(id="app-id", tracing=json.dumps({"enabled": False}))
    mock_db.get.return_value = app
    assert OpsTraceManager.get_ops_trace_instance("app-id") is None


def test_get_ops_trace_instance_invalid_provider(mock_db, monkeypatch: pytest.MonkeyPatch):
    app = SimpleNamespace(id="app-id", tracing=json.dumps({"enabled": True, "tracing_provider": "missing"}))
    mock_db.get.return_value = app
    monkeypatch.setattr("core.ops.ops_trace_manager.provider_config_map", FakeProviderMap({}))
    assert OpsTraceManager.get_ops_trace_instance("app-id") is None


def test_get_ops_trace_instance_success(monkeypatch: pytest.MonkeyPatch, mock_db):
    app = SimpleNamespace(id="app-id", tracing=json.dumps({"enabled": True, "tracing_provider": "dummy"}))
    mock_db.get.return_value = app
    monkeypatch.setattr(
        "core.ops.ops_trace_manager.OpsTraceManager.get_decrypted_tracing_config",
        classmethod(lambda cls, aid, provider: {"secret_value": "decrypted", "other_value": "info"}),
    )
    instance = OpsTraceManager.get_ops_trace_instance("app-id")
    assert instance is not None
    cached_instance = OpsTraceManager.get_ops_trace_instance("app-id")
    assert instance is cached_instance


def test_get_app_config_through_message_id_returns_none(mock_db):
    mock_db.scalar.return_value = None
    assert OpsTraceManager.get_app_config_through_message_id("m") is None


def test_get_app_config_through_message_id_prefers_override(mock_db):
    message = SimpleNamespace(conversation_id="conv")
    conversation = SimpleNamespace(app_model_config_id=None, override_model_configs={"foo": "bar"})
    app_config = SimpleNamespace(id="config-id")
    mock_db.scalar.side_effect = [message, conversation]
    result = OpsTraceManager.get_app_config_through_message_id("m")
    assert result == {"foo": "bar"}


def test_get_app_config_through_message_id_app_model_config(mock_db):
    message = SimpleNamespace(conversation_id="conv")
    conversation = SimpleNamespace(app_model_config_id="cfg", override_model_configs=None)
    mock_db.scalar.side_effect = [message, conversation, SimpleNamespace(id="cfg")]
    result = OpsTraceManager.get_app_config_through_message_id("m")
    assert result.id == "cfg"


def test_update_app_tracing_config_invalid_provider(mock_db, monkeypatch: pytest.MonkeyPatch):
    mock_db.get.return_value = None
    with pytest.raises(ValueError, match="Invalid tracing provider"):
        OpsTraceManager.update_app_tracing_config("app", True, "bad")
    with pytest.raises(ValueError, match="App not found"):
        OpsTraceManager.update_app_tracing_config("app", True, None)


def test_update_app_tracing_config_success(mock_db):
    app = SimpleNamespace(id="app-id", tracing="{}")
    mock_db.get.return_value = app
    OpsTraceManager.update_app_tracing_config("app-id", True, "dummy")
    assert app.tracing is not None
    mock_db.commit.assert_called_once()


def test_get_app_tracing_config_errors_when_missing(mock_db):
    mock_db.get.return_value = None
    with pytest.raises(ValueError, match="App not found"):
        OpsTraceManager.get_app_tracing_config("app", mock_db)


def test_get_app_tracing_config_returns_defaults(mock_db):
    mock_db.get.return_value = SimpleNamespace(tracing=None)
    assert OpsTraceManager.get_app_tracing_config("app-id", mock_db) == {"enabled": False, "tracing_provider": None}


def test_get_app_tracing_config_returns_payload(mock_db):
    payload = {"enabled": True, "tracing_provider": "dummy"}
    mock_db.get.return_value = SimpleNamespace(tracing=json.dumps(payload))
    assert OpsTraceManager.get_app_tracing_config("app-id", mock_db) == payload


def test_check_and_project_helpers(monkeypatch: pytest.MonkeyPatch):
    monkeypatch.setattr(
        "core.ops.ops_trace_manager.provider_config_map",
        FakeProviderMap(
            {
                "dummy": {
                    "config_class": DummyConfig,
                    "trace_instance": type(
                        "Trace",
                        (),
                        {
                            "__init__": lambda self, cfg: None,
                            "api_check": lambda self: True,
                            "get_project_key": lambda self: "key",
                            "get_project_url": lambda self: "url",
                        },
                    ),
                    "secret_keys": [],
                    "other_keys": [],
                }
            }
        ),
    )
    assert OpsTraceManager.check_trace_config_is_effective({}, "dummy")
    assert OpsTraceManager.get_trace_config_project_key({}, "dummy") == "key"
    assert OpsTraceManager.get_trace_config_project_url({}, "dummy") == "url"


def test_trace_task_conversation_and_extract(monkeypatch: pytest.MonkeyPatch):
    task = TraceTask(trace_type=TraceTaskName.CONVERSATION_TRACE, message_id="msg")
    assert task.conversation_trace(foo="bar") == {"foo": "bar"}
    assert task._extract_streaming_metrics(make_message_data(message_metadata="not json")) == {}


def test_trace_task_message_trace(trace_task_message, mock_db):
    task = TraceTask(trace_type=TraceTaskName.MESSAGE_TRACE, message_id="msg-id")
    result = task.message_trace("msg-id")
    assert result.message_id == "msg-id"


def test_trace_task_workflow_trace(workflow_repo_fixture, mock_db):
    DummySessionContext.scalar_values = ["wf-app-log", "message-ref"]
    execution = SimpleNamespace(id_="run-id", total_tokens=0)
    task = TraceTask(
        trace_type=TraceTaskName.WORKFLOW_TRACE, workflow_execution=execution, conversation_id="conv", user_id="user"
    )
    result = task.workflow_trace(workflow_run_id="run-id", conversation_id="conv", user_id="user")
    assert result.workflow_run_id == "run-id"
    assert result.workflow_id == "wf-1"


def test_trace_task_moderation_trace(trace_task_message):
    task = TraceTask(trace_type=TraceTaskName.MODERATION_TRACE, message_id="msg-id")
    moderation_result = SimpleNamespace(action="block", preset_response="no", query="q", flagged=True)
    timer = {"start": 1, "end": 2}
    result = task.moderation_trace("msg-id", timer, moderation_result=moderation_result, inputs={"src": "payload"})
    assert result.flagged is True
    assert result.message_id == "log-id"


def test_trace_task_suggested_question_trace(trace_task_message):
    task = TraceTask(trace_type=TraceTaskName.SUGGESTED_QUESTION_TRACE, message_id="msg-id")
    timer = {"start": 1, "end": 2}
    result = task.suggested_question_trace("msg-id", timer, suggested_question=["q1"])
    assert result.message_id == "log-id"
    assert "suggested_question" in result.__dict__


def test_trace_task_dataset_retrieval_trace(trace_task_message):
    task = TraceTask(trace_type=TraceTaskName.DATASET_RETRIEVAL_TRACE, message_id="msg-id")
    timer = {"start": 1, "end": 2}
    mock_doc = SimpleNamespace(model_dump=lambda: {"doc": "value"})
    result = task.dataset_retrieval_trace("msg-id", timer, documents=[mock_doc])
    assert result.documents == [{"doc": "value"}]


def test_trace_task_tool_trace(monkeypatch: pytest.MonkeyPatch, mock_db):
    custom_message = make_message_data(agent_thoughts=[make_agent_thought("tool-a", datetime(2025, 2, 20, 12, 1, 0))])
    monkeypatch.setattr("core.ops.ops_trace_manager.get_message_data", lambda _: custom_message)
    configure_db_scalar(mock_db, message_file=FakeMessageFile())
    task = TraceTask(trace_type=TraceTaskName.TOOL_TRACE, message_id="msg-id")
    timer = {"start": 1, "end": 5}
    result = task.tool_trace("msg-id", timer, tool_name="tool-a", tool_inputs={"foo": 1}, tool_outputs="result")
    assert result.tool_name == "tool-a"
    assert result.time_cost == 5


def test_trace_task_generate_name_trace():
    task = TraceTask(trace_type=TraceTaskName.GENERATE_NAME_TRACE, conversation_id="conv-id")
    timer = {"start": 1, "end": 2}
    assert task.generate_name_trace("conv-id", timer, tenant_id=None) == {}
    result = task.generate_name_trace(
        "conv-id", timer, tenant_id="tenant", generate_conversation_name="name", inputs="q"
    )
    assert result.outputs == "name"
    assert result.tenant_id == "tenant"


def test_extract_streaming_metrics_invalid_json():
    task = TraceTask(trace_type=TraceTaskName.MESSAGE_TRACE, message_id="msg-id")
    fake_message = make_message_data(message_metadata="invalid")
    assert task._extract_streaming_metrics(fake_message) == {}


def test_trace_queue_manager_add_and_collect(monkeypatch: pytest.MonkeyPatch):
    monkeypatch.setattr(
        "core.ops.ops_trace_manager.OpsTraceManager.get_ops_trace_instance", classmethod(lambda cls, aid: True)
    )
    manager = TraceQueueManager(app_id="app-id", user_id="user")
    task = TraceTask(trace_type=TraceTaskName.CONVERSATION_TRACE)
    manager.add_trace_task(task)
    tasks = manager.collect_tasks()
    assert tasks == [task]


def test_trace_queue_manager_run_invokes_send(monkeypatch: pytest.MonkeyPatch):
    monkeypatch.setattr(
        "core.ops.ops_trace_manager.OpsTraceManager.get_ops_trace_instance", classmethod(lambda cls, aid: True)
    )
    manager = TraceQueueManager(app_id="app-id", user_id="user")
    task = TraceTask(trace_type=TraceTaskName.CONVERSATION_TRACE)
    called = {}

    def fake_collect():
        return [task]

    def fake_send(tasks):
        called["tasks"] = tasks

    monkeypatch.setattr(TraceQueueManager, "collect_tasks", lambda self: fake_collect())
    monkeypatch.setattr(TraceQueueManager, "send_to_celery", lambda self, t: fake_send(t))
    manager.run()
    assert called["tasks"] == [task]


def test_trace_queue_manager_send_to_celery(monkeypatch: pytest.MonkeyPatch):
    monkeypatch.setattr(
        "core.ops.ops_trace_manager.OpsTraceManager.get_ops_trace_instance", classmethod(lambda cls, aid: True)
    )
    storage_save = MagicMock()
    process_delay = MagicMock()
    monkeypatch.setattr("core.ops.ops_trace_manager.storage.save", storage_save)
    monkeypatch.setattr("core.ops.ops_trace_manager.process_trace_tasks.delay", process_delay)
    monkeypatch.setattr("core.ops.ops_trace_manager.uuid4", MagicMock(return_value=SimpleNamespace(hex="file-123")))

    manager = TraceQueueManager(app_id="app-id", user_id="user")

    class DummyTraceInfo:
        def model_dump(self):
            return {"trace": "info"}

    class DummyTask:
        def __init__(self):
            self.app_id = "app-id"

        def execute(self):
            return DummyTraceInfo()

    task = DummyTask()
    manager.send_to_celery([task])
    storage_save.assert_called_once()
    process_delay.assert_called_once_with({"file_id": "file-123", "app_id": "app-id"})
