#!/usr/bin/env python3
"""
Migration script to add quetiapine columns to existing CSV data.
This script will backup the existing CSV and add the new columns.
"""

import os
import shutil
from datetime import datetime

import pandas as pd


def migrate_csv_add_quetiapine(csv_file: str = "thechart_data.csv"):
    """Add quetiapine and quetiapine_doses columns to existing CSV."""

    if not os.path.exists(csv_file):
        print(f"CSV file {csv_file} not found. No migration needed.")
        return

    # Create backup
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    backup_file = f"{csv_file}.backup_quetiapine_{timestamp}"
    shutil.copy2(csv_file, backup_file)
    print(f"Backup created: {backup_file}")

    # Load existing data
    try:
        df = pd.read_csv(csv_file)
        print(f"Loaded {len(df)} rows from {csv_file}")

        # Check if quetiapine columns already exist
        if "quetiapine" in df.columns:
            print("Quetiapine columns already exist. No migration needed.")
            return

        # Add new columns
        # Insert quetiapine columns before the note column
        note_col_index = (
            df.columns.get_loc("note") if "note" in df.columns else len(df.columns)
        )

        # Insert quetiapine column
        df.insert(note_col_index, "quetiapine", 0)
        df.insert(note_col_index + 1, "quetiapine_doses", "")

        # Save updated CSV
        df.to_csv(csv_file, index=False)
        print(f"Successfully added quetiapine columns to {csv_file}")
        print(f"New column order: {list(df.columns)}")

    except Exception as e:
        print(f"Error during migration: {e}")
        # Restore backup on error
        if os.path.exists(backup_file):
            shutil.copy2(backup_file, csv_file)
            print("Restored backup due to error")


if __name__ == "__main__":
    migrate_csv_add_quetiapine()