"""Security primitives for request validation, PII detection and audit logging."""
from __future__ import annotations

from collections import defaultdict, deque
from dataclasses import dataclass, field
from datetime import UTC, datetime, timedelta
import re
import time
from typing import Any

from logos_adk.config import RateLimitConfig
from logos_adk.errors import ConfigError, StateBackendError
from logos_adk.observe import get_observability_store
from logos_adk.postgres_utils import pooled_connect, resolve_database_url, safe_identifier


@dataclass
class ValidationResult:
    """Result of input validation."""

    passed: bool
    message: str = ""
    metadata: dict[str, Any] = field(default_factory=dict)


@dataclass
class AuditRetentionPolicy:
    """Retention contract for audit data."""

    ttl_days: int = 30
    delete_mode: str = "soft"


@dataclass
class DeletionRequest:
    """Deletion request tracked for retention/deletion workflows."""

    subject_id: str
    scope: str
    requested_at: str
    reason: str | None = None
    hard_delete: bool = False


class InputValidator:
    """Validate raw prompt text before provider execution."""

    def __init__(
        self,
        *,
        min_length: int = 1,
        max_length: int = 20_000,
        reject_null_bytes: bool = True,
    ) -> None:
        self.min_length = min_length
        self.max_length = max_length
        self.reject_null_bytes = reject_null_bytes

    async def validate(self, prompt: str) -> ValidationResult:
        if not isinstance(prompt, str):
            return ValidationResult(False, "Prompt must be a string")
        if len(prompt.strip()) < self.min_length:
            return ValidationResult(False, "Prompt must not be empty")
        if len(prompt) > self.max_length:
            return ValidationResult(
                False,
                f"Prompt exceeds max length of {self.max_length} characters",
                {"length": len(prompt)},
            )
        if self.reject_null_bytes and "\x00" in prompt:
            return ValidationResult(False, "Prompt contains null bytes")
        return ValidationResult(True)


@dataclass
class PromptShieldResult:
    """Result of prompt injection analysis."""

    is_injection: bool
    confidence: float = 0.0
    reason: str = ""
    matched_patterns: list[str] = field(default_factory=list)


class PromptShield:
    """Heuristic prompt injection detector."""

    DEFAULT_PATTERNS = [
        r"ignore (all|any|previous|earlier) instructions",
        r"disregard (all|any|previous|earlier) instructions",
        r"reveal (your|the) system prompt",
        r"show (your|the) hidden prompt",
        r"developer message",
        r"bypass (guardrails|safety|security)",
        r"print.*(api key|secret|token|credentials)",
        r"exfiltrate",
    ]

    def __init__(self, patterns: list[str] | None = None, *, threshold: float = 0.5) -> None:
        self.threshold = threshold
        self._patterns = [re.compile(pattern, re.IGNORECASE) for pattern in (patterns or self.DEFAULT_PATTERNS)]

    _MAX_PROMPT_LENGTH = 100_000

    async def analyze(self, prompt: str) -> PromptShieldResult:
        # Truncate to prevent regex DoS on extremely long inputs
        text = prompt[:self._MAX_PROMPT_LENGTH] if len(prompt) > self._MAX_PROMPT_LENGTH else prompt
        matched = [pattern.pattern for pattern in self._patterns if pattern.search(text)]
        if not matched:
            return PromptShieldResult(is_injection=False)
        confidence = min(1.0, 0.35 + (0.2 * len(matched)))
        return PromptShieldResult(
            is_injection=confidence >= self.threshold,
            confidence=confidence,
            reason="Prompt matched suspicious instruction-override patterns",
            matched_patterns=matched,
        )


@dataclass
class PIIFinding:
    """Single PII finding."""

    pii_type: str
    match: str
    start: int
    end: int
    severity: str = "high"


@dataclass
class PIIScanResult:
    """PII scan result."""

    has_pii: bool
    blocked: bool
    findings: list[PIIFinding] = field(default_factory=list)


class PIIDetector:
    """Regex-based PII detector."""

    DEFAULT_PATTERNS = {
        "email": (re.compile(r"\b[A-Z0-9._%+-]+@[A-Z0-9.-]+\.[A-Z]{2,}\b", re.IGNORECASE), "medium"),
        "phone": (re.compile(r"\b(?:\+?\d{1,2}[-.\s]?)?(?:\(?\d{3}\)?[-.\s]?){2}\d{4}\b"), "medium"),
        "ssn": (re.compile(r"\b\d{3}-\d{2}-\d{4}\b"), "high"),
        # Linear-time credit card pattern: 13-16 digits with optional separators
        "credit_card": (re.compile(r"\b\d(?:[ -]?\d){12,15}\b"), "high"),
    }

    def __init__(self, *, blocked_types: set[str] | None = None) -> None:
        self.blocked_types = blocked_types or {"ssn", "credit_card"}

    _MAX_SCAN_LENGTH = 100_000

    async def scan(self, text: str) -> PIIScanResult:
        # Truncate to prevent regex DoS on extremely long inputs
        scan_text = text[:self._MAX_SCAN_LENGTH] if len(text) > self._MAX_SCAN_LENGTH else text
        findings: list[PIIFinding] = []
        for pii_type, (pattern, severity) in self.DEFAULT_PATTERNS.items():
            for match in pattern.finditer(scan_text):
                findings.append(
                    PIIFinding(
                        pii_type=pii_type,
                        match=match.group(),
                        start=match.start(),
                        end=match.end(),
                        severity=severity,
                    )
                )
        blocked = any(finding.pii_type in self.blocked_types for finding in findings)
        return PIIScanResult(has_pii=bool(findings), blocked=blocked, findings=findings)


