from __future__ import annotations

from dataclasses import dataclass
import re
from typing import Iterable, Protocol

from .models import Artifact, Finding, Severity


class Rule(Protocol):
    rule_id: str
    title: str
    description: str
    default_severity: Severity

    def evaluate(self, artifact: Artifact) -> Iterable[Finding]: ...


def _location(text: str, offset: int) -> tuple[int, int]:
    line = text.count("\n", 0, offset) + 1
    last_break = text.rfind("\n", 0, offset)
    return line, offset - last_break


@dataclass(frozen=True, slots=True)
class RegexSignal:
    name: str
    pattern: re.Pattern[str]
    severity: Severity


class SecretRule:
    rule_id = "AAISA-SECRET-001"
    title = "Potential embedded secret"
    description = "Finds credential-shaped values and reports only a redacted signal."
    default_severity = Severity.HIGH
    artifact_kinds = ("text", "configuration", "source-code")
    standards = ("CWE-798", "NIST AI 600-1")
    references = (
        "https://cwe.mitre.org/data/definitions/798.html",
        "https://doi.org/10.6028/NIST.AI.600-1",
    )
    signals = (
        RegexSignal("private key", re.compile(r"-----BEGIN (?:RSA |EC |OPENSSH )?PRIVATE KEY-----"), Severity.CRITICAL),
        RegexSignal("AWS access key", re.compile(r"\bAKIA[0-9A-Z]{16}\b"), Severity.HIGH),
        RegexSignal("GitHub token", re.compile(r"\bgh[pousr]_[A-Za-z0-9]{20,255}\b"), Severity.HIGH),
        RegexSignal(
            "credential assignment",
            re.compile(r"(?im)^\s*(?:api[_-]?key|secret|token|password)\s*[:=]\s*['\"]?([^\s'\"#]{12,})"),
            Severity.HIGH,
        ),
    )

    def evaluate(self, artifact: Artifact) -> Iterable[Finding]:
        if artifact.text is None:
            return ()
        findings: list[Finding] = []
        for signal in self.signals:
            for match in signal.pattern.finditer(artifact.text):
                line, column = _location(artifact.text, match.start())
                findings.append(Finding(
                    self.rule_id,
                    self.title,
                    signal.severity,
                    f"Detected a {signal.name} pattern. The matched value is redacted.",
                    artifact.relative_path,
                    line,
                    column,
                    f"{signal.name}: [REDACTED]",
                    "Remove the value, rotate it if it may be live, and use an approved secret store.",
                    ("credentials", "redaction"),
                ))
        return findings


class PromptInjectionRule:
    rule_id = "AAISA-PROMPT-001"
    title = "Prompt injection or tool-exfiltration instruction"
    description = "Finds instructions commonly used to override controls or disclose protected context."
    default_severity = Severity.MEDIUM
    artifact_kinds = ("prompt", "configuration", "document", "dataset-metadata")
    standards = ("MITRE ATLAS AML.T0051", "OWASP Agentic AI Threats and Mitigations")
    references = (
        "https://atlas.mitre.org/techniques/AML.T0051",
        "https://genai.owasp.org/resource/agentic-ai-threats-and-mitigations/",
    )
    pattern = re.compile(
        r"(?i)\b(ignore|disregard|override)\b.{0,50}\b(previous|prior|system|developer)\b.{0,30}\b(instruction|prompt|message)s?\b"
        r"|\b(reveal|print|exfiltrate|send)\b.{0,50}\b(system prompt|secret|credential|token)s?\b",
        re.DOTALL,
    )

    def evaluate(self, artifact: Artifact) -> Iterable[Finding]:
        if artifact.text is None or artifact.kind not in {"prompt", "configuration", "document", "dataset-metadata"}:
            return ()
        results: list[Finding] = []
        for match in self.pattern.finditer(artifact.text):
            line, column = _location(artifact.text, match.start())
            results.append(Finding(
                self.rule_id,
                self.title,
                Severity.MEDIUM,
                "Artifact contains language associated with instruction override or protected-context disclosure.",
                artifact.relative_path,
                line,
                column,
                "Matched instruction category; surrounding content intentionally omitted.",
                "Review the trust boundary, quote or neutralize untrusted instructions, and enforce tool authorization outside the prompt.",
                ("prompt-security", "human-review"),
            ))
        return results


