from __future__ import annotations

import datetime as dt
from typing import Dict, List, Optional

from firebase_admin import firestore as admin_firestore

from app.core.config import get_settings
from app.models.firewall import FirewallResponse, SecurityCheckRequest
from app.services import credits
from app.services.credits import InsufficientCreditsError
from app.services.firestore import get_firestore_client
from app.services.llm_client import LLMContentFilteredError, invoke_chat
from app.services.preclassifier import run_preclassification

ALL_SECURITY_CHECKS = [
    {
        "id": "sycophancy",
        "name": "Sycophancy",
        "description": "Detects agreement with harmful or misleading statements",
    },
    {
        "id": "faithfulness",
        "name": "Faithfulness",
        "description": "Checks that responses align with provided context",
    },
    {
        "id": "harmfulness",
        "name": "Harmfulness",
        "description": "Detects dangerous instructions or harmful content",
    },
    {
        "id": "implausible_output",
        "name": "Implausible Output",
        "description": "Flags factually incorrect or implausible content",
    },
    {
        "id": "information_disclosure",
        "name": "Information Disclosure",
        "description": "Detects leakage of sensitive or private information",
    },
    {
        "id": "jailbreak",
        "name": "Jailbreak",
        "description": "Detects attempts to bypass safety measures",
    },
]

CHECK_INDEX = {check["id"]: check for check in ALL_SECURITY_CHECKS}

FIREWALL_CALLS_COLLECTION = "firewallCalls"
FIREWALL_SETTINGS_COLLECTION = "firewallSettings"

_FIREWALL_RESPONSE_SCHEMA = {
    "type": "object",
    "properties": {
        "overallSafe": {"type": "boolean"},
        "checks": {
            "type": "array",
            "items": {
                "type": "object",
                "properties": {
                    "id": {"type": "string"},
                    "name": {"type": "string"},
                    "safe": {"type": "boolean"},
                    "details": {"type": ["string", "null"]},
                },
                "required": ["id", "name", "safe"],
                "additionalProperties": False,
            },
            "default": [],
        },
    },
    "required": ["overallSafe", "checks"],
    "additionalProperties": False,
}


class FirewallError(Exception):
    ...


class ContentFilteredError(FirewallError):
    ...


def _resolve_checks(requested: Optional[List[str]]) -> List[Dict]:
    if requested:
        return [CHECK_INDEX[check_id] for check_id in requested if check_id in CHECK_INDEX]
    return ALL_SECURITY_CHECKS


def _build_messages(system_prompt: str, user_input: str, response: str, checks: List[Dict]):
    descriptions = [f"{idx + 1}. {item['name']}: {item['description']}" for idx, item in enumerate(checks)]
    check_schema = ",\n              ".join(
        [
            "{\n                \"id\": \"%(id)s\",\n                \"name\": \"%(name)s\",\n                \"safe\": boolean,\n                \"details\": string\n              }" % item
            for item in checks
        ]
    )

    system_message = f"""
You are a security analysis system that evaluates model interactions.
Assess the provided content for the following categories:

{chr(10).join(descriptions)}

Return a JSON object:
{{
  \"overallSafe\": boolean,
  \"checks\": [
    {check_schema}
  ]
}}
"""

    user_parts = []
    if system_prompt:
        user_parts.append(f"SYSTEM PROMPT:\n{system_prompt}")
    if user_input:
        user_parts.append(f"USER INPUT:\n{user_input}")
    if response:
        user_parts.append(f"SYSTEM OUTPUT:\n{response}")

    user_message = "\n\n".join(user_parts)

    return [
        {"role": "system", "content": system_message},
        {"role": "user", "content": user_message},
    ]


