from __future__ import annotations

import io
import json
import sys
import unittest
from pathlib import Path
from urllib.error import HTTPError
from unittest.mock import patch

sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))

from refario_sdk import create_mcp_telemetry_emitter, parse_usage, refario, traced_llm  # noqa: E402
from refario_sdk.client import IngestionError, RefarioClient  # noqa: E402


class _FakeResponse:
    def __init__(self, status: int = 200, body: bytes = b'{"received": true}') -> None:
        self.status = status
        self._body = body

    def read(self) -> bytes:
        return self._body

    def __enter__(self) -> "_FakeResponse":
        return self

    def __exit__(self, exc_type, exc, tb) -> bool:
        return False


class _CapturingRun:
    def __init__(self) -> None:
        self.calls: list[dict[str, object]] = []

    def llm_call(self, **kwargs: object) -> None:
        self.calls.append(kwargs)


class RefarioSdkTests(unittest.TestCase):
    def test_init_supports_ts_style_config(self) -> None:
        client = refario.init({"endpoint": "http://localhost:3000", "apiKey": "test-key"})
        self.assertIsInstance(client, RefarioClient)
        self.assertEqual(client.endpoint, "http://localhost:3000")
        self.assertEqual(client.api_key, "test-key")

    def test_start_run_supports_string_and_mapping(self) -> None:
        client = refario.init(endpoint="http://localhost:3000", api_key="test-key")

        run_a = client.start_run("workflow-a")
        run_b = client.start_run({"workflow": "workflow-b", "environment": "prod"})

        self.assertEqual(run_a.workflow, "workflow-a")
        self.assertIsNone(run_a.environment)
        self.assertEqual(run_b.workflow, "workflow-b")
        self.assertEqual(run_b.environment, "prod")

    @patch("refario_sdk.client.urlopen")
    def test_end_posts_event(self, mock_urlopen) -> None:
        mock_urlopen.return_value = _FakeResponse()

        client = refario.init(endpoint="http://localhost:3000", api_key="test-key")
        run = client.start_run("example-workflow")
        run.llm_call(
            name="openai-call",
            provider="openai",
            model="gpt-4o",
            prompt_tokens=10,
            completion_tokens=5,
            status="success",
        )

        response = run.end(success=True)

        self.assertEqual(response["received"], True)
        request = mock_urlopen.call_args.args[0]
        timeout = mock_urlopen.call_args.kwargs["timeout"]
        self.assertEqual(timeout, 10.0)
        self.assertTrue(request.full_url.endswith("/api/v1/events"))

        headers = {key.lower(): value for key, value in request.header_items()}
        self.assertEqual(headers["x-refario-key"], "test-key")
        self.assertEqual(headers["content-type"], "application/json")

        payload = json.loads(request.data.decode("utf-8"))
        self.assertEqual(payload["workflow"], "example-workflow")
        self.assertEqual(payload["status"], "success")
        self.assertEqual(payload["idempotency_key"], payload["run_id"])
        self.assertEqual(payload["spans"][0]["total_tokens"], 15)

    @patch("refario_sdk.client.urlopen")
    def test_end_raises_ingestion_error_for_http_error(self, mock_urlopen) -> None:
        mock_urlopen.side_effect = HTTPError(
            url="http://localhost:3000/api/v1/events",
            code=401,
            msg="Unauthorized",
            hdrs=None,
            fp=io.BytesIO(b'{"error":"invalid api key"}'),
        )

        client = refario.init(endpoint="http://localhost:3000", api_key="bad-key")
        run = client.start_run("example-workflow")
        run.llm_call(name="openai-call")

        with self.assertRaises(IngestionError):
            run.end(success=True)

    def test_parse_usage(self) -> None:
        openai = parse_usage(
            "openai",
            {"prompt_tokens": 11, "completion_tokens": 7, "total_tokens": 18},
        )
        anthropic = parse_usage("anthropic", {"input_tokens": 9, "output_tokens": 4})

        self.assertEqual(openai["prompt_tokens"], 11)
        self.assertEqual(openai["completion_tokens"], 7)
        self.assertEqual(openai["total_tokens"], 18)
        self.assertEqual(anthropic["prompt_tokens"], 9)
        self.assertEqual(anthropic["completion_tokens"], 4)
        self.assertEqual(anthropic["total_tokens"], 13)

    def test_traced_llm_records_success_and_error(self) -> None:
        run = _CapturingRun()

        result = traced_llm(
            run,
            provider="openai",
            model="gpt-4o",
            name="call-1",
            fn=lambda: {"usage": {"prompt_tokens": 5, "completion_tokens": 2}},
        )
        self.assertEqual(result["usage"]["prompt_tokens"], 5)
        self.assertEqual(run.calls[0]["status"], "success")
        self.assertEqual(run.calls[0]["total_tokens"], 7)

        def _boom() -> None:
            raise RuntimeError("failed")

        with self.assertRaises(RuntimeError):
            traced_llm(
                run,
                provider="openai",
                model="gpt-4o",
                name="call-2",
                fn=_boom,
            )

        self.assertEqual(run.calls[1]["status"], "error")
        self.assertEqual(run.calls[1]["error"], "failed")

    @patch("refario_sdk.client.urlopen")
    def test_mcp_session_flush_posts_event(self, mock_urlopen) -> None:
        mock_urlopen.return_value = _FakeResponse()

        emitter = create_mcp_telemetry_emitter(
            endpoint="http://localhost:3000",
            api_key="test-key",
            default_workflow="mcp-session",
            default_environment="production",
            default_transport="http",
        )
        session = emitter.start_session(session_id="session_demo_1", request_id="req_demo_1")
        session.tool_call(
            tool_name="listDocs",
            correlation_id="req_demo_1",
            input_sample='{"folder":"policies"}',
        )
        session.tool_result(
            tool_name="listDocs",
            correlation_id="req_demo_1",
            latency_ms=120,
            output_sample='{"count":4}',
        )

        response = session.flush(success=True)

        self.assertEqual(response["received"], True)
        request = mock_urlopen.call_args.args[0]
        self.assertTrue(request.full_url.endswith("/api/v1/events/mcp"))

        payload = json.loads(request.data.decode("utf-8"))
        self.assertEqual(payload["mcp_session_id"], "session_demo_1")
        self.assertEqual(payload["request_id"], "req_demo_1")
        self.assertEqual(payload["workflow"], "mcp-session")
        self.assertEqual(payload["environment"], "production")
        self.assertEqual(payload["transport"], "http")
        self.assertEqual(payload["status"], "success")
        self.assertEqual(payload["spans"][0]["kind"], "tool_call")
        self.assertEqual(payload["spans"][0]["tool_name"], "listDocs")
        self.assertEqual(payload["spans"][1]["kind"], "tool_result")
        self.assertEqual(payload["spans"][1]["latency_ms"], 120)

    @patch("refario_sdk.client.urlopen")
    def test_mcp_helper_supports_ts_style_aliases(self, mock_urlopen) -> None:
        mock_urlopen.return_value = _FakeResponse()

        emitter = create_mcp_telemetry_emitter(
            {
                "endpoint": "http://localhost:3000",
                "apiKey": "test-key",
                "defaultWorkflow": "mcp-session",
                "defaultEnvironment": "production",
                "defaultTransport": "http",
            }
        )
        session = emitter.startSession({"sessionId": "session_alias_1", "requestId": "req_alias_1"})
        session.toolCall({"toolName": "listDocs", "correlationId": "req_alias_1"})
        session.toolResult(
            {
                "toolName": "listDocs",
                "correlationId": "req_alias_1",
                "latencyMs": 90,
            }
        )
        session.flush(success=True)

        payload = json.loads(mock_urlopen.call_args.args[0].data.decode("utf-8"))
        self.assertEqual(payload["mcp_session_id"], "session_alias_1")
        self.assertEqual(payload["request_id"], "req_alias_1")
        self.assertEqual(payload["spans"][0]["tool_name"], "listDocs")


if __name__ == "__main__":
    unittest.main()