class AuditLogger:
    """Structured audit logger backed by the observability store."""

    def __init__(
        self,
        source: str = "security",
        *,
        retention: AuditRetentionPolicy | None = None,
        redact_pii: bool = True,
    ) -> None:
        self.source = source
        self.retention = retention or AuditRetentionPolicy()
        self.redact_pii = redact_pii

    def _sanitize_attributes(self, attributes: dict[str, Any] | None) -> dict[str, Any]:
        payload = dict(attributes or {})
        findings = payload.get("findings")
        if not self.redact_pii or not isinstance(findings, list):
            return payload
        redacted: list[dict[str, Any]] = []
        for finding in findings:
            if isinstance(finding, dict):
                redacted.append({**finding, "match": "[REDACTED]"})
            else:
                redacted.append({"match": "[REDACTED]"})
        payload["findings"] = redacted
        return payload

    def log(
        self,
        *,
        event: str,
        outcome: str,
        agent_name: str | None = None,
        session_id: str | None = None,
        trace_id: str | None = None,
        request_id: str | None = None,
        attributes: dict[str, Any] | None = None,
    ) -> None:
        observability = get_observability_store()
        sanitized_attributes = self._sanitize_attributes(attributes)
        payload = {
            "event": event,
            "outcome": outcome,
            "agent_name": agent_name,
            "session_id": session_id,
            "retention_ttl_days": self.retention.ttl_days,
            "delete_mode": self.retention.delete_mode,
            **sanitized_attributes,
        }
        observability.log(
            level="warning" if outcome != "allowed" else "info",
            message=f"audit:{event}:{outcome}",
            source=self.source,
            agent_name=agent_name,
            session_id=session_id,
            trace_id=trace_id,
            request_id=request_id,
            attributes=payload,
        )

    def record_deletion_request(
        self,
        request: DeletionRequest,
        *,
        agent_name: str | None = None,
        request_id: str | None = None,
    ) -> None:
        self.log(
            event="deletion_request",
            outcome="accepted",
            agent_name=agent_name,
            request_id=request_id,
            attributes=request.__dict__,
        )


class RateLimitStore:
    """Abstract rate limit store."""

    async def allow(self, key: str, *, limit: int, window_seconds: int, burst: int) -> tuple[bool, float]:
        raise NotImplementedError

    async def close(self) -> None:
        return None


class InMemoryRateLimitStore(RateLimitStore):
    """In-memory sliding window store."""

    def __init__(self) -> None:
        self._events: dict[str, deque[float]] = defaultdict(deque)

    async def allow(self, key: str, *, limit: int, window_seconds: int, burst: int) -> tuple[bool, float]:
        now = time.time()
        window_start = now - window_seconds
        events = self._events[key]
        while events and events[0] < window_start:
            events.popleft()
        allowed = len(events) < (limit + burst)
        retry_after = 0.0
        if not allowed and events:
            retry_after = max(0.0, window_seconds - (now - events[0]))
        if allowed:
            events.append(now)
        return allowed, retry_after


class PostgreSQLRateLimitStore(RateLimitStore):
    """PostgreSQL-backed sliding window store."""

    def __init__(self, dsn: str, table_name: str = "rate_limit_events") -> None:
        self.dsn = dsn
        self.table_name = safe_identifier(table_name)

    def _connect(self) -> Any:
        try:
            return pooled_connect(self.dsn)
        except Exception as exc:
            raise StateBackendError(f"Failed to open PostgreSQL rate limit store: {exc}") from exc

    def _ensure_schema(self, connection: Any) -> None:
        with connection.cursor() as cursor:
            cursor.execute(
                f"""
                CREATE TABLE IF NOT EXISTS {self.table_name} (
                    rate_key TEXT NOT NULL,
                    occurred_at TIMESTAMPTZ NOT NULL
                )
                """
            )
            cursor.execute(
                f"CREATE INDEX IF NOT EXISTS idx_{self.table_name}_key_time "
                f"ON {self.table_name} (rate_key, occurred_at)"
            )
        connection.commit()

    async def allow(self, key: str, *, limit: int, window_seconds: int, burst: int) -> tuple[bool, float]:
        window_start = datetime.now(UTC) - timedelta(seconds=window_seconds)
        with self._connect() as connection:
            self._ensure_schema(connection)
            with connection.cursor() as cursor:
                cursor.execute(
                    f"DELETE FROM {self.table_name} WHERE occurred_at < %s",
                    (window_start,),
                )
                cursor.execute(
                    f"SELECT COUNT(*), MIN(occurred_at) FROM {self.table_name} WHERE rate_key = %s",
                    (key,),
                )
                count, earliest = cursor.fetchone()
                allowed = int(count or 0) < (limit + burst)
                retry_after = 0.0
                if allowed:
                    cursor.execute(
                        f"INSERT INTO {self.table_name} (rate_key, occurred_at) VALUES (%s, %s)",
                        (key, datetime.now(UTC)),
                    )
                elif earliest is not None:
                    retry_after = max(0.0, window_seconds - (datetime.now(UTC) - earliest).total_seconds())
            connection.commit()
        return allowed, retry_after


