"""
Add a favorite_counts table to aleviate the query time.
"""

from yoyo import step

__depends__ = {"20240223_00_ASDAS-reset_relation_id_seq"}

steps = [
    step(
        """
		CREATE TABLE IF NOT EXISTS favorite_counts (
            service varchar NOT NULL,
            artist_id varchar NOT NULL,
            favorite_count INTEGER DEFAULT 0,
            PRIMARY KEY (service, artist_id)
        );
        """
    ),
    step(
        """
        CREATE OR REPLACE FUNCTION update_favorite_count() 
        RETURNS TRIGGER AS $$
        BEGIN
            -- For INSERT event: increment the count
            IF TG_OP = 'INSERT' THEN
                INSERT INTO favorite_counts (service, artist_id, favorite_count)
                VALUES (NEW.service, NEW.artist_id, 1)
                ON CONFLICT (service, artist_id) 
                DO UPDATE SET favorite_count = favorite_counts.favorite_count + 1;

            -- For DELETE event: decrement the count
            ELSIF TG_OP = 'DELETE' THEN
                UPDATE favorite_counts 
                SET favorite_count = favorite_count - 1
                WHERE service = OLD.service AND artist_id = OLD.artist_id;
            END IF;

            RETURN NULL; -- triggers on INSERT/DELETE return NULL
        END;
        $$ LANGUAGE plpgsql;
    """
    ),
    step(
        """
        CREATE TRIGGER update_favorite_count_trigger
        AFTER INSERT OR DELETE ON account_artist_favorite
        FOR EACH ROW
        EXECUTE FUNCTION update_favorite_count();
        """
    ),
    step(
        """
        INSERT INTO favorite_counts (service, artist_id, favorite_count)
        SELECT service, artist_id, COUNT(*) AS favorite_count
        FROM account_artist_favorite
        GROUP BY service, artist_id
        ON CONFLICT (service, artist_id)
        DO UPDATE SET favorite_count = EXCLUDED.favorite_count;
        """
    ),
]