from __future__ import annotations

import os
import unittest
from types import SimpleNamespace
from unittest.mock import patch

from src.langfuse_control import LangfuseControlError, LangfuseControlService, parse_control_message


class _FakeTraceApi:
    def list(self, **_kwargs):
        return SimpleNamespace(data=[{"id": "tr_1"}])

    def get(self, trace_id: str):
        return {"id": trace_id, "name": "demo"}


class _FakeObservationApi:
    def get_many(self, **_kwargs):
        return SimpleNamespace(data=[{"id": "obs_1"}])


class _FakeApi:
    def __init__(self):
        self.trace = _FakeTraceApi()
        self.observations = _FakeObservationApi()


class _FakeClient:
    def __init__(self):
        self.api = _FakeApi()

    def get_prompt(self, name: str, **_kwargs):
        return {"name": name, "prompt": "hi"}

    def create_prompt(self, **kwargs):
        return {"name": kwargs["name"], "version": 2}


class _FakeTransport:
    def __init__(self):
        self.calls: list[dict] = []

    def request_json(self, **kwargs):
        self.calls.append(kwargs)
        return {"ok": True, "path": kwargs.get("path"), "method": kwargs.get("method")}


class _FakeHttpClient:
    def __init__(self):
        self._transport = _FakeTransport()


