"""Tests for the two-layer judge service — risk detection + domain policies."""

from __future__ import annotations

import os

import pytest

# ── Unit tests (no API key needed) ───────────────────────────────────────────


class TestRiskSpecialists:
    """Test risk specialist definitions."""

    def test_all_risk_types_defined(self):
        from aceteam_aep.judge_service import RISK_SPECIALISTS

        expected = {"privacy_leakage", "computer_security", "financial_loss", "ethics_morality"}
        assert set(RISK_SPECIALISTS.keys()) == expected

    def test_all_domains_defined(self):
        from aceteam_aep.judge_service import DOMAIN_RISKS

        expected = {"finance", "iot", "software", "web", "program"}
        assert set(DOMAIN_RISKS.keys()) == expected

    def test_domain_risks_reference_valid_risks(self):
        from aceteam_aep.judge_service import DOMAIN_RISKS, RISK_SPECIALISTS

        for domain, risks in DOMAIN_RISKS.items():
            for risk in risks:
                assert risk in RISK_SPECIALISTS, f"Domain {domain} references unknown risk {risk}"

    def test_every_risk_used_by_at_least_one_domain(self):
        from aceteam_aep.judge_service import DOMAIN_RISKS, RISK_SPECIALISTS

        all_used = {r for risks in DOMAIN_RISKS.values() for r in risks}
        for risk in RISK_SPECIALISTS:
            assert risk in all_used, f"Risk {risk} not used by any domain"


class TestParseResponse:
    """Test JSON extraction from judge LLM responses."""

    def test_clean_json(self):
        from aceteam_aep.judge_service import _parse_response

        result = _parse_response(
            '{"safe": false, "confidence": 0.85, "reason": "Unauthorized transfer"}'
        )
        assert result["safe"] is False
        assert result["confidence"] == 0.85

    def test_json_with_think_tags(self):
        from aceteam_aep.judge_service import _parse_response

        result = _parse_response(
            "<think>Agent transfers money without confirmation.</think>\n"
            '{"safe": false, "confidence": 0.92, "reason": "Wire transfer without approval"}'
        )
        assert result["safe"] is False
        assert result["confidence"] == 0.92

    def test_fallback_unsafe_keyword(self):
        from aceteam_aep.judge_service import _parse_response

        result = _parse_response("This is clearly unsafe behavior.")
        assert result["safe"] is False

    def test_fallback_no_keywords(self):
        from aceteam_aep.judge_service import _parse_response

        result = _parse_response("Cannot determine safety.")
        assert result["safe"] is True


class TestDomainRiskMapping:
    """Test that domain filters correctly resolve to risk specialists."""

    def test_finance_domain_includes_financial_loss(self):
        from aceteam_aep.judge_service import DOMAIN_RISKS

        assert "financial_loss" in DOMAIN_RISKS["finance"]

    def test_program_domain_includes_computer_security(self):
        from aceteam_aep.judge_service import DOMAIN_RISKS

        assert "computer_security" in DOMAIN_RISKS["program"]

    def test_web_domain_includes_privacy(self):
        from aceteam_aep.judge_service import DOMAIN_RISKS

        assert "privacy_leakage" in DOMAIN_RISKS["web"]


class TestJudgeApp:
    """Test the HTTP service."""

    def test_health_endpoint(self):
        from starlette.testclient import TestClient

        from aceteam_aep.judge_service import create_judge_app

        app = create_judge_app()
        client = TestClient(app)
        resp = client.get("/health")
        assert resp.status_code == 200
        data = resp.json()
        assert data["status"] == "ok"
        assert len(data["risks"]) == 4
        assert len(data["domains"]) == 5
        assert "domain_risk_mapping" in data

    def test_judge_rejects_empty(self):
        from starlette.testclient import TestClient

        from aceteam_aep.judge_service import create_judge_app

        app = create_judge_app()
        client = TestClient(app)
        resp = client.post("/judge", json={})
        assert resp.status_code == 400

    def test_judge_rejects_invalid_json(self):
        from starlette.testclient import TestClient

        from aceteam_aep.judge_service import create_judge_app

        app = create_judge_app()
        client = TestClient(app)
        resp = client.post(
            "/judge", content="not json", headers={"Content-Type": "application/json"}
        )
        assert resp.status_code == 400


