"""Fail-closed signed deterministic execution amendment for R7."""
from datetime import datetime, timezone
import hashlib, json, sqlite3
from pathlib import Path
from typing import Mapping
from aios_task_terminal_recovery import DeterministicPlan, RecoveryError, canonical_json

AMENDMENT_SCHEMA="aios.deterministic-execution-amendment.v1"
APPLY_SCHEMA="aios.deterministic-execution-amendment-apply.v1"
NAMESPACE="aios-owner-approval-v1"
EVENT="TASK_RECOVERY_EXECUTION_AMENDMENT_APPLIED"
APPLY_STATES=frozenset({"FAILED", "QUEUED"})


def _err(code): raise RecoveryError(code)
def _sha(v): return hashlib.sha256(canonical_json(v).encode()).hexdigest()
def amendment_body(x):
    return {k:v for k,v in x.items() if k not in {"signature","amendment_id"}}
def amendment_id(x): return "sha256:"+_sha(amendment_body(x))
def signed_payload(x):
    return canonical_json({"amendment_id":x["amendment_id"],**amendment_body(x)}).encode()


def parse_amendment(raw):
    x=dict(raw)
    required={"schema","amendment_id","nonce","task_id","job_id","generation",
      "repository","branch","expected_branch_oid","recovery_epoch","decision_id",
      "plan_id","plan_sha256","expected_apply_status","expected_apply_revision",
      "expected_apply_attempt","expected_current_capabilities","envelope",
      "owner_identity","key_id","issued_at","not_before","expires_at",
      "one_time_use","signature"}
    if set(x)<required or x.get("schema")!=AMENDMENT_SCHEMA: _err("AMENDMENT_SCHEMA_INVALID")
    if x["expected_apply_status"] not in APPLY_STATES: _err("AMENDMENT_APPLY_STATUS_INVALID")
    if x.get("one_time_use") is not True: _err("AMENDMENT_ONE_TIME_REQUIRED")
    if not isinstance(x["envelope"],Mapping): _err("AMENDMENT_ENVELOPE_INVALID")
    if not isinstance(x["expected_current_capabilities"],list): _err("AMENDMENT_CAPABILITIES_INVALID")
    for k in ("generation","expected_apply_revision","expected_apply_attempt"):
        if isinstance(x[k],bool) or not isinstance(x[k],int) or x[k]<0: _err("AMENDMENT_CAS_INVALID")
    for k in ("issued_at","not_before","expires_at"):
        try:d=datetime.fromisoformat(str(x[k]).replace("Z","+00:00"))
        except ValueError as e: raise RecoveryError("AMENDMENT_TIME_INVALID") from e
        if d.tzinfo is None:_err("AMENDMENT_TIME_INVALID")
        x[k]=d.astimezone(timezone.utc).isoformat().replace("+00:00","Z")
    x["expected_current_capabilities"]=sorted(set(x["expected_current_capabilities"]))
    x["envelope"]=dict(x["envelope"])
    return x


def _validate_envelope(env,plan,x):
    try:
        from aios_online_agent_execution import GovernedExecutionBinding,validate_execution_envelope
        binding=GovernedExecutionBinding.from_mapping(env)
        e=dict(validate_execution_envelope(env,binding=binding,worker_capabilities=plan.required_capabilities))
    except Exception as z: raise RecoveryError("AMENDMENT_EXECUTION_ENVELOPE_INVALID") from z
    if e.get("action")!="deterministic" or e.get("adapter")!="deterministic":_err("AMENDMENT_EXECUTION_ACTION_INVALID")
    if tuple(sorted(e.get("capability_scope") or ()))!=plan.required_capabilities:_err("AMENDMENT_EXECUTION_CAPABILITIES_MISMATCH")
    if not e.get("change_set") or not e.get("verification"):_err("AMENDMENT_EXECUTION_OPERATIONS_REQUIRED")
    if e.get("runtime_input") not in (None,{}) or e.get("source_signals") not in (None,[]):_err("AMENDMENT_STALE_RUNTIME_FIELDS_FORBIDDEN")
    for k,w in (("task_id",x["task_id"]),("job_id",x["job_id"]),("repository",x["repository"]),
      ("generation",x["generation"]),("target_branch",x["branch"]),
      ("expected_branch_oid",x["expected_branch_oid"]),("base_oid",x["expected_branch_oid"]),
      ("plan_id",x["plan_id"]),("plan_sha256",x["plan_sha256"]),("recovery_epoch",x["recovery_epoch"])):
        if e.get(k)!=w:_err("AMENDMENT_EXECUTION_BINDING_MISMATCH")
    g=e.get("governance") or {}
    if g.get("writer_lease") is not True:_err("AMENDMENT_WRITER_LEASE_REQUIRED")
    if any(g.get(k,False) for k in ("main_write","production_write","tag_write","baseline_write","force_push","protected_ref_bypass")):_err("AMENDMENT_PROTECTED_WRITE_FORBIDDEN")
    return e


