#!/usr/bin/env python3
"""
Migration: add security columns.
  users                   — totp_secret, totp_enabled
  schwab_connections      — refresh_token_expires_at

Run once:  python scripts/add_security_columns.py
Safe to re-run — skips columns that already exist.
"""
import sys, os
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from app import create_app
from app.extensions import db

app = create_app()

COLUMNS = [
    ("users", "totp_secret",
     "VARCHAR(64) NULL DEFAULT NULL"),
    ("users", "totp_enabled",
     "TINYINT(1) NOT NULL DEFAULT 0"),
    ("schwab_connections", "refresh_token_expires_at",
     "DATETIME NULL DEFAULT NULL"),
]

with app.app_context():
    with db.engine.connect() as conn:
        for table, column, definition in COLUMNS:
            exists = conn.execute(db.text(
                "SELECT COUNT(*) FROM information_schema.columns "
                "WHERE table_schema = DATABASE() "
                "AND table_name = :t AND column_name = :c"
            ), {"t": table, "c": column}).scalar()

            if exists:
                print(f"  {table}.{column} already exists — skipped.")
            else:
                conn.execute(db.text(
                    f"ALTER TABLE {table} ADD COLUMN {column} {definition}"
                ))
                conn.commit()
                print(f"  Added {table}.{column}.")

print("Done.")