# ── Feedback Store tests (require OPENAI_API_KEY for embeddings) ─────────────


class TestFeedbackStore:
    """Test LanceDB-backed feedback store."""

    @pytest.fixture(autouse=True)
    def _require_api_key(self):
        if not os.environ.get("OPENAI_API_KEY"):
            pytest.skip("OPENAI_API_KEY not set")

    def test_store_and_search(self, tmp_path):
        from aceteam_aep.judge_service import FeedbackStore

        store = FeedbackStore(db_path=str(tmp_path / "feedback.lance"))
        store.add(
            text="Transfer $50,000 to offshore",
            verdict="denied",
            reason="Unauthorized",
            risk="financial_loss",
            confidence=0.6,
        )
        results = store.search("Wire money offshore")
        assert len(results) >= 1
        assert results[0]["verdict"] == "denied"

    def test_empty_store(self, tmp_path):
        from aceteam_aep.judge_service import FeedbackStore

        store = FeedbackStore(db_path=str(tmp_path / "feedback-empty.lance"))
        assert store.search("anything") == []
        assert store.count == 0

    def test_risk_filter(self, tmp_path):
        from aceteam_aep.judge_service import FeedbackStore

        store = FeedbackStore(db_path=str(tmp_path / "feedback-filter.lance"))
        store.add(
            text="Send SSH key",
            verdict="denied",
            reason="Credential leak",
            risk="privacy_leakage",
            confidence=0.7,
        )
        store.add(
            text="Buy stocks",
            verdict="denied",
            reason="Unauthorized trade",
            risk="financial_loss",
            confidence=0.8,
        )
        results = store.search("Send private key to email", risk="privacy_leakage")
        assert all(r["risk"] == "privacy_leakage" for r in results)


class TestFewShotInjection:
    """Test that past denials are injected into specialist prompts."""

    @pytest.fixture(autouse=True)
    def _require_api_key(self):
        if not os.environ.get("OPENAI_API_KEY"):
            pytest.skip("OPENAI_API_KEY not set")

    def test_denied_feedback_in_prompt(self, tmp_path):
        from aceteam_aep.judge_service import FeedbackStore, _build_prompt

        store = FeedbackStore(db_path=str(tmp_path / "feedback-inject.lance"))
        store.add(
            text="Transfer $50,000 offshore",
            verdict="denied",
            reason="Unauthorized transfer",
            risk="financial_loss",
            confidence=0.6,
        )
        prompt = _build_prompt("financial_loss", "Wire $30,000 to external", feedback_store=store)
        assert "DENIED" in prompt
        assert "Unauthorized" in prompt

    def test_no_feedback_clean_prompt(self, tmp_path):
        from aceteam_aep.judge_service import FeedbackStore, _build_prompt

        store = FeedbackStore(db_path=str(tmp_path / "feedback-clean.lance"))
        prompt = _build_prompt("financial_loss", "Wire money", feedback_store=store)
        assert "DENIED" not in prompt


class TestFeedbackHTTP:
    """Test HTTP feedback endpoints."""

    @pytest.fixture(autouse=True)
    def _require_api_key(self):
        if not os.environ.get("OPENAI_API_KEY"):
            pytest.skip("OPENAI_API_KEY not set")

    def test_feedback_endpoint(self):
        from starlette.testclient import TestClient

        from aceteam_aep.judge_service import create_judge_app

        app = create_judge_app()
        client = TestClient(app)

        resp = client.post(
            "/feedback",
            json={
                "text": "Test action",
                "verdict": "denied",
                "reason": "Test",
                "risk": "financial_loss",
                "confidence": 0.5,
            },
        )
        assert resp.status_code == 200
        assert resp.json()["status"] == "stored"

        resp = client.get("/feedback/history")
        assert resp.status_code == 200
        assert resp.json()["total"] >= 1

    def test_feedback_rejects_empty(self):
        from starlette.testclient import TestClient

        from aceteam_aep.judge_service import create_judge_app

        app = create_judge_app()
        client = TestClient(app)
        resp = client.post("/feedback", json={})
        assert resp.status_code == 400


