"""
Migrate location and state_label from asset_log_entries to asset_state_nodes.
1. Add columns to asset_state_nodes
2. For each interval: migrate location/state_label to first state node (or create one)
3. Drop columns from asset_log_entries
"""
import os, psycopg, uuid
from datetime import datetime, timezone
from dotenv import load_dotenv
from pathlib import Path

load_dotenv(Path(__file__).resolve().parents[1] / ".env")
conn = psycopg.connect(
    host=os.getenv("PGHOST"), port=os.getenv("PGPORT"),
    dbname=os.getenv("PGDATABASE"), user=os.getenv("PGUSER"),
    password=os.getenv("PGPASSWORD"))
cur = conn.cursor()
tz = timezone.utc
NOW = datetime.now(tz)

# Step 1: Add columns
print("Step 1: Adding location and state_label to asset_state_nodes...")
try:
    cur.execute("ALTER TABLE asset_state_nodes ADD COLUMN location VARCHAR(255)")
    print("  Added location column")
except Exception as e:
    if 'already exists' in str(e):
        print("  location column already exists")
    else:
        raise
try:
    cur.execute("ALTER TABLE asset_state_nodes ADD COLUMN state_label VARCHAR(120)")
    print("  Added state_label column")
except Exception as e:
    if 'already exists' in str(e):
        print("  state_label column already exists")
    else:
        raise

# Step 2: Migrate data
print("\nStep 2: Migrating location and state_label to state nodes...")
cur.execute("""
    SELECT le.id, le.location, le.state_label, le.title, le.start_at,
           (SELECT COUNT(*) FROM asset_state_nodes sn WHERE sn.entry_id = le.id) as nc
    FROM asset_log_entries le
    WHERE le.location IS NOT NULL OR le.state_label IS NOT NULL
""")
rows = cur.fetchall()
print(f"  Found {len(rows)} entries with location/state_label")

created = 0
updated = 0
for row in rows:
    entry_id, loc, st_lbl, title, start_at, nc = row
    if nc > 0:
        # Update the first (earliest) state node of this interval
        cur.execute("""
            UPDATE asset_state_nodes
            SET location = COALESCE(%s, location),
                state_label = COALESCE(%s, state_label),
                updated_at = %s
            WHERE id = (
                SELECT id FROM asset_state_nodes
                WHERE entry_id = %s
                ORDER BY at_time ASC LIMIT 1
            )
        """, (loc, st_lbl, NOW, entry_id))
        updated += 1
    else:
        # No state nodes — create one at the interval's start_at
        cur.execute("""
            INSERT INTO asset_state_nodes (id, entry_id, title, detail, at_time, location, state_label, created_at, updated_at)
            VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s)
        """, (str(uuid.uuid4()), entry_id, title or '(Start)',
              f'Auto-created during location migration. Original interval: {title}',
              start_at, loc, st_lbl, NOW, NOW))
        created += 1
print(f"  Updated {updated} existing state nodes, created {created} new ones")

# Step 3: Drop columns
print("\nStep 3: Dropping location and state_label from asset_log_entries...")
cur.execute("ALTER TABLE asset_log_entries DROP COLUMN location")
print("  Dropped location")
cur.execute("ALTER TABLE asset_log_entries DROP COLUMN state_label")
print("  Dropped state_label")

conn.commit()
print("\nMigration complete.")
conn.close()