"""Close WP12 aggregation, correction and allocation integrity gaps."""

from alembic import op

revision = "20260818_549"
down_revision = "20260818_548"
branch_labels = None
depends_on = None


def upgrade() -> None:
    op.get_bind().exec_driver_sql(r'''
    ALTER FUNCTION public.metering_showback_control_write(text,text,jsonb) RENAME TO metering_showback_control_write_legacy;
    CREATE FUNCTION public.metering_showback_control_write(p_action text,p_principal text,p_payload jsonb) RETURNS jsonb LANGUAGE plpgsql SECURITY DEFINER SET search_path=pg_catalog,public AS $guard$
    DECLARE n integer; distinct_n integer; valid_n integer;
    BEGIN
      IF p_action='allocation' THEN
        IF jsonb_typeof(p_payload)<>'object' OR jsonb_typeof(p_payload->'allocations')<>'array' THEN RAISE EXCEPTION 'metering_allocation_closed'; END IF;
        SELECT count(*),count(DISTINCT value->>'target'),count(*) FILTER (WHERE jsonb_typeof(value)='object' AND value-ARRAY['target','weight_micros']='{}'::jsonb AND value->>'target' ~ '^[a-z][a-z0-9-]{0,62}$' AND value->>'weight_micros' ~ '^[1-9][0-9]{0,6}$' AND (value->>'weight_micros')::integer<=1000000) INTO n,distinct_n,valid_n FROM jsonb_array_elements(p_payload->'allocations');
        IF n<>distinct_n THEN RAISE EXCEPTION 'metering_allocation_target_duplicate'; END IF;
        IF n<>valid_n THEN RAISE EXCEPTION 'metering_allocation_weight_invalid'; END IF;
      END IF;
      RETURN public.metering_showback_control_write_legacy(p_action,p_principal,p_payload);
    END; $guard$;
    ALTER FUNCTION public.metering_showback_control_write(text,text,jsonb) OWNER TO dataops_tenant_foundation_owner;
    REVOKE ALL ON FUNCTION public.metering_showback_control_write_legacy(text,text,jsonb),public.metering_showback_control_write(text,text,jsonb) FROM PUBLIC,dataops_app,dataops_app_runtime,dataops_bi_ai_catalog_control;
    GRANT EXECUTE ON FUNCTION public.metering_showback_control_write(text,text,jsonb) TO dataops_bi_ai_catalog_control;

    CREATE FUNCTION public.metering_showback_correction_guard() RETURNS trigger LANGUAGE plpgsql SECURITY DEFINER SET search_path=pg_catalog,public AS $correction$
    DECLARE original record;
    BEGIN
      IF NEW.evidence_reference !~ '^local-fixture://wp12/v1(/[a-z][a-z0-9-]{0,62})?$' OR NEW.evidence_reference ~ '(sql|script|select|insert|update|delete|drop|exec|curl|wget|token|secret|credential|password)' THEN RAISE EXCEPTION 'metering_evidence_reference_invalid'; END IF;
      IF NEW.correction_of IS NULL THEN RETURN NEW; END IF;
      SELECT * INTO original FROM public.metering_events WHERE event_uid=NEW.correction_of FOR KEY SHARE;
      IF NOT FOUND OR original.tenant_ref<>NEW.tenant_ref OR original.domain_ref<>NEW.domain_ref OR original.correction_of IS NOT NULL OR original.event_kind<>NEW.event_kind OR original.unit<>NEW.unit OR original.window_start<>NEW.window_start OR original.window_end<>NEW.window_end OR original.department_ref<>NEW.department_ref OR original.project_ref<>NEW.project_ref OR original.cost_center_ref<>NEW.cost_center_ref THEN RAISE EXCEPTION 'metering_correction_scope_invalid'; END IF;
      RETURN NEW;
    END; $correction$;
    CREATE TRIGGER metering_showback_correction_guard BEFORE INSERT ON public.metering_events FOR EACH ROW EXECUTE FUNCTION public.metering_showback_correction_guard();

    CREATE FUNCTION public.metering_showback_rule_window_guard() RETURNS trigger LANGUAGE plpgsql SECURITY DEFINER SET search_path=pg_catalog,public AS $window$
    BEGIN
      IF EXISTS(SELECT 1 FROM public.metering_allocation_rules r WHERE r.tenant_ref=NEW.tenant_ref AND r.domain_ref=NEW.domain_ref AND r.department_ref=NEW.department_ref AND r.project_ref=NEW.project_ref AND r.cost_center_ref=NEW.cost_center_ref AND (r.rule_uid,r.rule_version)<>(NEW.rule_uid,NEW.rule_version) AND tstzrange(r.effective_start,r.effective_end,'[)') && tstzrange(NEW.effective_start,NEW.effective_end,'[)')) THEN RAISE EXCEPTION 'metering_allocation_window_overlap'; END IF;
      RETURN NEW;
    END; $window$;
    CREATE TRIGGER metering_showback_rule_window_guard BEFORE INSERT OR UPDATE ON public.metering_allocation_rules FOR EACH ROW EXECUTE FUNCTION public.metering_showback_rule_window_guard();

    CREATE FUNCTION public.metering_showback_rollup(p_payload jsonb) RETURNS jsonb LANGUAGE plpgsql SECURITY DEFINER SET search_path=pg_catalog,public AS $rollup$
    DECLARE c record; source_micros bigint; allocated_micros bigint; details jsonb;
    BEGIN
      IF NOT pg_has_role(session_user,'dataops_app_runtime','MEMBER') OR jsonb_typeof(p_payload)<>'object' OR p_payload-ARRAY['request_claim','window']<>'{}'::jsonb OR p_payload->>'request_claim' !~ '^[0-9a-f-]{36}$' OR p_payload->>'window' !~ '^[0-9]{4}-(0[1-9]|1[0-2])$' THEN RAISE EXCEPTION 'metering_read_denied'; END IF;
      SELECT * INTO c FROM public.metering_runtime_claims WHERE claim_uid=(p_payload->>'request_claim')::uuid AND action_name='read' AND consumed_at IS NULL AND expires_at>clock_timestamp() FOR UPDATE;
      IF NOT FOUND THEN RAISE EXCEPTION 'metering_claim_denied'; END IF;
      WITH scoped AS (SELECT e.*,r.rule_uid,r.rule_version,r.rule_digest FROM public.metering_events e LEFT JOIN LATERAL (SELECT x.rule_uid,x.rule_version,x.rule_digest FROM public.metering_allocation_rules x WHERE x.tenant_ref=e.tenant_ref AND x.domain_ref=e.domain_ref AND x.department_ref=e.department_ref AND x.project_ref=e.project_ref AND x.cost_center_ref=e.cost_center_ref AND x.effective_start<=e.window_start AND x.effective_end>e.window_start ORDER BY x.rule_version DESC,x.rule_uid DESC LIMIT 1) r ON true WHERE e.tenant_ref=c.tenant_ref AND e.domain_ref=c.domain_ref AND to_char(e.window_start AT TIME ZONE 'UTC','YYYY-MM')=p_payload->>'window'), rules AS (SELECT rule_uid,rule_version,rule_digest,department_ref,project_ref,cost_center_ref,sum(quantity_micros)::bigint AS source FROM scoped WHERE rule_uid IS NOT NULL GROUP BY rule_uid,rule_version,rule_digest,department_ref,project_ref,cost_center_ref), weights AS (SELECT rs.*,a.target_ref,(rs.source*a.weight_micros/1000000)::bigint AS base FROM rules rs JOIN public.metering_allocations a ON a.rule_uid=rs.rule_uid AND a.rule_version=rs.rule_version), expanded AS (SELECT w.*,w.base+CASE WHEN w.target_ref=min(w.target_ref) OVER (PARTITION BY w.rule_uid,w.rule_version) THEN w.source-sum(w.base) OVER (PARTITION BY w.rule_uid,w.rule_version) ELSE 0 END AS amount FROM weights w), per_rule AS (SELECT rule_uid,rule_version,rule_digest,department_ref,project_ref,cost_center_ref,max(source) AS source,sum(amount)::bigint AS allocated,jsonb_agg(jsonb_build_object('target',target_ref,'quantity_micros',amount) ORDER BY target_ref) AS allocations FROM expanded GROUP BY rule_uid,rule_version,rule_digest,department_ref,project_ref,cost_center_ref) SELECT COALESCE((SELECT sum(quantity_micros) FROM scoped),0),COALESCE(sum(allocated),0),COALESCE(jsonb_agg(jsonb_build_object('rule_uid',rule_uid,'rule_version',rule_version,'rule_digest',rule_digest,'mapping',jsonb_build_object('department',department_ref,'business_domain',c.domain_ref,'project',project_ref,'cost_center',cost_center_ref),'source_micros',source,'allocated_micros',allocated,'variance_micros',source-allocated,'allocations',allocations) ORDER BY rule_uid,rule_version),'[]'::jsonb) INTO source_micros,allocated_micros,details FROM per_rule;
      UPDATE public.metering_runtime_claims SET consumed_at=clock_timestamp() WHERE claim_uid=c.claim_uid;
      INSERT INTO public.metering_audit_events(tenant_ref,domain_ref,event_type,actor_ref,payload_digest) VALUES(c.tenant_ref,c.domain_ref,'showback.read',c.principal_ref,encode(sha256(convert_to('rollup|'||p_payload::text,'utf8')),'hex'));
      RETURN jsonb_build_object('window',p_payload->>'window','source_micros',source_micros,'allocated_micros',allocated_micros,'unallocated_micros',source_micros-allocated_micros,'variance_micros',source_micros-allocated_micros,'allocation_rules',details,'mode','ENGINEERING_EVIDENCE_ONLY','chargeback_enabled',false);
    END; $rollup$;
    ALTER FUNCTION public.metering_showback_rollup(jsonb) OWNER TO dataops_tenant_foundation_owner;
    REVOKE ALL ON FUNCTION public.metering_showback_rollup(jsonb),public.metering_showback_correction_guard(),public.metering_showback_rule_window_guard() FROM PUBLIC,dataops_app,dataops_bi_ai_catalog_control;
    GRANT EXECUTE ON FUNCTION public.metering_showback_rollup(jsonb) TO dataops_app_runtime;
    ALTER FUNCTION public.metering_showback_runtime_read(text,jsonb) RENAME TO metering_showback_runtime_read_legacy;
    CREATE FUNCTION public.metering_showback_runtime_read(p_action text,p_payload jsonb) RETURNS jsonb LANGUAGE plpgsql SECURITY DEFINER SET search_path=pg_catalog,public AS $read_guard$
    BEGIN
      IF p_action='chargeback' THEN RAISE EXCEPTION 'chargeback_disabled'; END IF;
      IF p_action IN ('showback','reconciliation') THEN RAISE EXCEPTION 'metering_rollup_required'; END IF;
      IF p_action<>'audit' THEN RAISE EXCEPTION 'metering_read_denied'; END IF;
      RETURN public.metering_showback_runtime_read_legacy(p_action,p_payload);
    END; $read_guard$;
    ALTER FUNCTION public.metering_showback_runtime_read(text,jsonb) OWNER TO dataops_tenant_foundation_owner;
    REVOKE ALL ON FUNCTION public.metering_showback_runtime_read_legacy(text,jsonb),public.metering_showback_runtime_read(text,jsonb) FROM PUBLIC,dataops_app,dataops_bi_ai_catalog_control;
    GRANT EXECUTE ON FUNCTION public.metering_showback_runtime_read(text,jsonb) TO dataops_app_runtime;
    ''')