class RedisRateLimitStore(RateLimitStore):
    """Redis-backed sliding window store."""

    def __init__(self, url: str, namespace: str = "logos_adk_rate_limit") -> None:
        try:
            from redis import Redis
        except ImportError as exc:
            raise ImportError(
                "Redis rate limiting requires the 'redis' package. Install it with: pip install logos-adk[cache-redis]"
            ) from exc
        self._redis = Redis.from_url(url)
        self.namespace = namespace

    async def allow(self, key: str, *, limit: int, window_seconds: int, burst: int) -> tuple[bool, float]:
        now = time.time()
        bucket = f"{self.namespace}:{key}"
        window_start = now - window_seconds
        pipe = self._redis.pipeline()
        pipe.zremrangebyscore(bucket, 0, window_start)
        pipe.zcard(bucket)
        _, count = pipe.execute()
        allowed = int(count or 0) < (limit + burst)
        retry_after = 0.0
        if allowed:
            member = f"{now}:{time.monotonic_ns()}"
            self._redis.zadd(bucket, {member: now})
            self._redis.expire(bucket, window_seconds)
        else:
            earliest = self._redis.zrange(bucket, 0, 0, withscores=True)
            if earliest:
                retry_after = max(0.0, window_seconds - (now - float(earliest[0][1])))
        return allowed, retry_after

    async def close(self) -> None:
        try:
            self._redis.close()
        except Exception:
            return None


class RateLimiter:
    """Sliding-window rate limiter with pluggable backend."""

    def __init__(
        self,
        *,
        limit: int = 120,
        window_seconds: int = 60,
        burst: int = 30,
        store: RateLimitStore | None = None,
    ) -> None:
        self.limit = max(1, limit)
        self.window_seconds = max(1, window_seconds)
        self.burst = max(0, burst)
        self.store = store or InMemoryRateLimitStore()

    async def allow(self, key: str) -> tuple[bool, float]:
        return await self.store.allow(
            key,
            limit=self.limit,
            window_seconds=self.window_seconds,
            burst=self.burst,
        )

    async def close(self) -> None:
        await self.store.close()


def build_rate_limiter(config: RateLimitConfig | None = None) -> RateLimiter:
    """Build a rate limiter from config."""
    cfg = config or RateLimitConfig()
    backend = cfg.backend.lower()
    if backend in {"postgres", "postgresql"}:
        dsn = resolve_database_url(
            cfg.dsn,
            env_vars=("LOGOS_ADK_RATE_LIMIT_DSN", "LOGOS_ADK_DATABASE_URL", "DATABASE_URL"),
        )
        if not dsn:
            raise ConfigError("PostgreSQL rate limiting requires a DSN or DATABASE_URL")
        store: RateLimitStore = PostgreSQLRateLimitStore(dsn, table_name=cfg.table_name)
    elif backend == "redis":
        dsn = cfg.dsn or resolve_database_url(
            None,
            env_vars=("LOGOS_ADK_REDIS_URL", "REDIS_URL"),
        )
        if not dsn:
            raise ConfigError("Redis rate limiting requires a DSN or REDIS_URL")
        store = RedisRateLimitStore(dsn, namespace=cfg.namespace)
    elif backend == "memory":
        store = InMemoryRateLimitStore()
    else:
        raise ConfigError(f"Unsupported rate limit backend: {cfg.backend}")
    return RateLimiter(
        limit=cfg.limit_per_minute,
        window_seconds=cfg.window_seconds,
        burst=cfg.burst,
        store=store,
    )


__all__ = [
    "AuditRetentionPolicy",
    "AuditLogger",
    "DeletionRequest",
    "InputValidator",
    "InMemoryRateLimitStore",
    "PostgreSQLRateLimitStore",
    "PIIDetector",
    "PIIFinding",
    "PIIScanResult",
    "PromptShield",
    "PromptShieldResult",
    "RateLimiter",
    "RateLimitStore",
    "ValidationResult",
    "RedisRateLimitStore",
    "build_rate_limiter",
]
