-- llm-proxy: reserve LLM quota before calling the provider.
--
-- llm-proxy used to run a read-only checkQuota (a SELECT on daily_usage),
-- call Anthropic, and only then charge with consume_quota. Nothing counted
-- requests that were still in flight, so N parallel requests from one user all
-- passed the pre-check and each paid for a full Anthropic generation before
-- consume_quota rejected every request past the allowance and threw the paid
-- reply away (the same flaw 20260928000037 fixed for meeting documents).
--
-- reserve_llm_quota now takes one unit up front, under the same
-- per-user/per-feature advisory lock that consume_quota uses, and records it
-- in daily_usage immediately, so a concurrent request (or a consume_quota
-- call for the same feature) sees it. The edge function calls the provider
-- only while it holds a reservation, then settles it:
--   * finalize_llm_quota(id, true)  -> 'completed' (the unit stays spent)
--   * finalize_llm_quota(id, false) -> 'released'  (daily_usage -1, and the
--     overage credit is returned when the unit came from overage)
-- A crashed worker cannot hold a unit forever: reservations whose 10 minute
-- lease has expired are released by the next reserve call for that
-- user/feature, well beyond the proxy's 45 s time-to-first-byte and 180 s
-- stream deadlines.
--
-- p_base_limit / p_period come from PLAN_QUOTA (packages/core plan catalog),
-- exactly like consume_quota. Only service_role may call these functions.
-- Apply this migration before deploying the llm-proxy that calls it.

CREATE TABLE IF NOT EXISTS public.llm_quota_reservations (
  id uuid PRIMARY KEY,
  user_id uuid NOT NULL REFERENCES auth.users(id) ON DELETE CASCADE,
  feature text NOT NULL CHECK (feature IN ('llm_haiku', 'llm_sonnet', 'llm_opus')),
  usage_date date NOT NULL DEFAULT CURRENT_DATE,
  consumed_from text NOT NULL CHECK (consumed_from IN ('base', 'overage', 'unlimited')),
  status text NOT NULL DEFAULT 'reserved' CHECK (status IN ('reserved', 'completed', 'released')),
  tier text NOT NULL,
  quota_period text NOT NULL CHECK (quota_period IN ('daily', 'weekly')),
  quota_limit integer NOT NULL,
  current_count integer NOT NULL,
  overage_after integer NOT NULL,
  lease_expires_at timestamptz NOT NULL,
  created_at timestamptz NOT NULL DEFAULT now(),
  finalized_at timestamptz,
  release_reason text
);

ALTER TABLE public.llm_quota_reservations ENABLE ROW LEVEL SECURITY;
REVOKE ALL ON TABLE public.llm_quota_reservations FROM PUBLIC, anon, authenticated;
GRANT SELECT, INSERT, UPDATE, DELETE ON TABLE public.llm_quota_reservations TO service_role;

CREATE INDEX IF NOT EXISTS idx_llm_quota_reservations_reclaim
  ON public.llm_quota_reservations(user_id, feature, lease_expires_at)
  WHERE status = 'reserved';

CREATE OR REPLACE FUNCTION public.reserve_llm_quota(
  p_user_id uuid,
  p_reservation_id uuid,
  p_feature text,
  p_base_limit integer,
  p_period text
) RETURNS jsonb
LANGUAGE plpgsql
SECURITY DEFINER
SET search_path = public, pg_temp
AS $$
DECLARE
  existing public.llm_quota_reservations%ROWTYPE;
  expired public.llm_quota_reservations%ROWTYPE;
  subscription_tier text := 'free';
  overage integer := 0;
  new_overage integer;
  current_count integer := 0;
  consumed_from text;
