# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/17_utils_migrate.ipynb.

# %% auto 0
__all__ = ['logger', 'Migration', 'parse_migration', 'discover_migrations', 'MigrationStatus', 'get_status', 'apply_migrations',
           'DestructiveRollbackError', 'rollback', 'run_migrations', 'print_status']

# %% ../nbs/17_utils_migrate.ipynb
from dataclasses import dataclass
from typing import Optional, List
from pathlib import Path
from datetime import datetime
import hashlib
import re
import logging
import os
from fastsql import Database
from sqlalchemy import text

logger = logging.getLogger(__name__)

# %% ../nbs/17_utils_migrate.ipynb
@dataclass
class Migration:
    version: int
    name: str
    filepath: Path
    up_sql: str
    down_sql: str
    checksum: str
    
    @property
    def is_destructive(self) -> bool:
        patterns = [r'\bDROP\s+TABLE\b', r'\bDROP\s+COLUMN\b', r'\bTRUNCATE\b', r'\bDELETE\s+FROM\b']
        for pattern in patterns:
            if re.search(pattern, self.down_sql, re.IGNORECASE):
                return True
        return False

# %% ../nbs/17_utils_migrate.ipynb
def parse_migration(filepath: Path) -> Migration:
    match = re.match(r'^(\d+)_(.+)\.sql$', filepath.name)
    if not match:
        raise ValueError(f"Invalid migration filename: {filepath.name}. Expected: 001_description.sql")
    version = int(match.group(1))
    name = match.group(2)
    content = filepath.read_text(encoding='utf-8')
    up_match = re.search(r'--\s*UP\s*--\s*\n(.*?)(?=--\s*DOWN\s*--|$)', content, re.DOTALL | re.IGNORECASE)
    down_match = re.search(r'--\s*DOWN\s*--\s*\n(.*?)$', content, re.DOTALL | re.IGNORECASE)
    if not up_match:
        raise ValueError(f"Migration {filepath.name} missing '-- UP --' section")
    if not down_match:
        raise ValueError(f"Migration {filepath.name} missing '-- DOWN --' section")
    up_sql = up_match.group(1).strip()
    down_sql = down_match.group(1).strip()
    if not up_sql:
        raise ValueError(f"Migration {filepath.name} has empty UP section")
    checksum = hashlib.sha256(up_sql.encode()).hexdigest()[:16]
    return Migration(version=version, name=name, filepath=filepath, up_sql=up_sql, down_sql=down_sql, checksum=checksum)

# %% ../nbs/17_utils_migrate.ipynb
def discover_migrations(migrations_dir) -> List[Migration]:
    path = Path(migrations_dir)
    if not path.exists():
        logger.warning(f"Migrations directory not found: {path}")
        return []
    migrations = []
    for sql_file in sorted(path.glob('*.sql')):
        try:
            migrations.append(parse_migration(sql_file))
        except ValueError as e:
            logger.warning(f"Skipping invalid migration: {e}")
    migrations.sort(key=lambda m: m.version)
    versions = [m.version for m in migrations]
    if len(versions) != len(set(versions)):
        raise ValueError("Duplicate migration versions found")
    return migrations

# %% ../nbs/17_utils_migrate.ipynb
@dataclass
class MigrationStatus:
    current_version: int
    pending_count: int
    applied_migrations: List[dict]
    pending_migrations: List[Migration]

# %% ../nbs/17_utils_migrate.ipynb
def _ensure_migrations_table(db: Database) -> None:
    db.execute(text("CREATE TABLE IF NOT EXISTS _migrations (version INTEGER PRIMARY KEY, name TEXT NOT NULL, checksum TEXT NOT NULL, applied_at TEXT NOT NULL)"))
    db.conn.commit()

def _get_applied_versions(db: Database) -> dict:
    _ensure_migrations_table(db)
    rows = db.execute(text("SELECT version, name, checksum, applied_at FROM _migrations ORDER BY version")).fetchall()
    return {row[0]: {'name': row[1], 'checksum': row[2], 'applied_at': row[3]} for row in rows}

# %% ../nbs/17_utils_migrate.ipynb
def get_status(db: Database, migrations_dir="migrations") -> MigrationStatus:
    applied = _get_applied_versions(db)
    all_migrations = discover_migrations(migrations_dir)
    current_version = max(applied.keys()) if applied else 0
    pending = [m for m in all_migrations if m.version not in applied]
    return MigrationStatus(
        current_version=current_version,
        pending_count=len(pending),
        applied_migrations=[{'version': v, **info} for v, info in applied.items()],
        pending_migrations=pending
    )

