from __future__ import annotations

import asyncio
import logging
import re
import threading
import time
from dataclasses import dataclass
from typing import Any, Dict, Literal, Mapping, Optional, Tuple

import httpx

from .errors import SecurityViolationException
from .pii import PIIManager

logger = logging.getLogger("agentid")
AGENTID_SDK_VERSION_HEADER = "py-1.1.0"

PYTHON_GENERAL_RCE_PATTERN = re.compile(
    r"(import\s+(os|sys|subprocess)|from\s+(os|sys|subprocess)\s+import|exec\s*\(|eval\s*\(|__import__)",
    re.IGNORECASE,
)
SHELL_BASH_RCE_PATTERN = re.compile(
    r"(wget\s+|curl\s+|rm\s+-rf|chmod\s+\+x|cat\s+\/etc\/passwd|\/bin\/sh|\/bin\/bash)",
    re.IGNORECASE,
)
JAVASCRIPT_RCE_PATTERN = re.compile(
    r"(new\s+Function\(|process\.env|child_process)",
    re.IGNORECASE,
)
SQL_DATABASE_ACCESS_PATTERN = re.compile(
    r"\b(SELECT|INSERT|UPDATE|DELETE|DROP|UNION|ALTER)\b[\s\S]+?\b(FROM|INTO|TABLE|DATABASE|VIEW|INDEX)\b",
    re.IGNORECASE,
)
_PLACEHOLDER_PATTERN = re.compile(r"<[A-Z_]+_\d+>")

ViolationType = Literal[
    "PII_LEAKAGE_STRICT",
    "SQL_INJECTION_ATTEMPT",
    "RCE_ATTEMPT",
]
ActionTaken = Literal["BLOCKED", "REDACTED"]


@dataclass(frozen=True)
class CapabilityFlags:
    block_pii_leakage: bool
    block_db_access: bool
    block_code_execution: bool
    block_toxicity: bool
    strict_security_mode: bool = False
    failure_mode: Literal["fail_open", "fail_close"] = "fail_open"


@dataclass(frozen=True)
class _CapabilityCacheEntry:
    flags: CapabilityFlags
    expires_at: float


@dataclass(frozen=True)
class LocalSecurityEvent:
    violation_type: ViolationType
    action_taken: ActionTaken


@dataclass(frozen=True)
class LocalSecurityEnforcementResult:
    sanitized_input: str
    events: Tuple[LocalSecurityEvent, ...]


DEFAULT_STRICT_CAPABILITY_FLAGS = CapabilityFlags(
    block_pii_leakage=True,
    block_db_access=True,
    block_code_execution=True,
    block_toxicity=True,
    strict_security_mode=True,
    failure_mode="fail_close",
)

DEFAULT_FAIL_OPEN_CAPABILITY_FLAGS = CapabilityFlags(
    block_pii_leakage=False,
    block_db_access=False,
    block_code_execution=False,
    block_toxicity=False,
    strict_security_mode=False,
    failure_mode="fail_open",
)

DEFAULT_TIMEOUT_SECONDS = 2.0
DEFAULT_TTL_SECONDS = 300.0

_CONFIG_CACHE: Dict[str, _CapabilityCacheEntry] = {}
_CACHE_STATE_LOCK = threading.Lock()
_SYNC_FETCH_LOCKS: Dict[str, threading.Lock] = {}
_ASYNC_FETCH_LOCKS: Dict[tuple[str, int], asyncio.Lock] = {}


class LocalPolicyBlockError(SecurityViolationException):
    def __init__(
        self,
        *,
        violation_type: ViolationType,
        action_taken: ActionTaken = "BLOCKED",
        message: Optional[str] = None,
    ) -> None:
        self.violation_type = violation_type
        self.action_taken = action_taken
        super().__init__(message or f"AgentID: Security policy blocked ({violation_type})")


class StrictPiiStreamingError(ValueError):
    def __init__(self) -> None:
        self.violation_type: ViolationType = "PII_LEAKAGE_STRICT"
        self.action_taken: ActionTaken = "BLOCKED"
        super().__init__(
            "AgentID: Streaming is not supported when Strict PII Mode is enabled. "
            "Please disable streaming or adjust security settings."
        )


class LocalSecurityEnforcer:
    def __init__(self, pii_manager: PIIManager):
        self._pii_manager = pii_manager

    def enforce(
        self,
        *,
        input_text: str,
        stream: bool,
        flags: CapabilityFlags,
    ) -> LocalSecurityEnforcementResult:
        violation_type = detect_capability_violation(input_text, flags)
        if violation_type is not None:
            raise LocalPolicyBlockError(violation_type=violation_type)

        if not flags.block_pii_leakage:
            return LocalSecurityEnforcementResult(
                sanitized_input=input_text,
                events=(),
            )

        if stream:
            # We must not silently mutate stream=True to stream=False.
            # That would change the return type (Generator -> Object) and break caller iteration.
            raise StrictPiiStreamingError()

        redacted_text, was_redacted = redact_pii_strict(self._pii_manager, input_text)
        events: Tuple[LocalSecurityEvent, ...] = ()
        if was_redacted:
            events = (
                LocalSecurityEvent(
                    violation_type="PII_LEAKAGE_STRICT",
                    action_taken="REDACTED",
                ),
            )

        return LocalSecurityEnforcementResult(
            sanitized_input=redacted_text,
            events=events,
        )


