from __future__ import annotations

from contextlib import contextmanager
from contextvars import ContextVar
from dataclasses import dataclass, field
from pathlib import Path
from typing import Dict, Iterable, List, Optional, Tuple
import os
import re
import shlex

from armorer.profiles import NetworkPolicy


_SECRET_PATTERNS = [
    re.compile(r"(sk-[A-Za-z0-9]{20,})"),
    re.compile(r"(AIza[0-9A-Za-z\-_]{30,})"),
    re.compile(r"(ghp_[A-Za-z0-9]{20,})"),
    re.compile(r"(xox[baprs]-[A-Za-z0-9-]{10,})"),
    re.compile(r"(?i)(api[_-]?key\s*[=:]\s*['\"]?[^\s'\"]+)"),
    re.compile(r"(?i)(password\s*[=:]\s*['\"]?[^\s'\"]+)"),
]

_PROMPT_INJECTION_PATTERNS = [
    re.compile(r"(?i)ignore\s+all\s+previous\s+instructions"),
    re.compile(r"(?i)reveal\s+system\s+prompt"),
    re.compile(r"(?i)disable\s+security"),
    re.compile(r"(?i)run\s+arbitrary\s+command"),
]

_DENY_COMMAND_PATTERNS = [
    re.compile(r"(^|\s)rm\s+-rf\s+/"),
    re.compile(r"(^|\s)mkfs\b"),
    re.compile(r"(^|\s)shutdown\b"),
    re.compile(r"(^|\s)reboot\b"),
    re.compile(r"(^|\s)curl\s+.*\|\s*(sh|bash)\b"),
    re.compile(r"(^|\s)wget\s+.*\|\s*(sh|bash)\b"),
    re.compile(r"(^|\s)chmod\s+777\s+/"),
    re.compile(r"(^|\s)chown\s+-R\s+root\s+/"),
]

_NETWORK_COMMAND_PATTERNS = [
    re.compile(r"(^|\s)(curl|wget|nc|nmap|telnet|httpie)\b"),
    re.compile(r"(^|\s)(pip|npm|pnpm|yarn|go|cargo|git)\s+(install|clone|fetch|pull)\b"),
    re.compile(r"(^|\s)docker\s+(pull|search|buildx)\b"),
]


@dataclass
class PolicyDecision:
    allowed: bool
    reason: str = ""


