import os
from collections.abc import Sequence
from pathlib import Path

from alembic import op


def get_revision() -> str:
    return Path(__file__).name.split("-")[0]


def get_down_revision() -> str | None:
    versions_path = Path(__file__).parent
    migration_files = sorted([f.name for f in versions_path.iterdir() if f.is_file()])
    for i, file in enumerate(migration_files):
        if file == Path(__file__).name:
            return None if i == 0 else migration_files[i - 1].split("-")[0]
    return None


# revision identifiers, used by Alembic.
revision: str = get_revision()
down_revision: str | None = get_down_revision()
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None

username = os.getenv("POSTGRES_CLIENT_DEVELOPER_USER", "client_developer")
password = os.getenv("POSTGRES_CLIENT_DEVELOPER_PASSWORD", "client_developer")


def upgrade() -> None:
    op.execute(
        """--sql
        CREATE OR REPLACE FUNCTION update_quota_allocation(
            p_client_name TEXT,
            p_amount      NUMERIC,
            p_used        NUMERIC
        )
        RETURNS void
        LANGUAGE plpgsql
        SECURITY DEFINER
        SET search_path = public  -- 必ず固定する（任意のスキーマ利用を防止）
        AS $$
        DECLARE
            v_tenant_id           UUID;
            v_quota_allocation_id UUID;
            v_new_status          quota_allocation_status;
        BEGIN
            -- 引数チェック
            IF p_used > p_amount THEN
                RAISE EXCEPTION 'used (%) cannot be greater than amount (%)', p_used, p_amount;
                RETURN;
            END IF;

            -- クライアント名から tenant_id を取得
            SELECT tenant_id
                INTO v_tenant_id
            FROM client_tenants
            WHERE name = p_client_name
            LIMIT 1;

            IF v_tenant_id IS NULL THEN
                RAISE EXCEPTION 'client % not found in client_tenants', p_client_name;
                RETURN;
            END IF;

            -- quota_allocations は 1 件だけ存在する前提
            BEGIN
                SELECT id
                INTO STRICT v_quota_allocation_id
                FROM quota_allocations
                WHERE tenant_id = v_tenant_id;

            EXCEPTION
                WHEN NO_DATA_FOUND THEN
                    RAISE EXCEPTION 'quota_allocation for tenant % not found', v_tenant_id;
                    RETURN;

                WHEN TOO_MANY_ROWS THEN
                    RAISE EXCEPTION 'found multiple quota_allocation_ids for tenant %, expected 1 (client: %)',
                        v_tenant_id,
                        p_client_name;
            END;

            -- status 判定
            IF p_used >= p_amount THEN
                v_new_status := 'exceeded'::quota_allocation_status;
            ELSE
                v_new_status := 'active'::quota_allocation_status;
            END IF;

            -- UPDATE 実行
            UPDATE quota_allocations
            SET amount = p_amount,
                used   = p_used,
                status = v_new_status
            WHERE id = v_quota_allocation_id;

        END;
        $$;
        """,
    )
    op.execute(f"GRANT EXECUTE ON FUNCTION update_quota_allocation(TEXT, NUMERIC, NUMERIC) TO {username}")


def downgrade() -> None:
    op.execute(f"REVOKE EXECUTE ON FUNCTION update_quota_allocation(TEXT, NUMERIC, NUMERIC) FROM {username}")
    op.execute("DROP FUNCTION IF EXISTS update_quota_allocation(TEXT, NUMERIC, NUMERIC);")