def _cache_key(api_key: str, base_url: str) -> str:
    return f"{base_url}|{api_key}"


def _is_fresh(entry: _CapabilityCacheEntry) -> bool:
    return entry.expires_at > time.monotonic()


def _read_bool(raw: Mapping[str, Any], primary_key: str, alias_key: str) -> bool:
    if primary_key in raw:
        value = raw.get(primary_key)
        if isinstance(value, bool):
            return value
        raise ValueError(f"Invalid capability flag: {primary_key}")
    if alias_key in raw:
        value = raw.get(alias_key)
        if isinstance(value, bool):
            return value
        raise ValueError(f"Invalid capability flag: {alias_key}")
    raise ValueError(f"Missing capability flag: {primary_key}")


def _normalize_flags(payload: Mapping[str, Any] | None) -> CapabilityFlags:
    if payload is None:
        raise ValueError("Missing capability config payload")
    raw = payload
    block_pii_leakage = _read_bool(raw, "block_pii_leakage", "block_pii")
    block_db_access = _read_bool(raw, "block_db_access", "block_db")
    block_code_execution = _read_bool(raw, "block_code_execution", "block_code")
    block_toxicity = _read_bool(raw, "block_toxicity", "block_toxic")
    strict_security_mode = bool(raw.get("strict_security_mode")) if isinstance(raw.get("strict_security_mode"), bool) else False
    failure_mode_raw = raw.get("failure_mode")
    failure_mode: Literal["fail_open", "fail_close"] = (
        failure_mode_raw
        if failure_mode_raw in {"fail_open", "fail_close"}
        else ("fail_close" if strict_security_mode else "fail_open")
    )
    effective_strict_mode = strict_security_mode or failure_mode == "fail_close"
    return CapabilityFlags(
        block_pii_leakage=block_pii_leakage,
        block_db_access=block_db_access,
        block_code_execution=block_code_execution,
        block_toxicity=block_toxicity,
        strict_security_mode=effective_strict_mode,
        failure_mode="fail_close" if effective_strict_mode else "fail_open",
    )


def _fetch_capability_flags_sync(
    *,
    api_key: str,
    base_url: str,
    timeout_seconds: float,
) -> CapabilityFlags:
    with httpx.Client(
        base_url=base_url,
        headers={
            "Content-Type": "application/json",
            "x-agentid-api-key": api_key,
            "X-AgentID-SDK-Version": AGENTID_SDK_VERSION_HEADER,
        },
        timeout=httpx.Timeout(timeout_seconds),
    ) as client:
        response = client.get("/agent/config")
        response.raise_for_status()
        body = response.json()
        if not isinstance(body, dict):
            raise ValueError("Invalid capability config payload")
        return _normalize_flags(body)


async def _fetch_capability_flags_async(
    *,
    api_key: str,
    base_url: str,
    timeout_seconds: float,
) -> CapabilityFlags:
    async with httpx.AsyncClient(
        base_url=base_url,
        headers={
            "Content-Type": "application/json",
            "x-agentid-api-key": api_key,
            "X-AgentID-SDK-Version": AGENTID_SDK_VERSION_HEADER,
        },
        timeout=httpx.Timeout(timeout_seconds),
    ) as client:
        response = await client.get("/agent/config")
        response.raise_for_status()
        body = response.json()
        if not isinstance(body, dict):
            raise ValueError("Invalid capability config payload")
        return _normalize_flags(body)


def _get_cached_entry(key: str) -> Optional[_CapabilityCacheEntry]:
    with _CACHE_STATE_LOCK:
        return _CONFIG_CACHE.get(key)


def _store_cache(key: str, flags: CapabilityFlags, ttl_seconds: float) -> None:
    with _CACHE_STATE_LOCK:
        _CONFIG_CACHE[key] = _CapabilityCacheEntry(
            flags=flags,
            expires_at=time.monotonic() + ttl_seconds,
        )


def _get_sync_fetch_lock(key: str) -> threading.Lock:
    with _CACHE_STATE_LOCK:
        lock = _SYNC_FETCH_LOCKS.get(key)
        if lock is None:
            lock = threading.Lock()
            _SYNC_FETCH_LOCKS[key] = lock
        return lock


