"""
test_viz_endpoints.py — Tests for all 6 visualization endpoints.

No real model. Uses mocked database.

Run: pytest tests/test_viz_endpoints.py -v
"""
import json
import pytest
from unittest.mock import MagicMock, patch
from fastapi.testclient import TestClient


# ── Test data factory ────────────────────────────────────────────────────────

def _make_run(run_id="suite-abc", model_target="Qwen/Qwen3-0.6B", status="completed"):
    """Create a mock run dict as returned by db.get_run()"""
    return {
        "run_id": run_id,
        "model_target": model_target,
        "status": status,
        "created_at": "2026-02-01T10:00:00Z",
        "summary_json": json.dumps({
            "total_attacks": 10,
            "passed": 4,
            "failed": 5,
            "warnings": 1,
            "critical_count": 2,
            "high_count": 3,
            "policy_gap_score": 55.0,
            "radar_scores": {
                "injection": 0.8,
                "jailbreak": 0.4,
                "bias": 0.2,
            },
            "compliance_coverage": {
                "GDPR_ART22": "covered",
                "HIPAA": "partial",
                "EU_AI_ACT_ART10": "not_tested",
            },
        }),
    }


def _make_results():
    """Create mock results as returned by db.get_results_for_run()"""
    return [
        {
            "attack_name": "prompt_injection_direct",
            "attack_type": "llm",
            "attack_category": "injection",
            "status": "fail",
            "severity": "high",
            "success_rate": 1.0,
            "turn_number": None,
            "safety_score": 0.1,
            "compliance_tags": ["GDPR_ART22"],
        },
        {
            "attack_name": "crescendo",
            "attack_type": "alignment",
            "attack_category": "alignment",
            "status": "fail",
            "severity": "critical",
            "success_rate": 1.0,
            "turn_number": 3,
            "safety_score": 0.2,
            "compliance_tags": ["EU_AI_ACT_ART13"],
        },
        {
            "attack_name": "crescendo",
            "attack_type": "alignment",
            "attack_category": "alignment",
            "status": "pass",
            "severity": "low",
            "success_rate": 0.0,
            "turn_number": 1,
            "safety_score": 0.9,
            "compliance_tags": [],
        },
        {
            "attack_name": "bias_fairness",
            "attack_type": "bias",
            "attack_category": "bias",
            "status": "pass",
            "severity": "low",
            "success_rate": 0.0,
            "turn_number": None,
            "safety_score": 1.0,
            "compliance_tags": ["HIPAA"],
        },
    ]


@pytest.fixture
def client():
    """Create test client with mocked database and bypassed auth."""
    with patch("sichgate.server.Database") as MockDB, \
         patch("sichgate.server._SESSION_TOKEN", "test-token"):
        mock_db = MagicMock()
        MockDB.return_value = mock_db
        
        # Store mock_db on fixture for per-test configuration
        from sichgate.server import create_app
        app = create_app()
        
        test_client = TestClient(app)
        test_client._mock_db = mock_db
        yield test_client


# ── Radar ────────────────────────────────────────────────────────────────────

class TestVizRadar:
    def test_returns_axes(self, client):
        client._mock_db.get_run.return_value = _make_run()
        r = client.get(
            "/api/v1/viz/radar/suite-abc",
            headers={"X-Session-Token": "test-token"}
        )
        assert r.status_code == 200
        data = r.json()
        assert "axes" in data
        assert len(data["axes"]) == 3
        assert all("category" in a and "score" in a and "label" in a for a in data["axes"])

    def test_scores_are_floats_0_to_1(self, client):
        client._mock_db.get_run.return_value = _make_run()
        r = client.get(
            "/api/v1/viz/radar/suite-abc",
            headers={"X-Session-Token": "test-token"}
        )
        for axis in r.json()["axes"]:
            assert 0.0 <= axis["score"] <= 1.0

    def test_404_on_missing_suite(self, client):
        client._mock_db.get_run.return_value = None
        r = client.get(
            "/api/v1/viz/radar/nonexistent",
            headers={"X-Session-Token": "test-token"}
        )
        assert r.status_code == 404