def downgrade() -> None:
    op.get_bind().exec_driver_sql("""
    DO $$ BEGIN
      IF to_regprocedure('public.metering_showback_runtime_read_legacy(text,jsonb)') IS NOT NULL THEN
        DROP FUNCTION public.metering_showback_runtime_read(text,jsonb);
        ALTER FUNCTION public.metering_showback_runtime_read_legacy(text,jsonb) RENAME TO metering_showback_runtime_read;
        GRANT EXECUTE ON FUNCTION public.metering_showback_runtime_read(text,jsonb) TO dataops_app_runtime;
      END IF;
      IF to_regprocedure('public.metering_showback_control_write_legacy(text,text,jsonb)') IS NOT NULL THEN
        DROP FUNCTION public.metering_showback_control_write(text,text,jsonb);
        ALTER FUNCTION public.metering_showback_control_write_legacy(text,text,jsonb) RENAME TO metering_showback_control_write;
        GRANT EXECUTE ON FUNCTION public.metering_showback_control_write(text,text,jsonb) TO dataops_bi_ai_catalog_control;
      END IF;
    END $$;
    DROP FUNCTION public.metering_showback_rollup(jsonb);
    DROP TRIGGER metering_showback_rule_window_guard ON public.metering_allocation_rules;
    DROP FUNCTION public.metering_showback_rule_window_guard();
    DROP TRIGGER metering_showback_correction_guard ON public.metering_events;
    DROP FUNCTION public.metering_showback_correction_guard();
    """)