import json
import os

from fastapi.testclient import TestClient

# storage_dual requires DATABASE_URL at import time.
os.environ.setdefault("DATABASE_URL", "postgresql://cascade:cascade@localhost:5432/cascade")

from backend import api as api_module
from backend.main import app


def _sensitive_attrs():
    return {
        "cascade.span_type": "llm",
        "llm.prompt": "visible",
        "llm.messages": "hidden",
        "llm.prompt_full": "hidden",
        "llm.prompt_last_user": "hidden",
        "ai.prompt.messages": "hidden",
    }


def test_spans_since_strips_sensitive_attrs(monkeypatch):
    monkeypatch.setattr(api_module, "get_project_from_request", lambda request: "org")
    monkeypatch.setattr(
        api_module.storage,
        "get_trace",
        lambda trace_id: [{"project": "org/test"}],
    )
    monkeypatch.setattr(
        api_module.storage,
        "get_spans_since",
        lambda trace_id, since_time, limit=100, include_content=False: {
            "spans": [{"span_id": "s1", "attributes": _sensitive_attrs()}],
            "server_time": 123,
        },
    )
    monkeypatch.setattr(api_module.storage, "is_trace_complete", lambda trace_id: True)

    client = TestClient(app)
    resp = client.get(
        "/api/traces/t1/spans/since?since_time=0&include_content=true",
        headers={"X-API-Key": "fake"},
    )
    assert resp.status_code == 200
    attrs = resp.json()["spans"][0]["attributes"]
    assert attrs["llm.prompt"] == "visible"
    assert "llm.messages" not in attrs
    assert "llm.prompt_full" not in attrs
    assert "llm.prompt_last_user" not in attrs
    assert "ai.prompt.messages" not in attrs


def test_get_span_strips_sensitive_attrs(monkeypatch):
    monkeypatch.setattr(api_module, "get_project_from_request", lambda request: "org")
    monkeypatch.setattr(
        api_module.storage,
        "get_span",
        lambda span_id: {"span_id": span_id, "project": "org/test", "attributes": _sensitive_attrs()},
    )

    client = TestClient(app)
    resp = client.get("/api/spans/s1", headers={"X-API-Key": "fake"})
    assert resp.status_code == 200
    attrs = resp.json()["attributes"]
    assert attrs["llm.prompt"] == "visible"
    assert "llm.messages" not in attrs
    assert "llm.prompt_full" not in attrs
    assert "llm.prompt_last_user" not in attrs
    assert "ai.prompt.messages" not in attrs


def test_trace_download_strips_sensitive_attrs(monkeypatch):
    monkeypatch.setattr(api_module, "get_project_from_request", lambda request: "org")

    # First call for ownership check, then full stream in generator.
    def _stream_trace_spans(trace_id):
        return iter(
            [
                {"span_id": "s1", "project": "org/test", "attributes": _sensitive_attrs()},
                {"span_id": "s2", "project": "org/test", "attributes": {"cascade.span_type": "tool"}},
            ]
        )

    monkeypatch.setattr(api_module.storage, "stream_trace_spans", _stream_trace_spans)

    client = TestClient(app)
    resp = client.get("/api/traces/t1/download", headers={"X-API-Key": "fake"})
    assert resp.status_code == 200
    lines = [ln for ln in resp.text.splitlines() if ln.strip()]
    first = json.loads(lines[0])
    attrs = first["attributes"]
    assert attrs["llm.prompt"] == "visible"
    assert "llm.messages" not in attrs
    assert "llm.prompt_full" not in attrs
    assert "llm.prompt_last_user" not in attrs
    assert "ai.prompt.messages" not in attrs
