ALTER TABLE medreport.account_balance_entries ADD COLUMN benefit_distribution_schedule_id uuid;

-- Also setting `benefit_distribution_schedule_id` value now
drop function if exists medreport.distribute_health_benefits(uuid, numeric, text);
create or replace function medreport.distribute_health_benefits(
    p_benefit_distribution_schedule_id uuid
)
returns void
language plpgsql
security definer
as $$
declare
    member_record record;
    expires_date timestamp with time zone;
    v_company_id uuid;
    v_benefit_amount numeric;
begin
    -- Expires on first day of next year.
    expires_date := date_trunc('year', now() + interval '1 year');

    -- Get company_id and benefit_amount from benefit_distribution_schedule
    select company_id, benefit_amount into v_company_id, v_benefit_amount
    from medreport.benefit_distribution_schedule
    where id = p_benefit_distribution_schedule_id;

    -- Get all personal accounts that are members of this company
    for member_record in
        select distinct a.id as personal_account_id
        from medreport.accounts a
        join medreport.accounts_memberships am on a.id = am.user_id
        where am.account_id = v_company_id
        and a.is_personal_account = true
    loop
        -- Check if there is already a balance entry for this personal account from the same company in same month
        if exists (
            select 1
            from medreport.account_balance_entries
            where entry_type = 'benefit'
            and account_id = member_record.personal_account_id
            and source_company_id = v_company_id
            and date_trunc('month', created_at) = date_trunc('month', now())
        ) then
            continue;
        end if;

        -- Insert balance entry for each personal account
        insert into medreport.account_balance_entries (
            account_id,
            amount,
            entry_type,
            description,
            source_company_id,
            created_by,
            expires_at,
            benefit_distribution_schedule_id
        ) values (
            member_record.personal_account_id,
            v_benefit_amount,
            'benefit',
            'Health benefit from company',
            v_company_id,
            auth.uid(),
            expires_date,
            p_benefit_distribution_schedule_id
        );
    end loop;
end;
$$;

grant execute on function medreport.distribute_health_benefits(uuid) to authenticated, service_role;

create or replace function medreport.process_periodic_benefit_distributions()
returns void
language plpgsql
as $$
declare
    schedule_record record;
    next_distribution_date timestamp with time zone;
begin
    -- Get all active schedules that are due for distribution
    for schedule_record in
        select *
        from medreport.benefit_distribution_schedule
        where is_active = true
        and next_distribution_at <= now()
    loop
        -- Distribute benefits
        perform medreport.distribute_health_benefits(
            schedule_record.id
        );
        
        -- Calculate next distribution date
        next_distribution_date := medreport.calculate_next_distribution_date(
            schedule_record.benefit_occurrence,
            now()
        );
        
        -- Update the schedule
        update medreport.benefit_distribution_schedule
        set 
            last_distributed_at = now(),
            next_distribution_at = next_distribution_date,
            updated_at = now()
        where id = schedule_record.id;
    end loop;
end;
$$;

DROP FUNCTION IF EXISTS medreport.upsert_benefit_distribution_schedule(uuid,numeric,text);
create or replace function medreport.upsert_benefit_distribution_schedule(
    p_company_id uuid,
    p_benefit_amount numeric,
    p_benefit_occurrence text
)
-- Return schedule row id
returns uuid
language plpgsql
as $$
declare
    next_distribution_date timestamp with time zone;
    record_id uuid;
begin
    -- Calculate next distribution date
    next_distribution_date := medreport.calculate_next_distribution_date(p_benefit_occurrence);
    
    -- Check if there's an existing record for this company
    select id into record_id
    from medreport.benefit_distribution_schedule
    where company_id = p_company_id
    limit 1;
    
    if record_id is not null then
        -- Update existing record
        update medreport.benefit_distribution_schedule
        set 
            benefit_amount = p_benefit_amount,
            benefit_occurrence = p_benefit_occurrence,
            next_distribution_at = next_distribution_date,
            is_active = true,
            updated_at = now()
        where id = record_id;
    else
        record_id := gen_random_uuid();
        
        -- Insert new record
        insert into medreport.benefit_distribution_schedule (
            id,
            company_id,
            benefit_amount,
            benefit_occurrence,
            next_distribution_at
        ) values (
            record_id,
            p_company_id,
            p_benefit_amount,
            p_benefit_occurrence,
            next_distribution_date
        );
    end if;

    return record_id;
end;
$$;

grant execute on function medreport.upsert_benefit_distribution_schedule(uuid, numeric, text) to authenticated, service_role;

create or replace function medreport.trigger_distribute_benefits()
returns trigger
language plpgsql
security definer
as $$
declare
    v_benefit_distribution_schedule_id uuid;
begin
    -- Only distribute if benefit_amount is set and greater than 0
    if new.benefit_amount is not null and new.benefit_amount > 0 then
        -- Create or update the distribution schedule for future distributions
        v_benefit_distribution_schedule_id := medreport.upsert_benefit_distribution_schedule(
            new.account_id,
            new.benefit_amount,
            coalesce(new.benefit_occurance, 'yearly')
        );

        -- Distribute benefits to all company members immediately
        perform medreport.distribute_health_benefits(
            v_benefit_distribution_schedule_id
        );
    else
        -- If benefit_amount is 0 or null, deactivate the schedule
        update medreport.benefit_distribution_schedule
        set is_active = false, updated_at = now()
        where company_id = new.account_id;
    end if;
    
    return new;
end;
$$;


CREATE OR REPLACE FUNCTION medreport.get_benefits_usages_for_company_members(p_account_id uuid)
returns table (
    personal_account_id uuid,
    benefit_amount numeric
)
language plpgsql
as $$
begin
    return query
        select
            abe.account_id as personal_account_id,
            sum(abe.amount) as benefit_amount
        from medreport.account_balance_entries abe
        where abe.source_company_id = p_account_id
        and abe.entry_type = 'benefit'
        group by abe.account_id;
end;
$$;

grant execute on function medreport.get_benefits_usages_for_company_members(uuid) to authenticated, service_role;