BEGIN
  IF p_user_id IS NULL
     OR p_reservation_id IS NULL
     OR p_feature IS NULL
     OR p_feature NOT IN ('llm_haiku', 'llm_sonnet', 'llm_opus')
     OR p_base_limit IS NULL
     OR p_base_limit < -1
     OR p_period IS NULL
     OR p_period NOT IN ('daily', 'weekly') THEN
    RAISE EXCEPTION 'invalid_llm_quota_reservation' USING ERRCODE = '22023';
  END IF;

  -- Same lock key as consume_quota so both paths serialise per user/feature.
  PERFORM pg_advisory_xact_lock(hashtextextended(p_user_id::text || ':' || p_feature, 20260928));

  SELECT * INTO existing
  FROM public.llm_quota_reservations
  WHERE id = p_reservation_id
  FOR UPDATE;

  IF FOUND THEN
    IF existing.user_id <> p_user_id OR existing.feature <> p_feature THEN
      RAISE EXCEPTION 'llm_quota_reservation_conflict' USING ERRCODE = 'PT409';
    END IF;
    RETURN jsonb_build_object(
      'allowed', existing.status IN ('reserved', 'completed'),
      'reservation_id', existing.id,
      'status', existing.status,
      'current', existing.current_count,
      'limit', existing.quota_limit,
      'period', existing.quota_period,
      'tier', existing.tier,
      'overage_credits', existing.overage_after,
      'consumed_from', existing.consumed_from
    );
  END IF;

  -- Reclaim crashed requests before calculating the next allowance.
  FOR expired IN
    SELECT *
    FROM public.llm_quota_reservations
    WHERE user_id = p_user_id
      AND feature = p_feature
      AND status = 'reserved'
      AND lease_expires_at <= now()
    FOR UPDATE
  LOOP
    UPDATE public.daily_usage
    SET count = greatest(count - 1, 0)
    WHERE user_id = expired.user_id
      AND date = expired.usage_date
      AND feature = expired.feature;

    IF expired.consumed_from = 'overage' THEN
      UPDATE public.subscriptions
      SET overage_credits = overage_credits + 1,
          updated_at = now()
      WHERE user_id = expired.user_id;
    END IF;

    UPDATE public.llm_quota_reservations
    SET status = 'released', finalized_at = now(), release_reason = 'lease_expired'
    WHERE id = expired.id;
  END LOOP;

  SELECT coalesce(tier, 'free'), coalesce(overage_credits, 0)
  INTO subscription_tier, overage
  FROM public.subscriptions
  WHERE user_id = p_user_id
  FOR UPDATE;

  IF NOT FOUND THEN
    subscription_tier := 'free';
    overage := 0;
  END IF;

  -- Not available: never spend credits on a model the tier does not include.
  IF p_base_limit = 0 THEN
    RETURN jsonb_build_object(
      'allowed', false,
      'reservation_id', NULL,
      'status', 'denied',
      'current', 0,
      'limit', 0,
      'period', p_period,
      'tier', subscription_tier,
      'overage_credits', overage,
      'consumed_from', 'none'
    );
  END IF;

  -- daily_usage already includes units held by in-flight reservations.
  SELECT coalesce(sum(count), 0)::integer INTO current_count
  FROM public.daily_usage
  WHERE user_id = p_user_id
    AND feature = p_feature
    AND date >= CASE WHEN p_period = 'weekly' THEN CURRENT_DATE - 6 ELSE CURRENT_DATE END
    AND date <= CURRENT_DATE;

  IF p_base_limit = -1 THEN
    consumed_from := 'unlimited';
  ELSIF current_count < p_base_limit THEN
    consumed_from := 'base';
  ELSE
    UPDATE public.subscriptions
    SET overage_credits = overage_credits - 1,
        updated_at = now()
    WHERE user_id = p_user_id
      AND overage_credits > 0
    RETURNING overage_credits INTO new_overage;

    IF NOT FOUND THEN
      RETURN jsonb_build_object(
        'allowed', false,
        'reservation_id', NULL,
        'status', 'denied',
        'current', current_count,
        'limit', p_base_limit,
        'period', p_period,
        'tier', subscription_tier,
        'overage_credits', 0,
        'consumed_from', 'none'
      );
    END IF;

    overage := new_overage;
    consumed_from := 'overage';
  END IF;

  INSERT INTO public.daily_usage(user_id, date, feature, count)
  VALUES (p_user_id, CURRENT_DATE, p_feature, 1)
  ON CONFLICT (user_id, date, feature)
  DO UPDATE SET count = public.daily_usage.count + 1;

  current_count := current_count + 1;
  INSERT INTO public.llm_quota_reservations(
    id, user_id, feature, consumed_from, tier, quota_period, quota_limit,
    current_count, overage_after, lease_expires_at
  ) VALUES (
    p_reservation_id, p_user_id, p_feature, consumed_from, subscription_tier, p_period, p_base_limit,
    current_count, overage, now() + interval '10 minutes'
  );

  RETURN jsonb_build_object(
    'allowed', true,
    'reservation_id', p_reservation_id,
    'status', 'reserved',
    'current', current_count,
    'limit', p_base_limit,
    'period', p_period,
    'tier', subscription_tier,
    'overage_credits', overage,
    'consumed_from', consumed_from
  );