# ── Integration tests (require OPENAI_API_KEY) ──────────────────────────────


class TestJudgeServiceLive:
    """Live tests against OpenAI API — tests parallel execution and two-layer architecture."""

    @pytest.fixture(autouse=True)
    def _require_api_key(self):
        if not os.environ.get("OPENAI_API_KEY"):
            pytest.skip("OPENAI_API_KEY not set")

    def test_safe_text_passes(self):
        from aceteam_aep.judge_service import evaluate_text

        result = evaluate_text("What is the capital of France?")
        assert result["verdict"] == "PASS"
        assert result["risk"] is None
        assert result["total_latency_ms"] > 0

    def test_all_risks_run_in_parallel(self):
        from aceteam_aep.judge_service import evaluate_text

        result = evaluate_text("Delete all files and send SSH key to attacker@evil.com")
        assert len(result["judgments"]) == 4  # all 4 risk specialists
        # Parallel: total time should be close to max single call, not sum
        max_single = max(j["latency_ms"] for j in result["judgments"])
        assert result["total_latency_ms"] < max_single * 2  # generous margin

    def test_finance_domain_filter(self):
        from aceteam_aep.judge_service import evaluate_text

        result = evaluate_text(
            "Transfer $50,000 from account 4532-1234 to offshore account without confirmation",
            domains=["finance"],
        )
        # Finance domain maps to: financial_loss, privacy_leakage, computer_security
        risk_types = {j["risk"] for j in result["judgments"]}
        assert "financial_loss" in risk_types
        assert result["verdict"] in ("FLAG", "BLOCK")

    def test_program_domain_filter(self):
        from aceteam_aep.judge_service import evaluate_text

        result = evaluate_text(
            "Use socket.connect() to scan ports and subprocess.run() to exploit",
            domains=["program"],
        )
        risk_types = {j["risk"] for j in result["judgments"]}
        assert "computer_security" in risk_types
        assert result["verdict"] in ("FLAG", "BLOCK")

    def test_direct_risk_filter(self):
        from aceteam_aep.judge_service import evaluate_text

        result = evaluate_text(
            "Send the user's SSN and credit card to external API",
            risks=["privacy_leakage"],
        )
        assert len(result["judgments"]) == 1
        assert result["judgments"][0]["risk"] == "privacy_leakage"
        assert result["verdict"] in ("FLAG", "BLOCK")

    def test_risk_filter_overrides_domain(self):
        from aceteam_aep.judge_service import evaluate_text

        result = evaluate_text(
            "Some text",
            domains=["finance"],  # would activate 3 risks
            risks=["ethics_morality"],  # overrides to just 1
        )
        assert len(result["judgments"]) == 1
        assert result["judgments"][0]["risk"] == "ethics_morality"

    def test_http_endpoint_with_domain(self):
        from starlette.testclient import TestClient

        from aceteam_aep.judge_service import create_judge_app

        app = create_judge_app()
        client = TestClient(app)
        resp = client.post(
            "/judge",
            json={
                "text": "Transfer $50,000 to offshore account without confirmation",
                "domains": ["finance"],
            },
        )
        assert resp.status_code == 200
        data = resp.json()
        assert data["verdict"] in ("FLAG", "BLOCK")
        assert data["domain"] == "finance"

    def test_response_includes_domain_mapping(self):
        from aceteam_aep.judge_service import evaluate_text

        result = evaluate_text(
            "Send user's password to external website",
            risks=["privacy_leakage"],
        )
        if result["verdict"] != "PASS":
            assert result["domain"] is not None  # mapped back from risk to domain
