#!/usr/bin/env python3
"""
update_labor_hours.py
Updates locationInfo.laborHours in MySQL from a labor budget xlsx spreadsheet.

Usage:
    python scripts/update_labor_hours.py <path_to_xlsx> [--dry-run]
"""
import argparse
import csv
import os
import sys
from datetime import datetime
from pathlib import Path

import openpyxl
import pymysql
from dotenv import load_dotenv

SCRIPT_DIR = Path(__file__).parent


def read_spreadsheet(xlsx_path):
    """
    Parse xlsx and return list of (location_id, weekly_hours) tuples.
    weekly_hours = round(daily_amount * 7). Zero-hour locations return 0.
    Columns: B (index 1) = LocationID, I (index 8) = Amount (daily hours).
    """
    wb = openpyxl.load_workbook(xlsx_path)
    ws = wb.active
    results = []
    for i, row in enumerate(ws.iter_rows(values_only=True)):
        if i == 0:
            continue  # skip header row
        location_id = row[1]
        amount = row[8]
        if location_id is None:
            continue
        weekly = round(amount * 7) if amount else 0
        results.append((location_id, weekly))
    return results


def backup_labor_hours(conn):
    """
    Write current locationInfo.laborHours to a timestamped CSV next to this script.
    Returns the Path of the created file. Raises on any failure.
    """
    timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
    backup_path = SCRIPT_DIR / f'labor_hours_backup_{timestamp}.csv'
    cursor = conn.cursor()
    cursor.execute(
        "SELECT locationID, locationName, laborHours FROM locationInfo ORDER BY locationID"
    )
    rows = cursor.fetchall()
    cursor.close()
    with open(backup_path, 'w', newline='') as f:
        writer = csv.writer(f)
        writer.writerow(['locationID', 'locationName', 'laborHours'])
        writer.writerows(rows)
    return backup_path


def apply_updates(conn, spreadsheet_data, dry_run=False):
    """
    Update locationInfo.laborHours for each (location_id, weekly_hours) pair.
    Prints a line per location. Returns summary dict with keys:
      updated, set_to_zero, not_found
    """
    cursor = conn.cursor()
    try:
        cursor.execute(
            "SELECT locationID, locationName, laborHours FROM locationInfo"
        )
        db_rows = {row[0]: (row[1], row[2]) for row in cursor.fetchall()}

        updated = 0
        set_to_zero = 0
        not_found = 0

        for location_id, weekly_hours in spreadsheet_data:
            if location_id not in db_rows:
                print(f"[WARN]    LocationID {location_id} not found in DB — skipping")
                not_found += 1
                continue

            name, current = db_rows[location_id]
            prefix = "[DRY RUN]" if dry_run else "[UPDATE] "
            print(f"{prefix} {location_id} {name}: {current} → {weekly_hours}")

            if not dry_run:
                cursor.execute(
                    "UPDATE locationInfo SET laborHours = %s WHERE locationID = %s",
                    (weekly_hours, location_id)
                )

            if weekly_hours == 0:
                set_to_zero += 1
            else:
                updated += 1

        if not dry_run:
            conn.commit()

        return {'updated': updated, 'set_to_zero': set_to_zero, 'not_found': not_found}
    finally:
        cursor.close()


def load_env():
    """Load DB credentials from scripts/.env. Exits if file not found."""
    env_path = SCRIPT_DIR / '.env'
    if not env_path.exists():
        print(f"ERROR: .env not found at {env_path}")
        print("Create scripts/.env with: DB_HOST, DB_PORT, DB_NAME, DB_USER, DB_PASS")
        sys.exit(1)
    load_dotenv(env_path)


def get_db_connection():
    """Return a pymysql connection using env vars. Raises on failure."""
    return pymysql.connect(
        host=os.environ['DB_HOST'],
        port=int(os.environ.get('DB_PORT', 3306)),
        db=os.environ['DB_NAME'],
        user=os.environ['DB_USER'],
        password=os.environ['DB_PASS'],
        charset='utf8mb4',
    )


def main():
    parser = argparse.ArgumentParser(
        description='Update locationInfo.laborHours from a labor budget xlsx spreadsheet.'
    )
    parser.add_argument('xlsx_path', help='Path to the xlsx budget spreadsheet')
    parser.add_argument(
        '--dry-run', action='store_true',
        help='Print planned changes without writing to the database'
    )
    args = parser.parse_args()

    load_env()

    print(f"Reading spreadsheet: {args.xlsx_path}")
    try:
        spreadsheet_data = read_spreadsheet(args.xlsx_path)
    except Exception as e:
        print(f"ERROR: Could not read spreadsheet: {e}")
        sys.exit(1)
    print(f"Found {len(spreadsheet_data)} locations in spreadsheet\n")

    try:
        conn = get_db_connection()
    except pymysql.Error as e:
        print(f"ERROR: Could not connect to database: {e}")
        sys.exit(1)

    if not args.dry_run:
        try:
            backup_path = backup_labor_hours(conn)
            print(f"Backup saved: {backup_path}\n")
        except Exception as e:
            print(f"ERROR: Backup failed — aborting before any writes: {e}")
            conn.close()
            sys.exit(1)

    summary = apply_updates(conn, spreadsheet_data, dry_run=args.dry_run)
    conn.close()

    mode = "DRY RUN " if args.dry_run else ""
    print(
        f"\n{mode}Summary: {summary['updated']} updated, "
        f"{summary['set_to_zero']} set to 0, "
        f"{summary['not_found']} not found in DB"
    )


if __name__ == '__main__':
    main()