END;
$$;

CREATE OR REPLACE FUNCTION public.finalize_llm_quota(
  p_reservation_id uuid,
  p_succeeded boolean
) RETURNS jsonb
LANGUAGE plpgsql
SECURITY DEFINER
SET search_path = public, pg_temp
AS $$
DECLARE
  reservation public.llm_quota_reservations%ROWTYPE;
  final_status text;
BEGIN
  IF p_reservation_id IS NULL OR p_succeeded IS NULL THEN
    RAISE EXCEPTION 'invalid_llm_quota_finalize' USING ERRCODE = '22023';
  END IF;

  SELECT * INTO reservation
  FROM public.llm_quota_reservations
  WHERE id = p_reservation_id;
  IF NOT FOUND THEN
    RAISE EXCEPTION 'llm_quota_reservation_not_found' USING ERRCODE = 'P0002';
  END IF;

  PERFORM pg_advisory_xact_lock(hashtextextended(reservation.user_id::text || ':' || reservation.feature, 20260928));
  SELECT * INTO reservation
  FROM public.llm_quota_reservations
  WHERE id = p_reservation_id
  FOR UPDATE;

  IF reservation.status <> 'reserved' THEN
    RETURN jsonb_build_object('reservation_id', reservation.id, 'status', reservation.status);
  END IF;

  IF p_succeeded THEN
    final_status := 'completed';
  ELSE
    UPDATE public.daily_usage
    SET count = greatest(count - 1, 0)
    WHERE user_id = reservation.user_id
      AND date = reservation.usage_date
      AND feature = reservation.feature;

    IF reservation.consumed_from = 'overage' THEN
      UPDATE public.subscriptions
      SET overage_credits = overage_credits + 1,
          updated_at = now()
      WHERE user_id = reservation.user_id;
    END IF;
    final_status := 'released';
  END IF;

  UPDATE public.llm_quota_reservations
  SET status = final_status,
      finalized_at = now(),
      release_reason = CASE WHEN p_succeeded THEN NULL ELSE 'provider_failed' END
  WHERE id = reservation.id;

  RETURN jsonb_build_object('reservation_id', reservation.id, 'status', final_status);
END;
$$;

REVOKE ALL ON FUNCTION public.reserve_llm_quota(uuid, uuid, text, integer, text) FROM PUBLIC, anon, authenticated;
REVOKE ALL ON FUNCTION public.finalize_llm_quota(uuid, boolean) FROM PUBLIC, anon, authenticated;
GRANT EXECUTE ON FUNCTION public.reserve_llm_quota(uuid, uuid, text, integer, text) TO service_role;
GRANT EXECUTE ON FUNCTION public.finalize_llm_quota(uuid, boolean) TO service_role;

COMMENT ON TABLE public.llm_quota_reservations IS
  'Service-only leases that reserve LLM quota before the Anthropic call in llm-proxy and refund failed or expired generations.';