class LangfuseControlTests(unittest.TestCase):
    def test_parse_control_message_from_key_value(self) -> None:
        payload = parse_control_message("op=prompt.get name=my.prompt label=latest")
        self.assertEqual(payload["operation"], "prompt.get")
        self.assertEqual(payload["params"]["name"], "my.prompt")

    def test_execute_traces_list(self) -> None:
        settings = {
            "host": "https://example.langfuse.com",
            "public_key": "pk-test",
            "secret_key": "sk-test",
            "environment": "prod",
            "timeout_sec": 10,
        }
        with patch("src.langfuse_control.build_langfuse_client", return_value=(_FakeClient(), "sdk", None)):
            service = LangfuseControlService(settings=settings)
            result = service.execute({"operation": "traces.list", "params": {"limit": 1}})

        self.assertEqual(result["operation"], "traces.list")
        self.assertEqual(result["result"]["backend"], "sdk")
        self.assertIn("data", result["result"])

    def test_raw_http_mutation_is_blocked_by_default(self) -> None:
        settings = {
            "host": "https://example.langfuse.com",
            "public_key": "pk-test",
            "secret_key": "sk-test",
            "environment": "prod",
            "timeout_sec": 10,
        }
        service = LangfuseControlService(settings=settings)
        with patch.dict(os.environ, {"AGENT_LANGFUSE_CONTROL_ALLOW_MUTATION": ""}, clear=False):
            with self.assertRaises(LangfuseControlError):
                service.raw_http(method="POST", path="api/public/v2/prompts", body={"name": "x"})

    def test_execute_prompt_list_uses_http_fallback(self) -> None:
        settings = {
            "host": "https://example.langfuse.com",
            "public_key": "pk-test",
            "secret_key": "sk-test",
            "environment": "prod",
            "timeout_sec": 10,
        }
        fake_http = _FakeHttpClient()
        service = LangfuseControlService(settings=settings)
        with patch.object(service, "_build_http_client", return_value=fake_http):
            result = service.execute({"operation": "prompt.list", "params": {"limit": 3}})
        self.assertEqual(result["operation"], "prompt.list")
        self.assertEqual(result["result"]["backend"], "http-fallback")
        self.assertEqual(fake_http._transport.calls[0]["path"], "api/public/v2/prompts")

    def test_execute_dataset_import_from_traces(self) -> None:
        settings = {
            "host": "https://example.langfuse.com",
            "public_key": "pk-test",
            "secret_key": "sk-test",
            "environment": "prod",
            "timeout_sec": 10,
        }
        fake_http = _FakeHttpClient()
        service = LangfuseControlService(settings=settings)
        with patch.object(service, "_build_http_client", return_value=fake_http):
            with patch.dict(os.environ, {"AGENT_LANGFUSE_CONTROL_ALLOW_MUTATION": "true"}, clear=False):
                result = service.execute(
                    {
                        "operation": "datasets.import_from_traces",
                        "params": {"dataset_id": "ds_1", "trace_ids": ["tr_1", "tr_2"]},
                    }
                )
        self.assertEqual(result["operation"], "datasets.import_from_traces")
        self.assertEqual(fake_http._transport.calls[0]["path"], "api/public/v2/datasets/ds_1/items")

    def test_score_create_requires_mutation_flag(self) -> None:
        settings = {
            "host": "https://example.langfuse.com",
            "public_key": "pk-test",
            "secret_key": "sk-test",
            "environment": "prod",
            "timeout_sec": 10,
        }
        service = LangfuseControlService(settings=settings)
        with patch.dict(os.environ, {"AGENT_LANGFUSE_CONTROL_ALLOW_MUTATION": ""}, clear=False):
            with self.assertRaises(LangfuseControlError):
                service.execute({"operation": "score.create", "params": {"name": "quality", "value": 1}})

    def test_openapi_list_operations(self) -> None:
        settings = {
            "host": "https://example.langfuse.com",
            "public_key": "pk-test",
            "secret_key": "sk-test",
            "environment": "prod",
            "timeout_sec": 10,
        }
        service = LangfuseControlService(settings=settings)
        fake_spec = {
            "openapi": "3.0.0",
            "paths": {
                "/api/public/projects": {
                    "get": {"operationId": "projects_get"},
                    "post": {"operationId": "projects_create"},
                }
            },
        }
        with patch.object(service, "_load_openapi_spec", return_value=fake_spec):
            result = service.execute({"operation": "openapi.list_operations", "params": {}})
        self.assertEqual(result["operation"], "openapi.list_operations")
        self.assertEqual(result["result"]["count"], 2)

    def test_openapi_execute_by_operation_id(self) -> None:
        settings = {
            "host": "https://example.langfuse.com",
            "public_key": "pk-test",
            "secret_key": "sk-test",
            "environment": "prod",
            "timeout_sec": 10,
        }
        service = LangfuseControlService(settings=settings)
        fake_spec = {
            "openapi": "3.0.0",
            "paths": {
                "/api/public/projects/{projectId}": {
                    "get": {
                        "operationId": "projects_getById",
                        "parameters": [{"name": "projectId", "in": "path"}],
                    }
                }
            },
        }
        fake_http = _FakeHttpClient()
        with patch.object(service, "_load_openapi_spec", return_value=fake_spec):
            with patch.object(service, "_build_http_client", return_value=fake_http):
                result = service.execute(
                    {
                        "operation": "openapi.execute",
                        "params": {"operation_id": "projects_getById", "projectId": "proj_1"},
                    }
                )
        self.assertEqual(result["operation"], "openapi.execute")
        self.assertEqual(result["result"]["meta"]["path"], "/api/public/projects/proj_1")

    def test_openapi_list_operations_applies_policy_filters(self) -> None:
        settings = {
            "host": "https://example.langfuse.com",
            "public_key": "pk-test",
            "secret_key": "sk-test",
            "environment": "prod",
            "timeout_sec": 10,
        }
        service = LangfuseControlService(settings=settings)
        fake_spec = {
            "openapi": "3.0.0",
            "paths": {
                "/api/public/projects": {"get": {"operationId": "projects_get"}},
                "/api/public/v2/prompts": {"get": {"operationId": "prompts_list"}},
            },
        }
        with patch.object(service, "_load_openapi_spec", return_value=fake_spec):
            with patch.dict(
                os.environ,
                {
                    "AGENT_LANGFUSE_OPENAPI_ALLOWLIST": "projects_*",
                    "AGENT_LANGFUSE_OPENAPI_DENYLIST": "",
                },
                clear=False,
            ):
                result = service.execute({"operation": "openapi.list_operations", "params": {}})
        operations = result["result"]["operations"]
        self.assertEqual(len(operations), 1)
        self.assertEqual(operations[0]["operation_id"], "projects_get")

    def test_openapi_execute_blocked_by_denylist_policy(self) -> None:
        settings = {
            "host": "https://example.langfuse.com",
            "public_key": "pk-test",
            "secret_key": "sk-test",
            "environment": "prod",
            "timeout_sec": 10,
        }
        service = LangfuseControlService(settings=settings)
        fake_spec = {
            "openapi": "3.0.0",
            "paths": {
                "/api/public/projects/{projectId}": {
                    "get": {
                        "operationId": "projects_getById",
                        "parameters": [{"name": "projectId", "in": "path"}],
                    }
                }
            },
        }
        fake_http = _FakeHttpClient()
        with patch.object(service, "_load_openapi_spec", return_value=fake_spec):
            with patch.object(service, "_build_http_client", return_value=fake_http):
                with patch.dict(
                    os.environ,
                    {
                        "AGENT_LANGFUSE_OPENAPI_ALLOWLIST": "",
                        "AGENT_LANGFUSE_OPENAPI_DENYLIST": "projects_*",
                    },
                    clear=False,
                ):
                    with self.assertRaises(LangfuseControlError):
                        service.execute(
                            {
                                "operation": "openapi.execute",
                                "params": {"operation_id": "projects_getById", "projectId": "proj_1"},
                            }
                        )


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