# ── Heatmap ──────────────────────────────────────────────────────────────────

class TestVizHeatmap:
    def test_returns_matrix_shape(self, client):
        run1 = _make_run("s1", "Qwen/Qwen3-0.6B")
        run2 = _make_run("s2", "Phi-2")
        client._mock_db.get_run.side_effect = lambda sid: {"s1": run1, "s2": run2}.get(sid)
        client._mock_db.get_results_for_run.return_value = _make_results()
        
        r = client.get(
            "/api/v1/viz/heatmap?suite_ids=s1&suite_ids=s2",
            headers={"X-Session-Token": "test-token"}
        )
        assert r.status_code == 200
        data = r.json()
        assert "rows" in data
        assert "columns" in data
        assert "matrix" in data
        assert len(data["columns"]) == 2

    def test_matrix_values_are_floats(self, client):
        client._mock_db.get_run.return_value = _make_run("s1")
        client._mock_db.get_results_for_run.return_value = _make_results()
        
        r = client.get(
            "/api/v1/viz/heatmap?suite_ids=s1",
            headers={"X-Session-Token": "test-token"}
        )
        for row_vals in r.json()["matrix"].values():
            for val in row_vals.values():
                assert isinstance(val, float)
                assert 0.0 <= val <= 1.0


# ── Gap Gauge ────────────────────────────────────────────────────────────────

class TestVizGap:
    def test_returns_score_and_grade(self, client):
        client._mock_db.get_run.return_value = _make_run()
        r = client.get(
            "/api/v1/viz/gap/suite-abc",
            headers={"X-Session-Token": "test-token"}
        )
        assert r.status_code == 200
        data = r.json()
        assert "policy_gap_score" in data
        assert "grade" in data
        assert data["grade"] in ("A", "B", "C", "D", "F")

    def test_grade_f_for_high_score(self, client):
        run = _make_run()
        summary = json.loads(run["summary_json"])
        summary["policy_gap_score"] = 95.0
        run["summary_json"] = json.dumps(summary)
        client._mock_db.get_run.return_value = run
        
        r = client.get(
            "/api/v1/viz/gap/suite-abc",
            headers={"X-Session-Token": "test-token"}
        )
        assert r.json()["grade"] == "F"

    def test_grade_a_for_low_score(self, client):
        run = _make_run()
        summary = json.loads(run["summary_json"])
        summary["policy_gap_score"] = 10.0
        run["summary_json"] = json.dumps(summary)
        client._mock_db.get_run.return_value = run
        
        r = client.get(
            "/api/v1/viz/gap/suite-abc",
            headers={"X-Session-Token": "test-token"}
        )
        assert r.json()["grade"] == "A"


# ── Trajectory ───────────────────────────────────────────────────────────────

class TestVizTrajectory:
    def test_returns_crescendo_turns(self, client):
        client._mock_db.get_run.return_value = _make_run()
        client._mock_db.get_results_for_run.return_value = _make_results()
        
        r = client.get(
            "/api/v1/viz/trajectory/suite-abc",
            headers={"X-Session-Token": "test-token"}
        )
        assert r.status_code == 200
        data = r.json()
        assert "attacks" in data
        # crescendo has turn_number set in our test data
        crescendo = next((a for a in data["attacks"] if a["attack_name"] == "crescendo"), None)
        assert crescendo is not None
        assert len(crescendo["turns"]) == 2  # turns 1 and 3

    def test_turns_sorted_ascending(self, client):
        client._mock_db.get_run.return_value = _make_run()
        client._mock_db.get_results_for_run.return_value = _make_results()
        
        r = client.get(
            "/api/v1/viz/trajectory/suite-abc",
            headers={"X-Session-Token": "test-token"}
        )
        for attack in r.json()["attacks"]:
            turn_nums = [t["turn"] for t in attack["turns"]]
            assert turn_nums == sorted(turn_nums)

    def test_single_turn_attacks_excluded(self, client):
        client._mock_db.get_run.return_value = _make_run()
        client._mock_db.get_results_for_run.return_value = _make_results()
        
        r = client.get(
            "/api/v1/viz/trajectory/suite-abc",
            headers={"X-Session-Token": "test-token"}
        )
        names = [a["attack_name"] for a in r.json()["attacks"]]
        assert "prompt_injection_direct" not in names  # turn_number=None


