#!/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"),
]

CREATE_TABLES = [
    """
    CREATE TABLE IF NOT EXISTS audit_logs (
        id          INT AUTO_INCREMENT PRIMARY KEY,
        timestamp   DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
        action      VARCHAR(64) NOT NULL,
        description VARCHAR(255),
        ip_address  VARCHAR(45),
        INDEX ix_audit_logs_timestamp (timestamp),
        INDEX ix_audit_logs_action    (action)
    ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4
    """
]

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}.")

        for sql in CREATE_TABLES:
            conn.execute(db.text(sql))
            conn.commit()
            print("  Created audit_logs table (or already exists).")

print("Done.")