class DependencyPinRule:
    rule_id = "AAISA-DEPENDENCY-001"
    title = "Unpinned Python dependency"
    description = "Finds requirements entries without an exact version or immutable reference."
    default_severity = Severity.MEDIUM
    artifact_kinds = ("dependency-manifest",)
    standards = ("SLSA",)
    references = ("https://slsa.dev/spec/v1.2/",)

    def evaluate(self, artifact: Artifact) -> Iterable[Finding]:
        if artifact.text is None or artifact.path.name.lower() != "requirements.txt":
            return ()
        results: list[Finding] = []
        for number, raw in enumerate(artifact.text.splitlines(), 1):
            line = raw.strip()
            if not line or line.startswith(("#", "-r", "--", "git+")):
                continue
            if "==" not in line or any(token in line for token in (">=", "<=", "~=", "!=", "*")):
                package = re.split(r"[<>=!~\s\[]", line, maxsplit=1)[0]
                results.append(Finding(
                    self.rule_id,
                    self.title,
                    Severity.MEDIUM,
                    f"Dependency {package!r} is not pinned to one exact version.",
                    artifact.relative_path,
                    number,
                    1,
                    f"requirement: {package or '[unparsed]'}",
                    "Pin the dependency and preserve a reviewed lock file or hash set.",
                    ("supply-chain", "dependencies"),
                ))
        return results


class ContainerTagRule:
    rule_id = "AAISA-CONTAINER-001"
    title = "Mutable container image reference"
    description = "Finds latest or untagged images in Dockerfiles and Compose files."
    default_severity = Severity.MEDIUM
    artifact_kinds = ("deployment",)
    standards = ("SLSA",)
    references = ("https://slsa.dev/spec/v1.2/",)
    from_pattern = re.compile(r"(?im)^\s*FROM\s+([^\s]+)")
    compose_pattern = re.compile(r"(?im)^\s*image:\s*([^\s#]+)")

    def evaluate(self, artifact: Artifact) -> Iterable[Finding]:
        if artifact.text is None or artifact.kind != "deployment":
            return ()
        findings: list[Finding] = []
        for pattern in (self.from_pattern, self.compose_pattern):
            for match in pattern.finditer(artifact.text):
                image = match.group(1)
                base = image.split("@", 1)[0]
                last_segment = base.rsplit("/", 1)[-1]
                mutable = "@sha256:" not in image and (":" not in last_segment or last_segment.endswith(":latest"))
                if mutable:
                    line, column = _location(artifact.text, match.start(1))
                    findings.append(Finding(
                        self.rule_id,
                        self.title,
                        Severity.MEDIUM,
                        f"Container image {image!r} is untagged or uses the mutable latest tag.",
                        artifact.relative_path,
                        line,
                        column,
                        f"image: {image}",
                        "Use a reviewed immutable digest and document the update process.",
                        ("supply-chain", "containers"),
                    ))
        return findings


class ProvenanceRule:
    rule_id = "AAISA-PROVENANCE-001"
    title = "Incomplete AI artifact provenance"
    description = "Checks model and dataset metadata for source, version, license, and SHA-256 evidence."
    default_severity = Severity.MEDIUM
    artifact_kinds = ("model-metadata", "dataset-metadata")
    standards = ("NIST AI 600-1", "CycloneDX ML-BOM")
    references = (
        "https://doi.org/10.6028/NIST.AI.600-1",
        "https://cyclonedx.org/guides/OWASP_CycloneDX-Authoritative-Guide-to-AI-ML-BOM-en.pdf",
    )
    required = ("source", "version", "license", "sha256")

    def evaluate(self, artifact: Artifact) -> Iterable[Finding]:
        if artifact.text is None or artifact.kind not in {"model-metadata", "dataset-metadata"}:
            return ()
        lowered = artifact.text.lower()
        missing = [field for field in self.required if not re.search(rf"(?m)^\s*[\"']?{field}[\"']?\s*[:=]", lowered)]
        if not missing:
            return ()
        return (Finding(
            self.rule_id,
            self.title,
            Severity.MEDIUM,
            "AI artifact metadata is missing provenance fields: " + ", ".join(missing) + ".",
            artifact.relative_path,
            1,
            1,
            "missing fields: " + ", ".join(missing),
            "Record the authorized source, immutable version, license, and verified SHA-256 digest.",
            ("provenance", artifact.kind),
        ),)


from .structured_rules import AgentToolPermissionRule, CycloneDxMLBomRule, SlsaProvenanceRule


BUILTIN_RULES: tuple[Rule, ...] = (
    SecretRule(),
    PromptInjectionRule(),
    DependencyPinRule(),
    ContainerTagRule(),
    ProvenanceRule(),
    AgentToolPermissionRule(),
    SlsaProvenanceRule(),
    CycloneDxMLBomRule(),
)


def select_rules(enabled: tuple[str, ...] | None) -> tuple[Rule, ...]:
    if enabled is None:
        return BUILTIN_RULES
    known = {rule.rule_id: rule for rule in BUILTIN_RULES}
    unknown = sorted(set(enabled) - known.keys())
    if unknown:
        raise ValueError("unknown enabled rule(s): " + ", ".join(unknown))
    return tuple(known[rule_id] for rule_id in enabled)