"""Split the bounty budget into two tracks (产研 / IT).

The programme is funded from two pots, not one: product-engineering
findings are charged to the R&D security budget, IT-owned systems to the
IT budget. One number could not answer "how much of MY budget is left",
and the two owners approve payment separately.

Adds ``bounty_reports.budget_track`` (``product`` default, matching where
the overwhelming majority of findings land) and rewrites the
``bounty.budget`` platform setting into a per-track shape, carrying any
existing flat default across into the product track so nothing resets.

Revision ID: 539
Revises: 538
"""
from __future__ import annotations

import json

import sqlalchemy as sa

from alembic import op

revision = "539"
down_revision = "538"
branch_labels = None
depends_on = None

REPORTS = "bounty_reports"
COLUMN = "budget_track"
SETTING = "bounty.budget"

TRACK_PRODUCT = "product"
TRACK_IT = "it"
DEFAULT_TRACK_BUDGET = 5000


def _column_exists(bind, table: str, column: str) -> bool:
    return bind.execute(
        sa.text(
            "SELECT 1 FROM information_schema.columns "
            "WHERE table_schema = DATABASE() AND table_name = :t "
            "AND column_name = :c"
        ),
        {"t": table, "c": column},
    ).scalar() is not None


def _index_exists(bind, table: str, name: str) -> bool:
    return bind.execute(
        sa.text(
            "SELECT 1 FROM information_schema.statistics "
            "WHERE table_schema = DATABASE() AND table_name = :t "
            "AND index_name = :n LIMIT 1"
        ),
        {"t": table, "n": name},
    ).scalar() is not None


def upgrade() -> None:
    bind = op.get_bind()

    if not _column_exists(bind, REPORTS, COLUMN):
        # NOT NULL with a server default: every existing row becomes
        # ``product``, which is where the historical findings actually
        # belong (they are all product/platform vulnerabilities). A
        # nullable column would leave an "unknown track" bucket that no
        # budget owner would claim.
        op.add_column(
            REPORTS,
            sa.Column(
                COLUMN, sa.String(16), nullable=False,
                server_default=TRACK_PRODUCT,
            ),
        )
    # The budget report groups by (track, fiscal_year, payment_status);
    # the existing ix_bounty_reports_budget covers the last two.
    if not _index_exists(bind, REPORTS, "ix_bounty_reports_track_budget"):
        op.create_index(
            "ix_bounty_reports_track_budget", REPORTS,
            [COLUMN, "fiscal_year", "payment_status"],
        )

    # ── Setting: flat → per-track, preserving whatever is configured ──
    row = bind.execute(
        sa.text("SELECT `value` FROM platform_settings WHERE `key` = :k"),
        {"k": SETTING},
    ).scalar()
    existing: dict = {}
    if row:
        try:
            existing = json.loads(row) if isinstance(row, str) else dict(row)
        except (TypeError, ValueError):
            existing = {}

    if "tracks" not in existing:
        # An operator-edited flat default carries into the product track
        # rather than being silently replaced by the shipped 5000.
        legacy_default = existing.get("default_usd", DEFAULT_TRACK_BUDGET)
        legacy_by_fy = existing.get("by_fy") or {}
        value = {
            "tracks": {
                TRACK_PRODUCT: {
                    "default_usd": legacy_default,
                    "by_fy": legacy_by_fy,
                },
                TRACK_IT: {
                    "default_usd": DEFAULT_TRACK_BUDGET,
                    "by_fy": {},
                },
            },
            "currency": existing.get("currency", "USD"),
        }
        description = (
            "漏洞赏金财年预算(USD),按预算块分开:product=产研,it=IT。"
            "每块可按财年覆盖 default_usd。"
        )
        if row:
            bind.execute(
                sa.text(
                    "UPDATE platform_settings SET `value` = :v, "
                    "description = :d, updated_at = UTC_TIMESTAMP() "
                    "WHERE `key` = :k"
                ),
                {"k": SETTING, "v": json.dumps(value), "d": description},
            )
        else:
            bind.execute(
                sa.text(
                    "INSERT INTO platform_settings "
                    "(`key`, `value`, description, created_at, updated_at) "
                    "VALUES (:k, :v, :d, UTC_TIMESTAMP(), UTC_TIMESTAMP())"
                ),
                {"k": SETTING, "v": json.dumps(value), "d": description},
            )


def downgrade() -> None:
    bind = op.get_bind()
    # Index before column: TiDB refuses DROP COLUMN while an index covers
    # it (error 8200), and MySQL's silent auto-rewrite hides that locally.
    if _index_exists(bind, REPORTS, "ix_bounty_reports_track_budget"):
        op.drop_index("ix_bounty_reports_track_budget", table_name=REPORTS)
    if _column_exists(bind, REPORTS, COLUMN):
        op.drop_column(REPORTS, COLUMN)
    # Collapse the setting back to the flat shape using the product track,
    # so a downgrade does not leave a shape the old code cannot read.
    row = bind.execute(
        sa.text("SELECT `value` FROM platform_settings WHERE `key` = :k"),
        {"k": SETTING},
    ).scalar()
    if row:
        try:
            current = json.loads(row) if isinstance(row, str) else dict(row)
        except (TypeError, ValueError):
            current = {}
        product = (current.get("tracks") or {}).get(TRACK_PRODUCT) or {}
        bind.execute(
            sa.text(
                "UPDATE platform_settings SET `value` = :v, "
                "updated_at = UTC_TIMESTAMP() WHERE `key` = :k"
            ),
            {
                "k": SETTING,
                "v": json.dumps({
                    "default_usd": product.get(
                        "default_usd", DEFAULT_TRACK_BUDGET,
                    ),
                    "by_fy": product.get("by_fy") or {},
                    "currency": current.get("currency", "USD"),
                }),
            },
        )