def _get_async_fetch_lock(key: str) -> asyncio.Lock:
    loop = asyncio.get_running_loop()
    lock_key = (key, id(loop))
    with _CACHE_STATE_LOCK:
        lock = _ASYNC_FETCH_LOCKS.get(lock_key)
        if lock is None:
            lock = asyncio.Lock()
            _ASYNC_FETCH_LOCKS[lock_key] = lock
        return lock


def ensure_capability_flags_sync(
    *,
    api_key: str,
    base_url: str,
    timeout_seconds: float = DEFAULT_TIMEOUT_SECONDS,
    ttl_seconds: float = DEFAULT_TTL_SECONDS,
) -> CapabilityFlags:
    key = _cache_key(api_key, base_url)
    cached = _get_cached_entry(key)
    if cached is not None and _is_fresh(cached):
        return cached.flags

    lock = _get_sync_fetch_lock(key)
    with lock:
        cached = _get_cached_entry(key)
        stale_flags = cached.flags if cached is not None else None
        if cached is not None and _is_fresh(cached):
            return cached.flags

        try:
            fetched = _fetch_capability_flags_sync(
                api_key=api_key,
                base_url=base_url,
                timeout_seconds=timeout_seconds,
            )
        except Exception as exc:
            logger.warning(
                "AgentID Config unreachable (%s). Defaulting to FAIL-OPEN MODE.",
                exc.__class__.__name__,
            )
            fetched = stale_flags or DEFAULT_FAIL_OPEN_CAPABILITY_FLAGS

        _store_cache(key, fetched, ttl_seconds)
        return fetched


async def ensure_capability_flags_async(
    *,
    api_key: str,
    base_url: str,
    timeout_seconds: float = DEFAULT_TIMEOUT_SECONDS,
    ttl_seconds: float = DEFAULT_TTL_SECONDS,
) -> CapabilityFlags:
    key = _cache_key(api_key, base_url)
    cached = _get_cached_entry(key)
    if cached is not None and _is_fresh(cached):
        return cached.flags

    lock = _get_async_fetch_lock(key)
    async with lock:
        cached = _get_cached_entry(key)
        stale_flags = cached.flags if cached is not None else None
        if cached is not None and _is_fresh(cached):
            return cached.flags

        try:
            fetched = await _fetch_capability_flags_async(
                api_key=api_key,
                base_url=base_url,
                timeout_seconds=timeout_seconds,
            )
        except Exception as exc:
            logger.warning(
                "AgentID Config unreachable (%s). Defaulting to FAIL-OPEN MODE.",
                exc.__class__.__name__,
            )
            fetched = stale_flags or DEFAULT_FAIL_OPEN_CAPABILITY_FLAGS

        _store_cache(key, fetched, ttl_seconds)
        return fetched


def get_or_fetch_capability_flags(
    *,
    api_key: str,
    base_url: str,
    timeout_seconds: float = DEFAULT_TIMEOUT_SECONDS,
    ttl_seconds: float = DEFAULT_TTL_SECONDS,
) -> CapabilityFlags:
    return ensure_capability_flags_sync(
        api_key=api_key,
        base_url=base_url,
        timeout_seconds=timeout_seconds,
        ttl_seconds=ttl_seconds,
    )


def get_cached_capability_flags(*, api_key: str, base_url: str) -> CapabilityFlags:
    key = _cache_key(api_key, base_url)
    cached = _get_cached_entry(key)
    if cached is not None:
        return cached.flags
    return DEFAULT_FAIL_OPEN_CAPABILITY_FLAGS


def detect_capability_violation(
    input_text: str,
    flags: CapabilityFlags,
) -> Optional[ViolationType]:
    if flags.block_db_access and SQL_DATABASE_ACCESS_PATTERN.search(input_text):
        return "SQL_INJECTION_ATTEMPT"

    if flags.block_code_execution and (
        PYTHON_GENERAL_RCE_PATTERN.search(input_text)
        or SHELL_BASH_RCE_PATTERN.search(input_text)
        or JAVASCRIPT_RCE_PATTERN.search(input_text)
    ):
        return "RCE_ATTEMPT"
    return None


def redact_pii_strict(pii_manager: PIIManager, text: str) -> Tuple[str, bool]:
    if not text:
        return text, False

    masked_text, mapping = pii_manager.anonymize(text)
    redacted = masked_text

    if mapping:
        for placeholder in sorted(mapping.keys(), key=len, reverse=True):
            redacted = redacted.replace(placeholder, "[REDACTED]")

    redacted = _PLACEHOLDER_PATTERN.sub("[REDACTED]", redacted)
    return redacted, redacted != text


def build_security_policy_violation_metadata(
    *,
    system_id: str,
    violation_type: ViolationType,
    action_taken: ActionTaken,
) -> Dict[str, str]:
    return {
        "event_type": "security_policy_violation",
        "severity": "high",
        "system_id": system_id,
        "violation_type": violation_type,
        "input_snippet": "[REDACTED_SAMPLE]",
        "action_taken": action_taken,
    }
