"""Tests for tool call telemetry."""

import json
import os
from dataclasses import dataclass
from pathlib import Path
from unittest.mock import patch

import pandas as pd
import pytest

from m4.core.context import M4ExecutionContext
from m4.core.datasets import DatasetDefinition, Modality
from m4.core.exceptions import M4Error
from m4.core.telemetry import (
    TELEMETRY_FILENAME,
    _agent_id_var,
    _get_terminal_session,
    _interface_var,
    _writer,
    invoke_tracked,
    set_agent_id,
    set_interface,
)
from m4.core.tools.base import ToolInput

# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------


@dataclass
class MockInput(ToolInput):
    sql_query: str = "SELECT 1"
    limit: int = 10


class MockTool:
    name = "mock_tool"
    description = "A mock tool for testing"
    input_model = MockInput
    required_modalities = frozenset({Modality.TABULAR})
    supported_datasets = None

    def __init__(self, return_value="ok", side_effect=None):
        self._return_value = return_value
        self._side_effect = side_effect

    def invoke(self, dataset, params, context):
        if self._side_effect:
            raise self._side_effect
        return self._return_value

    def is_compatible(self, dataset):
        return True


@pytest.fixture
def mock_dataset():
    return DatasetDefinition(
        name="test-dataset",
        description="Test dataset",
        modalities=frozenset({Modality.TABULAR}),
    )


@pytest.fixture(autouse=True)
def reset_telemetry_state():
    """Reset telemetry writer and context vars between tests."""
    _writer.reset()
    token_iface = _interface_var.set("unknown")
    token_agent = _agent_id_var.set(None)
    yield
    _writer.reset()
    _interface_var.reset(token_iface)
    _agent_id_var.reset(token_agent)


def _make_context(
    mock_dataset,
    *,
    interface="unknown",
    study_id=None,
    session_id=None,
    actor=None,
):
    return M4ExecutionContext(
        dataset=mock_dataset,
        backend_name="mock",
        backend=object(),
        interface=interface,
        study_id=study_id,
        session_id=session_id,
        actor=actor,
    )


def _capture_record(mock_dataset, tool=None, params=None, context=None):
    """Helper: invoke a tool and return the parsed telemetry record."""
    tool = tool or MockTool(return_value="ok")
    params = params or MockInput()
    context = context or _make_context(mock_dataset)
    with patch("m4.core.telemetry.logger") as mock_logger:
        invoke_tracked(tool, mock_dataset, params, context)
        return json.loads(mock_logger.info.call_args[0][0])


# ---------------------------------------------------------------------------
# Tests
# ---------------------------------------------------------------------------


class TestInvokeTracked:
    def test_returns_tool_result(self, mock_dataset):
        tool = MockTool(return_value={"data": [1, 2, 3]})
        result = invoke_tracked(
            tool, mock_dataset, MockInput(), _make_context(mock_dataset)
        )
        assert result == {"data": [1, 2, 3]}

    def test_success_record(self, mock_dataset):
        record = _capture_record(mock_dataset)

        assert record["success"] is True
        assert record["tool_name"] == "mock_tool"
        assert record["dataset_name"] == "test-dataset"
        assert record["error_type"] is None
        assert record["error_message"] is None
        assert record["duration_ms"] >= 0

    def test_failure_record_reraises(self, mock_dataset):
        tool = MockTool(side_effect=M4Error("something broke"))

        with patch("m4.core.telemetry.logger") as mock_logger:
            with pytest.raises(M4Error, match="something broke"):
                invoke_tracked(
                    tool, mock_dataset, MockInput(), _make_context(mock_dataset)
                )

            record = json.loads(mock_logger.info.call_args[0][0])

        assert record["success"] is False
        assert record["error_type"] == "M4Error"
        assert record["error_message"] == "something broke"

    def test_context_vars_in_record(self, mock_dataset):
        set_agent_id("agent-42")

        record = _capture_record(
            mock_dataset,
            context=_make_context(mock_dataset, interface="mcp"),
        )

        assert record["interface"] == "mcp"
        assert record["agent_id"] == "agent-42"
        assert record["actor"] == "agent-42"

    def test_params_captured(self, mock_dataset):
        params = MockInput(sql_query="SELECT * FROM patients", limit=5)
        record = _capture_record(mock_dataset, params=params)

        assert record["params_summary"]["sql_query"] == "SELECT * FROM patients"
        assert record["params_summary"]["limit"] == 5
        assert record["query_hash"] is not None

    def test_explicit_attribution_fields(self, mock_dataset):
        record = _capture_record(
            mock_dataset,
            context=_make_context(
                mock_dataset,
                study_id="study-1",
                session_id="session-1",
                actor="actor-1",
            ),
        )

        assert record["study_id"] == "study-1"
        assert record["session_id"] == "session-1"
        assert record["actor"] == "actor-1"

    def test_log_sql_hash_mode_redacts_full_sql(self, mock_dataset):
        with patch.dict(os.environ, {"M4_LOG_SQL": "hash"}):
            record = _capture_record(
                mock_dataset,
                params=MockInput(sql_query="SELECT * FROM patients", limit=5),
            )

        assert "sql_query" not in record["params_summary"]
        assert record["params_summary"]["sql_query_hash"] == record["query_hash"]