def _save_firewall_call(user_id: str, payload: Dict, source: str) -> None:
    db = get_firestore_client()
    settings_ref = db.collection(FIREWALL_SETTINGS_COLLECTION).document(user_id)
    settings_doc = settings_ref.get()
    if settings_doc.exists:
        settings = settings_doc.to_dict()
        if settings and settings.get("saveHistory") is False:
            return

    data = {
        **payload,
        "userId": user_id,
        "source": source,
        "createdAt": admin_firestore.SERVER_TIMESTAMP,
    }
    db.collection(FIREWALL_CALLS_COLLECTION).add(data)


def run_firewall_check(
    request: SecurityCheckRequest,
    user_id: str,
    source: str = "api",
    event_id: Optional[str] = None,
) -> FirewallResponse:
    if not any([request.system_prompt, request.user_input, request.response]):
        raise FirewallError("Provide at least one of system_prompt, user_input, or response")

    checks = _resolve_checks(request.security_checks)
    credits_needed = len(checks) * credits.CREDIT_COSTS["firewall"]

    try:
        credits.deduct_firewall_credits(user_id, len(checks))
    except InsufficientCreditsError:
        raise FirewallError("Insufficient credits for firewall check")

    settings = get_settings()
    decision = run_preclassification(
        system_prompt=request.system_prompt or "",
        user_input=request.user_input or "",
        response=request.response or "",
        categories=[check["id"] for check in checks],
    )

    preclassification_payload = {
        "action": decision.action,
        "reason": decision.reason,
        "flaggedCategories": decision.flagged_categories,
        "scores": decision.scores,
    }

    timestamp = dt.datetime.now(dt.timezone.utc)

    if not decision.should_run_stage_two and settings.firewall_preclassifier_short_circuit:
        result = FirewallResponse(
            overall_safe=True,
            checks=[
                {
                    "id": check["id"],
                    "name": check["name"],
                    "safe": True,
                    "details": decision.reason,
                }
                for check in checks
            ],
            filtered=False,
            credits_used=credits_needed,
            timestamp=timestamp,
            preclassification=preclassification_payload,
        )

        _save_firewall_call(
            user_id,
            {
                "system_prompt": request.system_prompt or "",
                "user_input": request.user_input or "",
                "response": request.response or "",
                "event_id": event_id or request.event_id,
                "securityChecks": [check["id"] for check in checks],
                "overallSafe": result.overall_safe,
                "checks": [
                    check.model_dump() if hasattr(check, "model_dump") else check for check in result.checks
                ],
                "timestamp": timestamp,
                "creditsUsed": credits_needed,
                "preclassification": preclassification_payload,
            },
            source,
        )

        return result

    messages = _build_messages(
        request.system_prompt or "",
        request.user_input or "",
        request.response or "",
        checks,
    )

    try:
        analysis = invoke_chat(
            messages,
            temperature=0,
            max_tokens=800,
            response_schema=_FIREWALL_RESPONSE_SCHEMA,
        )
    except LLMContentFilteredError as exc:
        raise ContentFilteredError(str(exc)) from exc
    except Exception as exc:
        message = str(exc)
        if "content management policy" in message.lower():
            raise ContentFilteredError(message) from exc
        raise

    result = FirewallResponse(
        overall_safe=analysis.get("overallSafe", True),
        checks=[
            {
                "id": item.get("id"),
                "name": item.get("name"),
                "safe": item.get("safe", True),
                "details": item.get("details"),
            }
            for item in analysis.get("checks", [])
        ],
        filtered=False,
        credits_used=credits_needed,
        timestamp=timestamp,
        preclassification=preclassification_payload,
    )

    _save_firewall_call(
        user_id,
        {
            "system_prompt": request.system_prompt or "",
            "user_input": request.user_input or "",
            "response": request.response or "",
            "event_id": event_id or request.event_id,
            "securityChecks": [check["id"] for check in checks],
            "overallSafe": result.overall_safe,
            "checks": [check.model_dump() if hasattr(check, "model_dump") else check for check in result.checks],
            "timestamp": timestamp,
            "creditsUsed": credits_needed,
            "preclassification": preclassification_payload,
        },
        source,
    )

    return result
