"""Bind WP09 control mutations to an authenticated, one-shot application claim."""

from alembic import op

revision = "20260813_500"
down_revision = "20260813_499"
branch_labels = None
depends_on = None


def upgrade() -> None:
    op.execute(r"""
    CREATE TABLE public.agent_runtime_control_claims (
      claim_uid uuid PRIMARY KEY,
      credential_uid uuid NOT NULL REFERENCES public.agent_credentials(uid),
      credential_token_digest char(64) NOT NULL,
      tenant_id text NOT NULL,
      agent_uid uuid NOT NULL,
      principal_id uuid NOT NULL,
      business_domain_uid uuid NOT NULL,
      environment text NOT NULL,
      operation text NOT NULL CHECK (operation IN ('recover','canary_promote','canary_rollback')),
      approval_task_uid uuid NOT NULL REFERENCES public.governance_tasks(uid),
      payload_digest char(64) NOT NULL,
      control_payload jsonb NOT NULL,
      nonce uuid NOT NULL UNIQUE,
      issued_at timestamptz NOT NULL DEFAULT clock_timestamp(),
      expires_at timestamptz NOT NULL,
      consumed_at timestamptz,
      status text NOT NULL CHECK (status IN ('issued','consumed')) DEFAULT 'issued'
    );
    ALTER TABLE public.agent_runtime_control_claims OWNER TO dataops_agent_runtime_owner;
    REVOKE ALL ON TABLE public.agent_runtime_control_claims FROM PUBLIC,dataops_agent_runtime,dataops_app_runtime;
    GRANT SELECT,INSERT,UPDATE ON TABLE public.agent_runtime_control_claims TO dataops_agent_runtime_owner;
    GRANT SELECT,UPDATE ON TABLE public.agent_credentials TO dataops_agent_runtime_owner;

    CREATE FUNCTION public.agent_runtime_control_issue_claim(p_claim jsonb)
    RETURNS uuid LANGUAGE plpgsql SECURITY DEFINER SET search_path=pg_catalog,public AS $issue$
    DECLARE c record; p jsonb; v_digest text; v_expiry timestamptz;
    BEGIN
      IF jsonb_typeof(p_claim)<>'object' OR NOT (p_claim ?& ARRAY['claim_uid','credential_uid','credential_token_digest','nonce','control_payload']) OR p_claim-ARRAY['claim_uid','credential_uid','credential_token_digest','nonce','control_payload']<>'{}'::jsonb THEN RAISE EXCEPTION 'control claim payload is closed'; END IF;
      IF (p_claim->>'claim_uid') !~ '^[0-9a-f-]{36}$' OR (p_claim->>'credential_uid') !~ '^[0-9a-f-]{36}$' OR (p_claim->>'nonce') !~ '^[0-9a-f-]{36}$' OR (p_claim->>'credential_token_digest') !~ '^[0-9a-f]{64}$' OR jsonb_typeof(p_claim->'control_payload')<>'object' THEN RAISE EXCEPTION 'control claim value invalid'; END IF;
      p:=p_claim->'control_payload';
      IF NOT (p ?& ARRAY['tenant_id','agent_uid','principal_id','business_domain_uid','environment','operation','approval_task_uid','idempotency_key','request_digest','expected_fence','incident_uid','from_state','to_state','generation','from_generation','to_generation','dataset_version','dataset_digest','metrics_digest','threshold_policy','issued_at']) OR p-ARRAY['tenant_id','agent_uid','principal_id','business_domain_uid','environment','operation','approval_task_uid','idempotency_key','request_digest','expected_fence','incident_uid','from_state','to_state','generation','from_generation','to_generation','dataset_version','dataset_digest','metrics_digest','threshold_policy','issued_at']<>'{}'::jsonb THEN RAISE EXCEPTION 'control claim context is closed'; END IF;
      IF (p->>'agent_uid') !~ '^[0-9a-f-]{36}$' OR (p->>'principal_id') !~ '^[0-9a-f-]{36}$' OR (p->>'business_domain_uid') !~ '^[0-9a-f-]{36}$' OR (p->>'approval_task_uid') !~ '^[0-9a-f-]{36}$' OR p->>'operation' NOT IN ('recover','canary_promote','canary_rollback') THEN RAISE EXCEPTION 'control claim context invalid'; END IF;
      PERFORM pg_advisory_xact_lock(hashtext((p->>'tenant_id') || ':' || (p->>'idempotency_key')));
      SELECT * INTO c FROM public.agent_credentials WHERE uid=(p_claim->>'credential_uid')::uuid FOR SHARE;
      IF NOT FOUND OR c.status<>'active' OR c.expires_at<=clock_timestamp() OR c.agent_uid<>(p->>'agent_uid')::uuid OR c.issued_by<>(p->>'principal_id')::uuid OR c.token_digest<>p_claim->>'credential_token_digest' THEN RAISE EXCEPTION 'control credential denied'; END IF;
      v_digest:=encode(sha256(convert_to(p::text,'utf8')),'hex');
      v_expiry:=LEAST(c.expires_at,clock_timestamp()+interval '2 minutes');
      INSERT INTO public.agent_runtime_control_claims(claim_uid,credential_uid,credential_token_digest,tenant_id,agent_uid,principal_id,business_domain_uid,environment,operation,approval_task_uid,payload_digest,control_payload,nonce,expires_at) VALUES((p_claim->>'claim_uid')::uuid,c.uid,c.token_digest,p->>'tenant_id',(p->>'agent_uid')::uuid,(p->>'principal_id')::uuid,(p->>'business_domain_uid')::uuid,p->>'environment',p->>'operation',(p->>'approval_task_uid')::uuid,v_digest,p,(p_claim->>'nonce')::uuid,v_expiry);
      RETURN (p_claim->>'claim_uid')::uuid;
    END; $issue$;
    ALTER FUNCTION public.agent_runtime_control_issue_claim(jsonb) OWNER TO dataops_agent_runtime_owner;
    REVOKE ALL ON FUNCTION public.agent_runtime_control_issue_claim(jsonb) FROM PUBLIC,dataops_agent_runtime;
    GRANT EXECUTE ON FUNCTION public.agent_runtime_control_issue_claim(jsonb) TO dataops_app_runtime,dataops_agent_runtime_owner;

    CREATE FUNCTION public.agent_runtime_control_claimed(p_payload jsonb,p_claim_uid uuid)
    RETURNS jsonb LANGUAGE plpgsql SECURITY DEFINER SET search_path=pg_catalog,public AS $claimed$
    DECLARE q record; c record; v_digest text; v_result jsonb;
    BEGIN
      IF jsonb_typeof(p_payload)<>'object' THEN RAISE EXCEPTION 'control claim payload is closed'; END IF;
      PERFORM pg_advisory_xact_lock(hashtext((p_payload->>'tenant_id') || ':' || (p_payload->>'idempotency_key')));
      v_digest:=encode(sha256(convert_to(p_payload::text,'utf8')),'hex');
      SELECT * INTO q FROM public.agent_runtime_control_claims WHERE claim_uid=p_claim_uid FOR UPDATE;
      IF NOT FOUND OR q.status<>'issued' OR q.expires_at<=clock_timestamp() OR q.payload_digest<>v_digest OR q.control_payload<>p_payload OR q.tenant_id<>p_payload->>'tenant_id' OR q.agent_uid<>(p_payload->>'agent_uid')::uuid OR q.principal_id<>(p_payload->>'principal_id')::uuid OR q.business_domain_uid<>(p_payload->>'business_domain_uid')::uuid OR q.environment<>p_payload->>'environment' OR q.operation<>p_payload->>'operation' OR q.approval_task_uid<>(p_payload->>'approval_task_uid')::uuid THEN RAISE EXCEPTION 'control claim denied'; END IF;
      SELECT * INTO c FROM public.agent_credentials WHERE uid=q.credential_uid FOR SHARE;
      IF NOT FOUND OR c.status<>'active' OR c.expires_at<=clock_timestamp() OR c.agent_uid<>q.agent_uid OR c.issued_by<>q.principal_id OR c.token_digest<>q.credential_token_digest THEN RAISE EXCEPTION 'control credential denied'; END IF;
      v_result:=public.agent_runtime_control(p_payload);
      UPDATE public.agent_runtime_control_claims SET status='consumed',consumed_at=clock_timestamp() WHERE claim_uid=q.claim_uid AND status='issued';
      IF NOT FOUND THEN RAISE EXCEPTION 'control claim already consumed'; END IF;
      RETURN v_result;
    END; $claimed$;
    ALTER FUNCTION public.agent_runtime_control_claimed(jsonb,uuid) OWNER TO dataops_agent_runtime_owner;
    REVOKE ALL ON FUNCTION public.agent_runtime_control(jsonb) FROM PUBLIC,dataops_agent_runtime,dataops_app_runtime;
    REVOKE ALL ON FUNCTION public.agent_runtime_control_claimed(jsonb,uuid) FROM PUBLIC,dataops_agent_runtime;
    GRANT EXECUTE ON FUNCTION public.agent_runtime_control_claimed(jsonb,uuid) TO dataops_app_runtime,dataops_agent_runtime_owner;
    """)


def downgrade() -> None:
    op.execute("""
    REVOKE ALL ON FUNCTION public.agent_runtime_control_claimed(jsonb,uuid) FROM PUBLIC,dataops_app_runtime;
    DROP FUNCTION IF EXISTS public.agent_runtime_control_claimed(jsonb,uuid);
    REVOKE ALL ON FUNCTION public.agent_runtime_control_issue_claim(jsonb) FROM PUBLIC,dataops_app_runtime;
    DROP FUNCTION IF EXISTS public.agent_runtime_control_issue_claim(jsonb);
    DROP TABLE IF EXISTS public.agent_runtime_control_claims;
    REVOKE ALL ON FUNCTION public.agent_runtime_control(jsonb) FROM PUBLIC;
    GRANT EXECUTE ON FUNCTION public.agent_runtime_control(jsonb) TO dataops_agent_runtime,dataops_agent_runtime_owner;
    """)