class DeterministicExecutionAmendmentController:
    def __init__(self,db_path,verifier,*,clock=None,envelope_validator=None):
        self.db_path=Path(db_path); self.verifier=verifier
        self.clock=clock or (lambda:datetime.now(timezone.utc)); self._init()

    def _db(self):
        c=sqlite3.connect(self.db_path,timeout=10,isolation_level=None); c.row_factory=sqlite3.Row
        c.execute("PRAGMA busy_timeout=10000"); return c

    def _now(self): return self.clock().astimezone(timezone.utc).isoformat().replace("+00:00","Z")

    def _init(self):
        with self._db() as c:
            c.execute("DROP TRIGGER IF EXISTS aios_r7_recovery_capability_projection")
            c.executescript("""
CREATE TABLE IF NOT EXISTS deterministic_execution_amendments(
 amendment_id TEXT PRIMARY KEY,nonce TEXT UNIQUE NOT NULL,task_id TEXT NOT NULL,
 decision_id TEXT NOT NULL,plan_id TEXT NOT NULL,plan_sha256 TEXT NOT NULL,
 envelope_digest TEXT NOT NULL,envelope_json TEXT NOT NULL,payload_json TEXT NOT NULL,
 created_at TEXT NOT NULL,UNIQUE(task_id,decision_id));
DROP TRIGGER IF EXISTS aios_r7_recovery_requires_executable_amendment;
CREATE TRIGGER aios_r7_recovery_requires_executable_amendment
BEFORE UPDATE OF status ON tasks
WHEN OLD.status='FAILED' AND NEW.status='QUEUED' BEGIN
 SELECT CASE WHEN NOT EXISTS(
  SELECT 1 FROM deterministic_execution_amendments a
  WHERE a.task_id=OLD.task_id
   AND a.decision_id=json_extract(NEW.metadata_json,'$.recovery.decision_id')
   AND a.amendment_id=json_extract(OLD.metadata_json,'$.recovery.execution_amendment_id')
   AND a.envelope_digest=json_extract(OLD.metadata_json,'$.execution.execution_digest')
 ) THEN RAISE(ABORT,'SIGNED_EXECUTION_AMENDMENT_REQUIRED') END;
END;""")

    def _plan(self,c,x):
        r=c.execute("SELECT * FROM deterministic_plans WHERE plan_id=?",(x["plan_id"],)).fetchone()
        if r is None or r["plan_sha256"]!=x["plan_sha256"]:_err("AMENDMENT_PLAN_NOT_FOUND_OR_MISMATCH")
        p=DeterministicPlan.from_dict(json.loads(r["payload_json"]))
        if (p.task_id,p.job_id,p.generation,p.repository,p.branch,p.expected_branch_oid,p.recovery_epoch)!=(x["task_id"],x["job_id"],x["generation"],x["repository"],x["branch"],x["expected_branch_oid"],x["recovery_epoch"]):_err("AMENDMENT_PLAN_BINDING_MISMATCH")
        return p

    def persist(self,raw):
        x=parse_amendment(raw)
        if x["amendment_id"]!=amendment_id(x):_err("AMENDMENT_ID_DIGEST_MISMATCH")
        now=self._now()
        if now<x["not_before"] or now>=x["expires_at"]:_err("AMENDMENT_EXPIRED_OR_NOT_YET_VALID")
        if not self.verifier.verify(namespace=NAMESPACE,payload=signed_payload(x),signature=x["signature"],key_id=x["key_id"]):_err("AMENDMENT_SIGNATURE_INVALID")
        with self._db() as c:
            c.execute("BEGIN IMMEDIATE"); e=_validate_envelope(x["envelope"],self._plan(c,x),x)
            ej=canonical_json(e); ed=e["execution_digest"]; stored=dict(x); stored["signature"]="<present>"
            old=c.execute("SELECT * FROM deterministic_execution_amendments WHERE amendment_id=?",(x["amendment_id"],)).fetchone()
            if old:
                if old["envelope_json"]!=ej:_err("AMENDMENT_ID_REUSE_MISMATCH")
                c.commit(); return {"status":"PRESENT","amendment_id":x["amendment_id"],"task_id":x["task_id"],"decision_id":x["decision_id"],"plan_id":x["plan_id"],"execution_digest":ed}
            if c.execute("SELECT 1 FROM deterministic_execution_amendments WHERE nonce=? OR (task_id=? AND decision_id=?)",(x["nonce"],x["task_id"],x["decision_id"])).fetchone():_err("AMENDMENT_REPLAY_OR_DECISION_REUSE")
            c.execute("INSERT INTO deterministic_execution_amendments VALUES(?,?,?,?,?,?,?,?,?,?)",(x["amendment_id"],x["nonce"],x["task_id"],x["decision_id"],x["plan_id"],x["plan_sha256"],ed,ej,canonical_json(stored),now)); c.commit()
        return {"status":"PERSISTED","amendment_id":x["amendment_id"],"task_id":x["task_id"],"decision_id":x["decision_id"],"plan_id":x["plan_id"],"execution_digest":ed}

    def apply(self,raw):
        if raw.get("schema")!=APPLY_SCHEMA:_err("AMENDMENT_APPLY_SCHEMA_INVALID")
        aid=raw.get("amendment_id"); tid=raw.get("task_id")
        with self._db() as c:
            c.execute("BEGIN IMMEDIATE"); a=c.execute("SELECT * FROM deterministic_execution_amendments WHERE amendment_id=? AND task_id=?",(aid,tid)).fetchone()
            if a is None:_err("AMENDMENT_NOT_FOUND")
            x=json.loads(a["payload_json"]); d=json.loads(a["envelope_json"])
            t=c.execute("SELECT * FROM tasks WHERE task_id=?",(tid,)).fetchone()
            if t is None:_err("TASK_NOT_FOUND")
            m=json.loads(t["metadata_json"] or "{}"); cp=tuple(sorted(json.loads(t["required_capabilities_json"] or "[]"))); dc=tuple(sorted(d.get("capability_scope") or ()))
            rec=m.get("recovery"); rec=dict(rec) if isinstance(rec,Mapping) else {}
            old=m.get("execution")
            if isinstance(old,Mapping) and old.get("execution_digest")==a["envelope_digest"] and cp==dc and rec.get("execution_amendment_id")==aid:
                c.commit(); return {"status":"PRESENT","task_id":tid,"amendment_id":aid,"execution_digest":a["envelope_digest"],"revision":int(t["revision"]),"attempt":int(t["attempt"])}
            if t["status"]!=x["expected_apply_status"] or int(t["revision"])!=x["expected_apply_revision"] or int(t["attempt"])!=x["expected_apply_attempt"] or cp!=tuple(x["expected_current_capabilities"]) or t["claimed_by"] is not None or t["lease_token_hash"] is not None:_err("AMENDMENT_APPLY_TASK_CAS_MISMATCH")
            if t["status"]=="QUEUED" and rec.get("decision_id")!=x["decision_id"]:_err("AMENDMENT_DECISION_BINDING_MISMATCH")
            m["execution"]=d; m["execution_projection"]={"schema":"aios.online-agent.task-projection.v1","task_id":tid,"job_id":d.get("job_id"),"execution_digest":a["envelope_digest"]}
            rec["execution_amendment_id"]=aid; m["recovery"]=rec; now=self._now()
            q=c.execute("""UPDATE tasks SET required_capabilities_json=?,metadata_json=?,revision=revision+1,updated_at=? WHERE task_id=? AND status=? AND revision=? AND attempt=? AND claimed_by IS NULL AND lease_token_hash IS NULL""",(canonical_json(list(dc)),canonical_json(m),now,tid,x["expected_apply_status"],x["expected_apply_revision"],x["expected_apply_attempt"]))
            if q.rowcount!=1:_err("AMENDMENT_APPLY_TASK_CAS_FAILED")
            c.execute("INSERT INTO events(task_id,agent_id,event_type,payload_json,created_at) VALUES(?,?,?,?,?)",(tid,None,EVENT,canonical_json({"amendment_id":aid,"decision_id":x["decision_id"],"plan_id":x["plan_id"],"execution_digest":a["envelope_digest"],"required_capabilities":list(dc)}),now)); c.commit()
        return {"status":"APPLIED","task_id":tid,"amendment_id":aid,"decision_id":x["decision_id"],"plan_id":x["plan_id"],"execution_digest":a["envelope_digest"],"required_capabilities":list(dc),"revision":x["expected_apply_revision"]+1,"attempt":x["expected_apply_attempt"]}