class TestJSONLFile:
    def test_jsonl_written(self, mock_dataset, tmp_path):
        """JSONL file is written with valid JSON per line."""
        with patch("m4.config.get_telemetry_dir", return_value=tmp_path):
            tool = MockTool(return_value="ok")
            invoke_tracked(tool, mock_dataset, MockInput(), _make_context(mock_dataset))
            invoke_tracked(
                tool,
                mock_dataset,
                MockInput(sql_query="SELECT 2"),
                _make_context(mock_dataset),
            )

        jsonl_path = tmp_path / "tool_calls.jsonl"
        assert jsonl_path.exists()

        lines = jsonl_path.read_text().strip().split("\n")
        assert len(lines) == 2

        for line in lines:
            record = json.loads(line)
            assert record["tool_name"] == "mock_tool"
            assert record["success"] is True

    def test_telemetry_off_suppresses_file(self, mock_dataset, tmp_path):
        """M4_TELEMETRY=off suppresses file output."""
        with patch.dict(os.environ, {"M4_TELEMETRY": "off"}):
            with patch("m4.config.get_telemetry_dir", return_value=tmp_path):
                tool = MockTool(return_value="ok")
                invoke_tracked(
                    tool, mock_dataset, MockInput(), _make_context(mock_dataset)
                )

        jsonl_path = tmp_path / "tool_calls.jsonl"
        assert not jsonl_path.exists()

    def test_event_log_env_overrides_default_path(self, mock_dataset, tmp_path):
        event_log = tmp_path / "events.jsonl"
        with patch.dict(os.environ, {"M4_EVENT_LOG": str(event_log)}):
            with patch("m4.config.get_telemetry_dir", return_value=tmp_path / "unused"):
                invoke_tracked(
                    MockTool(), mock_dataset, MockInput(), _make_context(mock_dataset)
                )

        assert event_log.exists()
        assert not (tmp_path / "unused" / "tool_calls.jsonl").exists()


class TestContextVars:
    def test_default_interface(self, mock_dataset):
        record = _capture_record(mock_dataset)

        assert record["interface"] == "unknown"
        assert record["agent_id"] is None

    def test_set_interface(self):
        set_interface("python_api")
        assert _interface_var.get() == "python_api"

    def test_set_agent_id(self):
        set_agent_id("my-agent")
        assert _agent_id_var.get() == "my-agent"


class TestTerminalSession:
    def test_terminal_session_present(self, mock_dataset):
        record = _capture_record(mock_dataset)
        session = record["terminal_session"]
        assert session is not None
        assert session.isdigit()

    def test_terminal_session_consistent(self, mock_dataset):
        r1 = _capture_record(mock_dataset)
        r2 = _capture_record(mock_dataset)
        assert r1["terminal_session"] == r2["terminal_session"]

    def test_terminal_session_none_without_getsid(self, mock_dataset):
        with patch("m4.core.telemetry.os.getsid", side_effect=AttributeError):
            record = _capture_record(mock_dataset)
        assert record["terminal_session"] is None

    def test_get_terminal_session_returns_string(self):
        result = _get_terminal_session()
        assert result is None or isinstance(result, str)


class TestRowCount:
    def test_row_count_for_dataframe(self, mock_dataset):
        df = pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]})
        tool = MockTool(return_value=df)
        record = _capture_record(mock_dataset, tool=tool)
        assert record["row_count"] == 3

    def test_row_count_none_for_dict(self, mock_dataset):
        tool = MockTool(return_value={"schema": "info"})
        record = _capture_record(mock_dataset, tool=tool)
        assert record["row_count"] is None

    def test_row_count_none_for_string(self, mock_dataset):
        tool = MockTool(return_value="ok")
        record = _capture_record(mock_dataset, tool=tool)
        assert record["row_count"] is None

    def test_row_count_none_on_error(self, mock_dataset):
        tool = MockTool(side_effect=M4Error("fail"))
        with patch("m4.core.telemetry.logger") as mock_logger:
            with pytest.raises(M4Error):
                invoke_tracked(
                    tool, mock_dataset, MockInput(), _make_context(mock_dataset)
                )
            record = json.loads(mock_logger.info.call_args[0][0])
        assert record["row_count"] is None

    def test_row_count_zero_for_empty_dataframe(self, mock_dataset):
        df = pd.DataFrame({"a": []})
        tool = MockTool(return_value=df)
        record = _capture_record(mock_dataset, tool=tool)
        assert record["row_count"] == 0


class TestTelemetryFilename:
    def test_constant_value(self):
        assert TELEMETRY_FILENAME == "tool_calls.jsonl"

    def test_jsonl_uses_constant(self, mock_dataset, tmp_path):
        with patch("m4.config.get_telemetry_dir", return_value=tmp_path):
            invoke_tracked(
                MockTool(), mock_dataset, MockInput(), _make_context(mock_dataset)
            )
        assert (tmp_path / TELEMETRY_FILENAME).exists()


class TestGetTelemetryPath:
    def test_returns_correct_path(self, tmp_path):
        with patch.dict(os.environ, {}, clear=False):
            os.environ.pop("M4_EVENT_LOG", None)
            with patch("m4.config.get_telemetry_dir", return_value=tmp_path):
                from m4.api import get_telemetry_path

                path = get_telemetry_path()
        assert path == tmp_path / "tool_calls.jsonl"

    def test_returns_event_log_env_path(self, tmp_path):
        event_log = tmp_path / "custom-events.jsonl"
        with patch.dict(os.environ, {"M4_EVENT_LOG": str(event_log)}):
            with patch("m4.config.get_telemetry_dir", return_value=tmp_path / "unused"):
                from m4.api import get_telemetry_path

                path = get_telemetry_path()
        assert path == event_log

    def test_expands_event_log_env_path(self):
        with patch.dict(os.environ, {"M4_EVENT_LOG": "~/custom-events.jsonl"}):
            from m4.api import get_telemetry_path

            path = get_telemetry_path()
        assert path == Path("~/custom-events.jsonl").expanduser()
