"""
Migrate state labels: Transfer/In Transit -> start-transit / transit-complete.
Set previous_location and transit_location on start-transit nodes.
"""
import os, psycopg, re
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()
NOW = datetime.now(timezone.utc)

# 1. Fetch all state nodes with their parent entry info, ordered by time
cur.execute("""
    SELECT sn.id::text, sn.title, sn.detail, sn.at_time, sn.location, sn.state_label,
           le.title as entry_title, le.mission_label, le.start_at, le.end_at,
           a.name as asset_name
    FROM asset_state_nodes sn
    JOIN asset_log_entries le ON sn.entry_id = le.id
    JOIN assets a ON le.asset_id = a.id
    ORDER BY a.name, le.start_at, sn.at_time
""")
rows = cur.fetchall()
print(f"Processing {len(rows)} state nodes")

# Group state nodes by entry
from collections import defaultdict
entries = defaultdict(list)
for r in rows:
    eid = (r[6], r[7], r[8])  # entry_title, mission_label, start_at
    entries[eid].append(r)

updated = 0

for ekey, nodes in entries.items():
    entry_title, mission_label, entry_start = ekey

    # Find pairs: departure -> arrival within this entry
    # A departure has keywords: depart, departure, leave, outbound
    # An arrival has keywords: arriv, arrival, inbound, return

    for i, node in enumerate(nodes):
        nid, ntitle, ndetail, ntime, nloc, nst, *_ = node
        title_lower = (ntitle or "").lower()
        detail_lower = (ndetail or "").lower()
        ctx = f"{title_lower} {detail_lower}"

        new_st = None
        new_prev_loc = None
        new_transit_loc = None

        is_departure = any(w in ctx for w in
            ["departure", "depart", "leave", "outbound", "departed"])
        is_arrival = any(w in ctx for w in
            ["arrival", "arriv", "arrive", "inbound", "return", "soi"])

        if is_departure and is_arrival:
            # Ambiguous - use position context. If it mentions Earth or a body + departure, it's departure
            if "departure" in ctx or "depart" in ctx or "leave" in ctx:
                is_arrival = False
            else:
                is_departure = False

        if is_departure:
            new_st = "start-transit"
            new_transit_loc = "Solar Orbit"

        elif is_arrival:
            new_st = "transit-complete"

        # Also check existing state_label
        if nst and nst.lower() in ("transfer", "in transit"):
            if not new_st:
                # Fallback: try to infer from title
                if any(w in ctx for w in ["departure", "depart", "leave"]):
                    new_st = "start-transit"
                    new_transit_loc = "Solar Orbit"
                elif any(w in ctx for w in ["arrival", "arriv", "arrive"]):
                    new_st = "transit-complete"
                else:
                    # Keep as Transfer if we can't determine
                    new_st = "start-transit"
                    new_transit_loc = "Solar Orbit"

        if new_st:
            # Determine previous_location for start-transit
            if new_st == "start-transit":
                # Get previous location from:
                # a) Entry-level context (where is the asset before this mission?)
                # b) Explicit departure body
                body_match = re.search(
                    r'\b(earth|mars|europa|jupiter|saturn|venus|mercury|neptune|uranus|vesta|ceres|titan|io|ganymede|callisto|triton|pluto|leo|geo)\b',
                    ctx
                )
                if body_match:
                    body = body_match.group(1)
                    body_map = {
                        "earth": "LEO", "leo": "LEO", "geo": "GEO",
                        "mars": "Mars Orbit", "venus": "Venus Orbit",
                        "mercury": "Mercury Orbit", "jupiter": "Jupiter System",
                        "europa": "Europa Orbit", "io": "Io Orbit",
                        "ganymede": "Ganymede Orbit", "callisto": "Callisto Orbit",
                        "saturn": "Saturn System", "titan": "Titan Orbit",
                        "uranus": "Uranus System", "neptune": "Neptune System",
                        "triton": "Triton Orbit", "vesta": "Vesta Orbit",
                        "ceres": "Ceres Orbit", "pluto": "Pluto Orbit",
                    }
                    new_prev_loc = body_map.get(body, "LEO")
                else:
                    new_prev_loc = nloc or "LEO"

                # transit_location
                if not new_transit_loc:
                    new_transit_loc = "Solar Orbit"

            # For transit-complete, the location IS the arrival location
            if new_st == "transit-complete" and not nloc:
                # Try to infer arrival location from title
                body_match = re.search(
                    r'\b(earth|mars|europa|jupiter|saturn|venus|mercury|neptune|uranus|vesta|titan|io|ganymede|callisto|triton)\b',
                    ctx
                )
                if body_match:
                    body = body_match.group(1)
                    body_map = {
                        "earth": "LEO", "mars": "Mars Orbit", "venus": "Venus Orbit",
                        "mercury": "Mercury Orbit", "jupiter": "Jupiter System",
                        "europa": "Europa Orbit", "io": "Io Orbit",
                        "ganymede": "Ganymede Orbit", "callisto": "Callisto Orbit",
                        "saturn": "Saturn System", "titan": "Titan Orbit",
                        "uranus": "Uranus System", "neptune": "Neptune System",
                        "triton": "Triton Orbit", "vesta": "Vesta Orbit",
                    }
                    nloc = body_map.get(body, nloc)

            cur.execute("""
                UPDATE asset_state_nodes SET
                    state_label = %s,
                    previous_location = COALESCE(%s, previous_location),
                    transit_location = COALESCE(%s, transit_location),
                    location = COALESCE(%s, location),
                    updated_at = %s
                WHERE id = %s
            """, (new_st, new_prev_loc, new_transit_loc, nloc, NOW, nid))
            updated += 1

conn.commit()
print(f"Updated {updated} state nodes with new state labels")

# Stats
cur.execute("SELECT state_label, COUNT(*) FROM asset_state_nodes GROUP BY state_label ORDER BY COUNT(*) DESC")
print("\nState label distribution:")
for r in cur.fetchall():
    print(f"  {r[0] or 'NULL':25s}: {r[1]}")

cur.execute("SELECT COUNT(*), COUNT(CASE WHEN previous_location IS NOT NULL THEN 1 END), COUNT(CASE WHEN transit_location IS NOT NULL THEN 1 END) FROM asset_state_nodes")
r = cur.fetchone()
print(f"\nTotal: {r[0]}, with prev_loc: {r[1]}, with transit_loc: {r[2]}")
conn.close()