@dataclass
class PolicyGuard:
    project_root: Path
    allowed_paths: List[Path]
    network_policy: NetworkPolicy = NetworkPolicy.OPEN
    command_allowlist: List[re.Pattern[str]] = field(default_factory=list)
    command_denylist: List[re.Pattern[str]] = field(default_factory=lambda: list(_DENY_COMMAND_PATTERNS))
    allow_docker_socket: bool = False
    install_scope: Optional[Path] = None
    instance_name: Optional[str] = None

    @classmethod
    def default(cls) -> "PolicyGuard":
        root = Path.cwd().resolve()
        return cls(
            project_root=root,
            allowed_paths=[root, root / "installed", Path.home() / ".armorer"],
            network_policy=NetworkPolicy.OPEN,
        )

    @classmethod
    def for_installation(
        cls,
        project_root: Path,
        install_dir: Path,
        network_policy: NetworkPolicy,
        extra_paths: Optional[Iterable[Path]] = None,
        allow_docker_socket: bool = False,
        instance_name: Optional[str] = None,
    ) -> "PolicyGuard":
        allowed = [project_root.resolve(), install_dir.resolve(), Path.home() / ".armorer"]
        if extra_paths:
            allowed.extend(Path(p).resolve() for p in extra_paths)
        allow_cmds = [
            re.compile(r"^docker\b"),
            re.compile(r"^git\b"),
            re.compile(r"^(ls|cat|grep|rg|find|pwd|echo|sed|awk|head|tail|env|printenv|python3?)\b"),
            re.compile(r"^(chmod|chown|cp|mv|mkdir|touch)\b"),
            re.compile(r"^(make|npm|pnpm|yarn|pip|uv)\b"),
            re.compile(r"^(bash|sh)\b"),
        ]
        return cls(
            project_root=project_root.resolve(),
            allowed_paths=allowed,
            network_policy=network_policy,
            command_allowlist=allow_cmds,
            allow_docker_socket=allow_docker_socket,
            install_scope=install_dir.resolve(),
            instance_name=(instance_name or "").strip() or None,
        )

    def sanitize_text(self, text: str) -> str:
        out = text
        for pat in _SECRET_PATTERNS:
            out = pat.sub("[REDACTED]", out)
        return out

    def filter_prompt_injection(self, text: str) -> Tuple[str, List[str]]:
        alerts: List[str] = []
        filtered = text
        for pat in _PROMPT_INJECTION_PATTERNS:
            if pat.search(filtered):
                alerts.append(pat.pattern)
                filtered = pat.sub("[BLOCKED_INJECTION_PATTERN]", filtered)
        return filtered, alerts

    def check_path(self, path: Path, op: str = "read") -> PolicyDecision:
        p = path.resolve()
        for allowed in self.allowed_paths:
            try:
                p.relative_to(allowed)
                return PolicyDecision(True)
            except Exception:
                continue
        return PolicyDecision(False, f"{op} denied for path: {p}")

    def _is_within(self, base: Path, p: Path) -> bool:
        try:
            p.resolve().relative_to(base.resolve())
            return True
        except Exception:
            return False

    def _check_docker_scope(self, command: str, cwd: Optional[Path]) -> PolicyDecision:
        """Prevent cross-instance destructive docker actions during install-time repair."""
        if not self.instance_name:
            return PolicyDecision(True)
        cmd = command.strip()
        if not cmd.lower().startswith("docker"):
            return PolicyDecision(True)

        try:
            tokens = shlex.split(cmd)
        except Exception:
            return PolicyDecision(False, "invalid docker command syntax")

        # Block direct destructive container control; require compose-scoped flows.
        if len(tokens) >= 2 and tokens[1] in {"rm", "stop", "kill", "restart"}:
            return PolicyDecision(False, "direct docker container control is blocked during scoped installation")

        # For docker compose commands, enforce project/path scope.
        if len(tokens) >= 3 and tokens[1] == "compose":
            project_name = None
            for i, tok in enumerate(tokens):
                if tok in {"-p", "--project-name"} and i + 1 < len(tokens):
                    project_name = tokens[i + 1]
                    break
            if project_name and project_name != self.instance_name:
                return PolicyDecision(False, f"compose project '{project_name}' is outside current instance scope")

            if not project_name:
                if self.install_scope is None:
                    return PolicyDecision(False, "compose command missing scoped project context")
                if cwd is None or not self._is_within(self.install_scope, cwd):
                    return PolicyDecision(False, "compose command must run within current instance install directory")

        return PolicyDecision(True)

    def check_command(self, command: str, cwd: Optional[Path] = None) -> PolicyDecision:
        cmd = command.strip()
        normalized = re.sub(r"^([A-Za-z_][A-Za-z0-9_]*=[^\\s]+\\s+)+", "", cmd)
        for pat in self.command_denylist:
            if pat.search(cmd):
                return PolicyDecision(False, f"command denied by security policy ({pat.pattern})")

        if self.command_allowlist:
            if not any(pat.search(normalized) for pat in self.command_allowlist):
                return PolicyDecision(False, "command not in allowlist")

        if self.network_policy == NetworkPolicy.NONE:
            if any(pat.search(cmd) for pat in _NETWORK_COMMAND_PATTERNS):
                return PolicyDecision(False, "network command denied by offline policy")

        if self.network_policy == NetworkPolicy.WHITELIST:
            if re.search(r"\b(curl|wget|httpie|nc|nmap|telnet)\b", cmd):
                return PolicyDecision(False, "direct network utility blocked under scoped network policy")

        docker_scope = self._check_docker_scope(cmd, cwd=cwd)
        if not docker_scope.allowed:
            return docker_scope

        return PolicyDecision(True)


_POLICY_CTX: ContextVar[PolicyGuard] = ContextVar("armorer_policy_guard", default=PolicyGuard.default())


def get_policy_guard() -> PolicyGuard:
    return _POLICY_CTX.get()


def set_policy_guard(guard: PolicyGuard):
    return _POLICY_CTX.set(guard)


def reset_policy_guard(token) -> None:
    _POLICY_CTX.reset(token)


@contextmanager
def policy_scope(guard: PolicyGuard):
    token = set_policy_guard(guard)
    try:
        yield
    finally:
        reset_policy_guard(token)