# ── Compliance ───────────────────────────────────────────────────────────────

class TestVizCompliance:
    def test_returns_all_tags(self, client):
        client._mock_db.get_run.return_value = _make_run()
        r = client.get(
            "/api/v1/viz/compliance/suite-abc",
            headers={"X-Session-Token": "test-token"}
        )
        assert r.status_code == 200
        data = r.json()
        assert "coverage" in data
        tags = [c["tag"] for c in data["coverage"]]
        assert "GDPR_ART22" in tags
        assert "HIPAA" in tags
        assert "EU_AI_ACT_ART10" in tags

    def test_status_values_are_valid(self, client):
        client._mock_db.get_run.return_value = _make_run()
        r = client.get(
            "/api/v1/viz/compliance/suite-abc",
            headers={"X-Session-Token": "test-token"}
        )
        valid = {"covered", "partial", "not_tested"}
        for item in r.json()["coverage"]:
            assert item["status"] in valid

    def test_each_item_has_framework(self, client):
        client._mock_db.get_run.return_value = _make_run()
        r = client.get(
            "/api/v1/viz/compliance/suite-abc",
            headers={"X-Session-Token": "test-token"}
        )
        for item in r.json()["coverage"]:
            assert "framework" in item
            assert item["framework"] in ("GDPR", "HIPAA", "EU AI Act", "NIST")


# ── Timeline ─────────────────────────────────────────────────────────────────

class TestVizTimeline:
    def test_returns_points_for_model(self, client):
        run1 = _make_run("s1", "Qwen/Qwen3-0.6B")
        run2 = _make_run("s2", "Qwen/Qwen3-0.6B")
        run3 = _make_run("s3", "Phi-2")  # different model — should be excluded
        client._mock_db.get_runs.return_value = [run1, run2, run3]
        
        r = client.get(
            "/api/v1/viz/timeline?model_target=Qwen/Qwen3-0.6B",
            headers={"X-Session-Token": "test-token"}
        )
        assert r.status_code == 200
        data = r.json()
        assert len(data["points"]) == 2
        assert all(p.get("policy_gap_score") is not None for p in data["points"])

    def test_excludes_incomplete_suites(self, client):
        run1 = _make_run("s1", "Qwen/Qwen3-0.6B", status="running")
        run2 = _make_run("s2", "Qwen/Qwen3-0.6B", status="completed")
        client._mock_db.get_runs.return_value = [run1, run2]
        
        r = client.get(
            "/api/v1/viz/timeline?model_target=Qwen/Qwen3-0.6B",
            headers={"X-Session-Token": "test-token"}
        )
        assert len(r.json()["points"]) == 1

    def test_points_sorted_by_timestamp(self, client):
        run1 = _make_run("s1", "Qwen/Qwen3-0.6B")
        run1["created_at"] = "2026-02-08T00:00:00Z"
        run2 = _make_run("s2", "Qwen/Qwen3-0.6B")
        run2["created_at"] = "2026-02-01T00:00:00Z"
        client._mock_db.get_runs.return_value = [run1, run2]
        
        r = client.get(
            "/api/v1/viz/timeline?model_target=Qwen/Qwen3-0.6B",
            headers={"X-Session-Token": "test-token"}
        )
        timestamps = [p["timestamp"] for p in r.json()["points"]]
        assert timestamps == sorted(timestamps)
