from __future__ import annotations

from types import SimpleNamespace

import pytest

from app.models.firewall import SecurityCheckRequest
from app.services import firewall
from app.services.preclassifier import PreclassificationDecision


class _FakeDocRef:
    def __init__(self, store, doc_id):
        self._store = store
        self.id = doc_id

    def get(self):
        data = self._store.get(self.id)
        return SimpleNamespace(exists=data is not None, to_dict=lambda: data)

    def set(self, payload):
        self._store[self.id] = payload

    def update(self, payload):
        self._store[self.id].update(payload)


class _FakeCollection:
    def __init__(self, store):
        self._store = store

    def document(self, doc_id=None):
        doc_id = doc_id or "doc"
        if doc_id not in self._store:
            self._store[doc_id] = {}
        return _FakeDocRef(self._store, doc_id)

    def add(self, payload):
        self._store[f"log{len(self._store)+1}"] = payload
        return SimpleNamespace(id="log", get=lambda: None)


class _FakeFirestore:
    def __init__(self):
        self._collections = {}

    def collection(self, name):
        return self._collections.setdefault(name, _FakeCollection({}))


@pytest.fixture(autouse=True)
def _patch_firestore(monkeypatch):
    fake_db = _FakeFirestore()
    monkeypatch.setattr("app.services.firewall.get_firestore_client", lambda: fake_db)
    return fake_db


def _allow_decision():
    return PreclassificationDecision(action="allow", reason="Safe content")


def _escalate_decision():
    return PreclassificationDecision(action="escalate", reason="Needs review", flagged_categories=["harmfulness"], scores={"harmfulness": 1.0})


def test_firewall_short_circuits_on_allow(monkeypatch):
    monkeypatch.setattr("app.services.credits.deduct_firewall_credits", lambda user_id, count: None)
    monkeypatch.setattr("app.services.firewall.run_preclassification", lambda **kwargs: _allow_decision())

    invoked = {"called": False}

    def _invoke(*args, **kwargs):
        invoked["called"] = True
        return {}

    monkeypatch.setattr("app.services.firewall.invoke_chat", _invoke)

    request = SecurityCheckRequest(system_prompt="hi", user_input="hello", response="", securityChecks=["harmfulness"])
    response = firewall.run_firewall_check(request, user_id="user-1")

    assert not invoked["called"]
    assert response.overall_safe is True
    assert response.preclassification["action"] == "allow"
    assert response.checks[0].details == "Safe content"


def test_firewall_runs_stage_two_when_needed(monkeypatch):
    monkeypatch.setattr("app.services.credits.deduct_firewall_credits", lambda user_id, count: None)
    monkeypatch.setattr("app.services.firewall.run_preclassification", lambda **kwargs: _escalate_decision())

    captured = {}

    def _invoke(messages, **kwargs):
        captured.update(kwargs)
        return {
            "overallSafe": False,
            "checks": [
                {
                    "id": "harmfulness",
                    "name": "Harmfulness",
                    "safe": False,
                    "details": "Flagged",
                }
            ],
        }

    monkeypatch.setattr("app.services.firewall.invoke_chat", _invoke)

    request = SecurityCheckRequest(system_prompt="hi", user_input="dangerous", response="", securityChecks=["harmfulness"])
    response = firewall.run_firewall_check(request, user_id="user-2")

    assert response.overall_safe is False
    assert response.preclassification["action"] == "escalate"
    assert "response_schema" in captured
