-- Fix the trigger to actually call the distribution function
-- First, drop the existing trigger and function
drop trigger if exists trigger_distribute_benefits_on_update on medreport.company_params;
drop function if exists medreport.trigger_distribute_benefits();

-- Create a new trigger function that actually distributes benefits
create or replace function medreport.trigger_distribute_benefits()
returns trigger
language plpgsql
as $$
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
        -- Distribute benefits to all company members
        perform medreport.distribute_health_benefits(
            new.account_id,
            new.benefit_amount,
            coalesce(new.benefit_occurance, 'yearly')
        );
    end if;
    
    return new;
end;
$$;

-- Recreate the trigger
create trigger trigger_distribute_benefits_on_update
    after update on medreport.company_params
    for each row
    when (old.benefit_amount is distinct from new.benefit_amount or old.benefit_occurance is distinct from new.benefit_occurance)
    execute function medreport.trigger_distribute_benefits();

-- Create a table to track periodic benefit distributions
create table "medreport"."benefit_distribution_schedule" (
    "id" uuid not null default gen_random_uuid(),
    "company_id" uuid not null,
    "benefit_amount" numeric not null,
    "benefit_occurrence" text not null,
    "last_distributed_at" timestamp with time zone,
    "next_distribution_at" timestamp with time zone not null,
    "is_active" boolean not null default true,
    "created_at" timestamp with time zone not null default now(),
    "updated_at" timestamp with time zone not null default now()
);

-- Add constraints
alter table "medreport"."benefit_distribution_schedule" add constraint "benefit_distribution_schedule_pkey" primary key (id);
alter table "medreport"."benefit_distribution_schedule" add constraint "benefit_distribution_schedule_company_id_fkey" 
    foreign key (company_id) references medreport.accounts(id) on delete cascade;

-- Add unique constraint on company_id for upsert functionality
create unique index "ix_benefit_distribution_schedule_company_id_unique" 
on "medreport"."benefit_distribution_schedule" (company_id) 
where is_active = true;

-- Add indexes
create index "ix_benefit_distribution_schedule_next_distribution" on "medreport"."benefit_distribution_schedule" (next_distribution_at) where is_active = true;
create index "ix_benefit_distribution_schedule_company_id" on "medreport"."benefit_distribution_schedule" (company_id);

-- Enable RLS
alter table "medreport"."benefit_distribution_schedule" enable row level security;

-- Create RLS policies
create policy "Service role can manage all distribution schedules"
    on "medreport"."benefit_distribution_schedule"
    for all
    to service_role
    using (true)
    with check (true);

create policy "Users can view distribution schedules for their companies"
    on "medreport"."benefit_distribution_schedule"
    for select
    to authenticated
    using (
        company_id in (
            select account_id 
            from medreport.accounts_memberships 
            where user_id = auth.uid()
        )
    );

-- Grant permissions
grant select, insert, update, delete on table "medreport"."benefit_distribution_schedule" to service_role, authenticated;

-- Function to calculate next distribution date
create or replace function medreport.calculate_next_distribution_date(
    p_occurrence text,
    p_current_date timestamp with time zone default now()
)
returns timestamp with time zone
language plpgsql
as $$
declare
    next_date timestamp with time zone;
    current_year integer;
    current_month integer;
    current_quarter integer;
begin
    case p_occurrence
        when 'yearly' then
            -- First day of next year
            current_year := extract(year from p_current_date);
            next_date := make_date(current_year + 1, 1, 1);
            return next_date;
            
        when 'monthly' then
            -- First day of next month
            current_year := extract(year from p_current_date);
            current_month := extract(month from p_current_date);
            if current_month = 12 then
                next_date := make_date(current_year + 1, 1, 1);
            else
                next_date := make_date(current_year, current_month + 1, 1);
            end if;
            return next_date;
            
        when 'quarterly' then
            -- First day of next quarter
            current_year := extract(year from p_current_date);
            current_month := extract(month from p_current_date);
            current_quarter := ((current_month - 1) / 3) + 1;
            
            if current_quarter = 4 then
                next_date := make_date(current_year + 1, 1, 1);
            else
                next_date := make_date(current_year, (current_quarter * 3) + 1, 1);
            end if;
            return next_date;
            
        else
            -- Default to yearly
            current_year := extract(year from p_current_date);
            next_date := make_date(current_year + 1, 1, 1);
            return next_date;
    end case;
end;
$$;

-- Grant execute permission
grant execute on function medreport.calculate_next_distribution_date(text, timestamp with time zone) to authenticated, service_role;

-- Function to process periodic benefit distributions
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.company_id,
            schedule_record.benefit_amount,
            schedule_record.benefit_occurrence
        );
        
        -- 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;
$$;

-- Grant execute permission
grant execute on function medreport.process_periodic_benefit_distributions() to service_role;

-- Function to create or update a benefit distribution schedule
create or replace function medreport.upsert_benefit_distribution_schedule(
    p_company_id uuid,
    p_benefit_amount numeric,
    p_benefit_occurrence text
)
returns void
language plpgsql
as $$
declare
    next_distribution_date timestamp with time zone;
    existing_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 existing_record_id
    from medreport.benefit_distribution_schedule
    where company_id = p_company_id
    limit 1;
    
    if existing_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 = existing_record_id;
    else
        -- Insert new record
        insert into medreport.benefit_distribution_schedule (
            company_id,
            benefit_amount,
            benefit_occurrence,
            next_distribution_at
        ) values (
            p_company_id,
            p_benefit_amount,
            p_benefit_occurrence,
            next_distribution_date
        );
    end if;
end;
$$;

-- Grant execute permission
grant execute on function medreport.upsert_benefit_distribution_schedule(uuid, numeric, text) to service_role, authenticated;
grant execute on function medreport.distribute_health_benefits(uuid, numeric, text) to service_role, authenticated;
grant execute on function medreport.calculate_next_distribution_date(text, timestamp with time zone) to service_role, authenticated;

-- Also grant permissions to the original functions from the first migration
grant execute on function medreport.get_account_balance(uuid) to service_role, authenticated;
grant execute on function medreport.consume_account_balance(uuid, numeric, text, text) to service_role, authenticated;

-- Update the trigger to also create/update the distribution schedule
create or replace function medreport.trigger_distribute_benefits()
returns trigger
language plpgsql
security definer
as $$
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
        -- Distribute benefits to all company members immediately
        perform medreport.distribute_health_benefits(
            new.account_id,
            new.benefit_amount,
            coalesce(new.benefit_occurance, 'yearly')
        );
        
        -- Create or update the distribution schedule for future distributions
        perform medreport.upsert_benefit_distribution_schedule(
            new.account_id,
            new.benefit_amount,
            coalesce(new.benefit_occurance, 'yearly')
        );
    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 a function to manually trigger benefit distribution (for testing)
create or replace function medreport.trigger_benefit_distribution(p_company_id uuid)
returns void
language plpgsql
as $$
declare
    company_params record;
begin
    -- Get company params
    select * into company_params
    from medreport.company_params
    where account_id = p_company_id;
    
    if found and company_params.benefit_amount > 0 then
        -- Distribute benefits
        perform medreport.distribute_health_benefits(
            p_company_id,
            company_params.benefit_amount,
            coalesce(company_params.benefit_occurance, 'yearly')
        );
    end if;
end;
$$;

-- Grant execute permission
grant execute on function medreport.trigger_benefit_distribution(uuid) to service_role, authenticated;
grant execute on function medreport.trigger_distribute_benefits() to service_role, authenticated;