"""CORTX-Secure — the security architecture for self-editing AI.

Threat model: the model itself, or a user prompting the model, attempts to
escalate self-editing beyond policy text into weights, code, identity, or
credentials. This module makes that physically impossible in code.

Defense layers:
  1. WHITELIST: only sections in EDITABLE_SECTIONS can ever be edited.
  2. LOCKED: security_rules, identity, weights, code, admin_token are rejected
     at the validator (enforced in code, not prompt text).
  3. SIZE CAPS: every edit bounded (chars per section, edits per turn).
  4. RATE LIMIT: max edits per rolling hour; violations logged and blocked.
  5. HASH CHAIN: every edit (and every rejected attempt) is appended to a
     tamper-evident audit log: entry.hash = sha256(prev.hash + entry.json).
     Admin can verify the chain; tampering is detectable.
  6. APPROVAL MODE: admin can require human approval before edits go live.
  7. ROLLBACK: version-pinned config; admin can restore any prior version.
  8. ISOLATION: the edit channel reaches ONLY policy text files on disk. It
     has no filesystem path capability, no exec, no network beyond the
     config repo. Weights/code/token are never in the same trust zone.
"""
import hashlib, json, time, threading, re

EDITABLE_SECTIONS = ("constitution", "tool_rules", "metacognitive_rules",
                     "refusal_rules", "playbook")
LOCKED_SECTIONS = ("security_rules", "identity", "weights", "code",
                   "admin_token", "system_prompt")
MAX_SECTION_CHARS = 8000
MAX_EDITS_PER_TURN = 8
MAX_EDITS_PER_HOUR = 40

EDIT_BLOCK_RE = re.compile(r"\[cortex_edit\]\s*(\{.*?\})\s*\[/cortex_edit\]", re.DOTALL)

class AuditLog:
    """Tamper-evident hash-chained audit log."""
    def __init__(self):
        self.entries = []
        self.lock = threading.Lock()
        self._prev_hash = "GENESIS"

    def append(self, kind, section, actor, detail, accepted):
        with self.lock:
            entry = {"ts": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
                     "seq": len(self.entries), "kind": kind, "section": section,
                     "actor": actor, "detail": str(detail)[:200],
                     "accepted": accepted, "prev_hash": self._prev_hash}
            payload = json.dumps(entry, sort_keys=True, ensure_ascii=False)
            entry["hash"] = hashlib.sha256(
                (self._prev_hash + payload).encode("utf-8")).hexdigest()
            self._prev_hash = entry["hash"]
            self.entries.append(entry)
            return entry

    def verify(self):
        prev = "GENESIS"
        for e in self.entries:
            payload = json.dumps({k: v for k, v in e.items() if k != "hash"},
                                 sort_keys=True, ensure_ascii=False)
            expect = hashlib.sha256((prev + payload).encode("utf-8")).hexdigest()
            if expect != e["hash"] or e["prev_hash"] != prev:
                return False, e["seq"]
            prev = e["hash"]
        return True, len(self.entries)

class SecureEditor:
    """Validates and applies self-edits under the security model above."""
    def __init__(self, config, audit=None, approval_required=False):
        self.config = config
        self.audit = audit or AuditLog()
        self.approval_required = approval_required
        self.pending = []
        self.edit_times = []
        self.lock = threading.Lock()

    def _rate_ok(self):
        now = time.time()
        self.edit_times = [t for t in self.edit_times if now - t < 3600]
        return len(self.edit_times) < MAX_EDITS_PER_HOUR

    def validate(self, section, content):
        """Returns (ok, reason). Pure function — no side effects."""
        section = str(section or "").strip().lower()
        content = str(content or "").strip()
        if section in LOCKED_SECTIONS:
            return False, "locked section: %s is outside the editable surface" % section
        if section not in EDITABLE_SECTIONS:
            return False, "unknown section: %s" % section
        if not content or len(content) < 20:
            return False, "content too short"
        if len(content) > MAX_SECTION_CHARS:
            return False, "content exceeds %d chars" % MAX_SECTION_CHARS
        return True, "ok"

    def apply_edits(self, model_text):
        """Parse [cortex_edit] blocks, validate, apply or reject.
        Returns (clean_text, applied, rejected)."""
        applied, rejected = [], []
        if not self.config.get("self_edit_enabled", True):
            return EDIT_BLOCK_RE.sub("", model_text), applied, rejected

        def _grab(m):
            try:
                obj = json.loads(re.sub(r",\s*([}\]])", r"\1", m.group(1)))
            except Exception:
                self.audit.append("edit", "?", "model", "unparseable json", False)
                rejected.append({"section": "?", "reason": "unparseable"})
                return ""
            sec = str(obj.get("section", "")).lower()
            content = str(obj.get("content", ""))
            ok, reason = self.validate(sec, content)
            self.audit.append("edit", sec, "model", reason, ok)
            if not ok:
                rejected.append({"section": sec, "reason": reason})
                return ""
            with self.lock:
                if len(applied) >= MAX_EDITS_PER_TURN:
                    rejected.append({"section": sec, "reason": "turn cap"})
                    return ""
                if not self._rate_ok():
                    rejected.append({"section": sec, "reason": "rate limit"})
                    return ""
                self.config[sec] = content
                self.config["version"] = int(self.config.get("version", 1)) + 1
                self.config["updated_utc"] = time.strftime(
                    "%Y-%m-%dT%H:%M:%SZ", time.gmtime())
                log = self.config.get("_edit_log", []) or []
                log.append({"ts": self.config["updated_utc"], "section": sec,
                            "chars": len(content), "preview": content[:120]})
                self.config["_edit_log"] = log[-50:]
                self.edit_times.append(time.time())
            applied.append({"section": sec, "chars": len(content)})
            return ""

        clean = EDIT_BLOCK_RE.sub(_grab, model_text)
        return clean, applied, rejected

    def rollback(self, version):
        """Admin-only: restore a pinned config version (caller supplies it)."""
        self.audit.append("rollback", "-", "admin", "rollback to v%s" % version, True)
        return True