# %% ../nbs/17_utils_migrate.ipynb
def apply_migrations(db: Database, migrations_dir="migrations", target_version: int = None, dry_run: bool = False) -> List[Migration]:
    status = get_status(db, migrations_dir)
    pending = status.pending_migrations
    if target_version is not None:
        pending = [m for m in pending if m.version <= target_version]
    if not pending:
        logger.info("No pending migrations")
        return []
    if dry_run:
        logger.info(f"[DRY RUN] Would apply {len(pending)} migration(s):")
        for m in pending:
            logger.info(f"  {m.version:03d}_{m.name}")
        return pending
    applied = []
    for migration in pending:
        try:
            logger.info(f"Applying {migration.version:03d}_{migration.name}...")
            db.execute(text(migration.up_sql))
            db.execute(text("INSERT INTO _migrations (version, name, checksum, applied_at) VALUES (:v, :n, :c, :a)"),
                       {"v": migration.version, "n": migration.name, "c": migration.checksum, "a": datetime.utcnow().isoformat()})
            db.conn.commit()
            applied.append(migration)
            logger.info(f"  Applied {migration.version:03d}_{migration.name}")
        except Exception as e:
            db.conn.rollback()
            logger.error(f"  Failed {migration.version:03d}_{migration.name}: {e}")
            raise Exception(f"Migration {migration.version:03d}_{migration.name} failed: {e}") from e
    logger.info(f"Successfully applied {len(applied)} migration(s)")
    return applied

# %% ../nbs/17_utils_migrate.ipynb
class DestructiveRollbackError(Exception):
    pass

# %% ../nbs/17_utils_migrate.ipynb
def rollback(db: Database, migrations_dir="migrations", steps: int = 1, target_version: int = None, force: bool = False, dry_run: bool = False) -> List[Migration]:
    applied = _get_applied_versions(db)
    if not applied:
        logger.info("No migrations to rollback")
        return []
    all_migrations = discover_migrations(migrations_dir)
    migrations_by_version = {m.version: m for m in all_migrations}
    applied_versions = sorted(applied.keys(), reverse=True)
    if target_version is not None:
        to_rollback = [v for v in applied_versions if v > target_version]
    else:
        to_rollback = applied_versions[:steps]
    if not to_rollback:
        logger.info("Nothing to rollback")
        return []
    rollback_migrations = []
    for version in to_rollback:
        if version not in migrations_by_version:
            raise ValueError(f"Migration file for version {version} not found")
        rollback_migrations.append(migrations_by_version[version])
    destructive = [m for m in rollback_migrations if m.is_destructive]
    if destructive and not force:
        names = ', '.join(f"{m.version:03d}_{m.name}" for m in destructive)
        raise DestructiveRollbackError(f"Destructive rollback in: {names}. Use force=True. WARNING: Data loss!")
    if dry_run:
        logger.info(f"[DRY RUN] Would rollback {len(rollback_migrations)} migration(s):")
        for m in rollback_migrations:
            flag = " DESTRUCTIVE" if m.is_destructive else ""
            logger.info(f"  {m.version:03d}_{m.name}{flag}")
        return rollback_migrations
    rolled_back = []
    for migration in rollback_migrations:
        try:
            flag = " (DESTRUCTIVE)" if migration.is_destructive else ""
            logger.info(f"Rolling back {migration.version:03d}_{migration.name}{flag}...")
            db.execute(text(migration.down_sql))
            db.execute(text("DELETE FROM _migrations WHERE version = :v"), {"v": migration.version})
            db.conn.commit()
            rolled_back.append(migration)
            logger.info(f"  Rolled back {migration.version:03d}_{migration.name}")
        except Exception as e:
            db.conn.rollback()
            logger.error(f"  Rollback failed {migration.version:03d}_{migration.name}: {e}")
            raise Exception(f"Rollback of {migration.version:03d}_{migration.name} failed: {e}") from e
    logger.info(f"Successfully rolled back {len(rolled_back)} migration(s)")
    return rolled_back

# %% ../nbs/17_utils_migrate.ipynb
def run_migrations(migrations_dir="migrations", db_url: str = None, dry_run: bool = False) -> MigrationStatus:
    import urllib.parse
    if db_url is None:
        db_type = os.getenv("DB_TYPE", "POSTGRESQL")
        if db_type == "POSTGRESQL":
            user = os.getenv("DB_USER", "postgres")
            password = os.getenv("DB_PASS", "")
            host = os.getenv("DB_HOST", "localhost")
            port = os.getenv("DB_PORT", "5432")
            name = os.getenv("DB_NAME", "app")
            encoded_pass = urllib.parse.quote_plus(password)
            db_url = f"postgresql://{user}:{encoded_pass}@{host}:{port}/{name}"
        else:
            name = os.getenv("DB_NAME", "app")
            db_url = f"sqlite:///{name}.db"
    db = Database(db_url)
    try:
        apply_migrations(db, migrations_dir, dry_run=dry_run)
        return get_status(db, migrations_dir)
    finally:
        try:
            db.conn.close()
        except:
            pass

# %% ../nbs/17_utils_migrate.ipynb
def print_status(db: Database, migrations_dir="migrations") -> None:
    status = get_status(db, migrations_dir)
    print("Migration Status")
    print("================")
    print(f"Current version: {status.current_version:03d}")
    print(f"Pending: {status.pending_count} migration(s)")
    if status.applied_migrations:
        print("\nApplied:")
        for m in status.applied_migrations:
            print(f"  {m['version']:03d}_{m['name']} ({m['applied_at'][:10]})")
    if status.pending_migrations:
        print("\nPending:")
        for m in status.pending_migrations:
            print(f"  {m.version:03d}_{m.name}")
