import csv
import math
from pathlib import Path
import re

# -----------------------
# CONFIG
# -----------------------
SCALE = 0.0254
FIXED_Y = 2.4
BUTTON_Y = 2.2
CONVEYOR_WIDTH = 1.524
EDGE_CLEARANCE = 0.45
BEAM_RANGE_ADJUSTMENT = 0.25

SCRIPT_DIR = Path(__file__).resolve().parent
csv_candidates = list(SCRIPT_DIR.glob("*.csv"))
if not csv_candidates:
    raise RuntimeError("No CSV found in script directory.")
if len(csv_candidates) > 1:
    raise RuntimeError(
        "Multiple CSV files found. Please keep only ONE CSV in the folder:\n"
        + "\n".join(c.name for c in csv_candidates)
    )
CSV_PATH = csv_candidates[0]

SCENE_DIR = SCRIPT_DIR.parent
OUTPUT_DIR = SCRIPT_DIR / "with_devices"
OUTPUT_DIR.mkdir(exist_ok=True)

SENSOR_SCENE_PATH = "res://parts/DiffuseSensor.tscn"
SENSOR_RES_ID = "auto_sensor"

# ✅ correct push button scene
BUTTON_SCENE_PATH = "res://parts/PushButton.tscn"
BUTTON_RES_ID = "auto_button"

# -----------------------
# LOAD CSV DEVICES
# -----------------------
tpe_devices = []
btn_devices = []
ss_devices = []
epc_devices = []

with open(CSV_PATH, newline="", encoding="utf-8") as f:
    reader = csv.DictReader(f)
    for row in reader:
        rt = (row.get("record_type") or "").strip().upper()
        if rt == "TPE":
            tpe_devices.append(row)
        elif rt == "S":
            btn_devices.append(row)
        elif rt == "SS":
            ss_devices.append(row)
        elif rt == "EPC":
            epc_devices.append(row)

if not tpe_devices:
    raise RuntimeError("No TPE records found in CSV.")

devices_by_conveyor = {}
for d in tpe_devices:
    key = (d.get("conveyor_key") or "").strip()
    if key:
        devices_by_conveyor.setdefault(key, []).append(d)

btn_by_conveyor = {}
for d in btn_devices:
    key = (d.get("conveyor_key") or "").strip()
    if key:
        btn_by_conveyor.setdefault(key, []).append(d)

ss_by_conveyor = {}
for d in ss_devices:
    key = (d.get("conveyor_key") or "").strip()
    if key:
        ss_by_conveyor.setdefault(key, []).append(d)

epc_by_conveyor = {}
for d in epc_devices:
    key = (d.get("conveyor_key") or "").strip()
    if key:
        epc_by_conveyor.setdefault(key, []).append(d)

# -----------------------
# SELECT SCENE
# -----------------------
scene_files = list(SCENE_DIR.glob("*.tscn"))
if not scene_files:
    raise RuntimeError("No .tscn files found in parent directory.")

print("\nAvailable scenes:")
for i, s in enumerate(scene_files, 1):
    print(f"[{i}] {s.name}")

choice = int(input("\nSelect scene number: ")) - 1
if choice < 0 or choice >= len(scene_files):
    raise RuntimeError("Invalid selection.")
SCENE_PATH = scene_files[choice]

# -----------------------
# DEVICE SELECTION MENU
# -----------------------
print("\nWhat devices do you want to place?")
print("[1] All")
print("[2] Sensors only")
print("[3] Buttons only (S)")
print("[4] SS only")
print("[5] EPC only")

device_choice = int(input("\nSelect option: "))
if device_choice < 1 or device_choice > 5:
    raise RuntimeError("Invalid selection.")

place_sensors = device_choice in (1, 2)
place_buttons = device_choice in (1, 3)
place_ss = device_choice in (1, 4)
place_epc = device_choice in (1, 5)

# -----------------------
# READ SCENE
# -----------------------
scene_text = SCENE_PATH.read_text(encoding="utf-8")
lines = scene_text.splitlines()

# -----------------------
# PARSE CONVEYORS
# -----------------------
conveyors = {}

node_name = None
has_size = False
pending_basis = None
pending_position = None
conveyor_width = CONVEYOR_WIDTH
conveyor_size = None

for line in lines:
    if line.startswith("[node name="):
        if node_name and pending_basis and pending_position:
            conveyors[node_name] = {
                "basis": pending_basis,
                "has_size": has_size,
                "width": conveyor_width,
                "position": pending_position,
                "size": conveyor_size
            }

        m = re.search(r'name="([^"]+)"', line)
        node_name = m.group(1) if m else None
        has_size = False
        pending_basis = None
        pending_position = None
        conveyor_width = CONVEYOR_WIDTH
        conveyor_size = None
        continue

    if node_name and line.strip().startswith("transform = Transform3D"):
        vals = [float(v) for v in re.search(r"\(([^)]+)\)", line).group(1).split(",")]
        pending_basis = (vals[0], vals[2], vals[6], vals[8])
        pending_position = (vals[9], vals[10], vals[11])

    if node_name and line.strip().startswith("size = Vector3"):
        has_size = True
        conveyor_size = tuple(
            float(v) for v in re.search(r"\(([^)]+)\)", line).group(1).split(",")
        )

    if node_name and line.strip().startswith("conveyor_width ="):
        conveyor_width = float(line.split("=")[1])

if node_name and pending_basis and pending_position:
    conveyors[node_name] = {
        "basis": pending_basis,
        "has_size": has_size,
        "width": conveyor_width,
        "position": pending_position,
        "size": conveyor_size
    }

# -----------------------
# ENSURE RESOURCES
# -----------------------
if place_sensors and SENSOR_SCENE_PATH not in scene_text:
    idx = max(i for i, l in enumerate(lines) if l.startswith("[ext_resource"))
    lines.insert(idx + 1,
        f'[ext_resource type="PackedScene" path="{SENSOR_SCENE_PATH}" id="{SENSOR_RES_ID}"]'
    )

if (place_buttons or place_ss or place_epc) and BUTTON_SCENE_PATH not in scene_text:
    idx = max(i for i, l in enumerate(lines) if l.startswith("[ext_resource"))
    lines.insert(idx + 1,
        f'[ext_resource type="PackedScene" path="{BUTTON_SCENE_PATH}" id="{BUTTON_RES_ID}"]'
    )

scene_text = "\n".join(lines) + "\n"

# -----------------------
# HELPERS
# -----------------------
def yaw_from_x_axis(xx, xz):
    return math.atan2(-xz, xx)

def yaw_from_z_axis(zx, zz):
    return math.atan2(zx, zz)

def compute_sensor_position(info, yaw, idx, total):
    px, _, pz = info["position"]

    if not info["size"]:
        return None, None

    length = info["size"][0]

    if total == 1:
        off = length / 2
    elif total == 2:
        off = (-length / 2) if idx == 0 else (length / 2)
    else:
        return None, None

    return (
        px + math.cos(yaw) * off,
        pz + math.sin(yaw) * off
    )

def godot_yaw_from_cad(cad_deg: float) -> float:
    return math.radians((270 - cad_deg) % 360)

def compute_edge_offset(sensor_yaw, width):
    side = sensor_yaw - math.pi / 2
    d = (width / 2) + EDGE_CLEARANCE
    return math.cos(side) * d, math.sin(side) * d

# --- BUTTON HELPERS (UNCHANGED) ---
def side_of_conveyor(info, conveyor_yaw, csv_x, csv_z):
    cx, _, cz = info["position"]
    rx = math.cos(conveyor_yaw - math.pi / 2)
    rz = math.sin(conveyor_yaw - math.pi / 2)
    vx = csv_x - cx
    vz = csv_z - cz
    return vx * rx + vz * rz

def project_to_conveyor(info, yaw, csv_x, csv_z):
    cx, _, cz = info["position"]
    fx = math.cos(yaw)
    fz = math.sin(yaw)
    vx = csv_x - cx
    vz = csv_z - cz
    t = vx * fx + vz * fz
    return cx + fx * t, cz + fz * t

# -----------------------
# PARSE EXISTING BUTTONS FROM SCENE
# -----------------------
existing_buttons = []

in_button_node = False
current_button_name = None
current_button_transform = None

for line in lines:
    # Detect button nodes
    if f'instance=ExtResource("{BUTTON_RES_ID}")' in line:
        in_button_node = True
        m = re.search(r'name="([^"]+)"', line)
        if m:
            current_button_name = m.group(1)
    
    # Extract transform if we're in a button node
    if in_button_node and line.strip().startswith("transform = Transform3D"):
        vals = [float(v) for v in re.search(r"\(([^)]+)\)", line).group(1).split(",")]
        current_button_transform = (vals[9], vals[11])  # x, z position
        
        if current_button_name and current_button_transform:
            existing_buttons.append({
                'name': current_button_name,
                'x': current_button_transform[0],
                'z': current_button_transform[1]
            })
        
        in_button_node = False
        current_button_name = None
        current_button_transform = None

print(f"\n📍 Found {len(existing_buttons)} existing buttons in scene")

# -----------------------
# COLLISION DETECTION HELPER
# -----------------------
MIN_BUTTON_SPACING = 0.5  # minimum distance between buttons (meters)

def check_collision_and_adjust(x, z, yaw, existing_buttons):
    """
    Check if position (x, z) collides with existing buttons.
    If collision detected, shift along conveyor direction (yaw).
    Returns adjusted (x, z) position.
    """
    max_attempts = 20
    shift_distance = 0.3  # how much to shift each attempt
    
    for attempt in range(max_attempts):
        collision = False
        
        for btn in existing_buttons:
            dist = math.sqrt((x - btn['x'])**2 + (z - btn['z'])**2)
            if dist < MIN_BUTTON_SPACING:
                collision = True
                break
        
        if not collision:
            return x, z
        
        # Shift forward along conveyor
        x += math.cos(yaw) * shift_distance
        z += math.sin(yaw) * shift_distance
    
    print(f"⚠ Warning: Could not find collision-free position after {max_attempts} attempts")
    return x, z

# -----------------------
# APPEND DEVICES
# -----------------------
node_blocks = []

# -------- SENSORS --------
if place_sensors:
    for key in sorted(devices_by_conveyor.keys()):
        devs = devices_by_conveyor[key]

        if key not in conveyors:
            print(f"⚠ Conveyor not found: {key}")
            continue

        info = conveyors[key]
        
        xx, xz, zx, zz = info["basis"]
        width = info["width"]

        yaw = yaw_from_x_axis(xx, xz) if info["has_size"] else yaw_from_z_axis(zx, zz)

        for i, d in enumerate(devs):
            name = f"{key}_TPE{i+1}"
            tag = f"{name}_OIP"

            max_range = width + BEAM_RANGE_ADJUSTMENT

            x, z = compute_sensor_position(info, yaw, i, len(devs))

            LONGITUDINAL_OFFSET = 0.15

            if x is not None:
                if len(devs) == 1:
                    x -= math.cos(yaw) * LONGITUDINAL_OFFSET
                    z -= math.sin(yaw) * LONGITUDINAL_OFFSET
                elif len(devs) == 2:
                    if i == 0:
                        x += math.cos(yaw) * LONGITUDINAL_OFFSET
                        z += math.sin(yaw) * LONGITUDINAL_OFFSET
                    else:
                        x -= math.cos(yaw) * LONGITUDINAL_OFFSET
                        z -= math.sin(yaw) * LONGITUDINAL_OFFSET
            else:
                x = float(d["tpe_x"]) * SCALE
                z = -float(d["tpe_y"]) * SCALE

            if "tpe_rotation" in d and d["tpe_rotation"] not in ("", None):
                tpe_rot = float(d["tpe_rotation"])
            else:
                tpe_rot = 0.0

            if "tpe_block_rotation" in d and d["tpe_block_rotation"] not in ("", None):
                blk_rot = float(d["tpe_block_rotation"])
            else:
                blk_rot = 0.0

            if tpe_rot != 0.0 or blk_rot != 0.0:
                cad_world_rot = (blk_rot + tpe_rot) % 360.0
                sensor_yaw = godot_yaw_from_cad(cad_world_rot)
            else:
                sensor_yaw = yaw + math.pi

            ox, oz = compute_edge_offset(sensor_yaw, width)
            x += ox
            z += oz

            c, s = math.cos(sensor_yaw), math.sin(sensor_yaw)

            transform = (
                f"Transform3D({c:.6f},0,{-s:.6f},0,1,0,{s:.6f},0,{c:.6f},"
                f"{x:.6f},{FIXED_Y:.6f},{z:.6f})"
            )

            node_blocks.append(
                f'\n[node name="{name}" parent="." instance=ExtResource("{SENSOR_RES_ID}")]\n'
                f"transform = {transform}\n"
                f"max_range = {max_range:.3f}\n"
                f"enable_comms = true\n"
                f'tag_name = "{tag}"\n'
            )

# -------- BUTTONS (S) --------
if place_buttons:
    for key in sorted(btn_by_conveyor.keys()):
        devs = btn_by_conveyor[key]

        def btn_index(d):
            name = (d.get("dev_name") or "")
            m = re.search(r"_S(\d+)$", name)
            return int(m.group(1)) if m else 999

        devs = sorted(devs, key=btn_index)

        if key not in conveyors:
            print(f"⚠ Conveyor not found for buttons: {key}")
            continue

        info = conveyors[key]
        xx, xz, zx, zz = info["basis"]
        width = info["width"]

        yaw = yaw_from_x_axis(xx, xz) if info["has_size"] else yaw_from_z_axis(zx, zz)

        for i, d in enumerate(devs):
            name = f"{key}_S{i+1}"

            csv_x = float(d["dev_x"]) * SCALE
            csv_z = -float(d["dev_y"]) * SCALE

            x, z = project_to_conveyor(info, yaw, csv_x, csv_z)
            side = side_of_conveyor(info, yaw, csv_x, csv_z)

            if abs(side) < 0.05:
                side_sign = 0.0
            else:
                side_sign = 1.0 if side > 0 else -1.0

            px = math.cos(yaw - math.pi / 2)
            pz = math.sin(yaw - math.pi / 2)

            BUTTON_CLEARANCE = 0.1
            d_off = (width / 2) + BUTTON_CLEARANCE
            x += px * d_off * side_sign
            z += pz * d_off * side_sign

            if side_sign >= 0:
                btn_yaw = yaw + math.pi
            else:
                btn_yaw = yaw

            # Check for collisions with existing buttons
            x, z = check_collision_and_adjust(x, z, yaw, existing_buttons)

            c, s = math.cos(btn_yaw), math.sin(btn_yaw)

            transform = (
                f"Transform3D({c:.6f},0,{-s:.6f},0,1,0,{s:.6f},0,{c:.6f},"
                f"{x:.6f},{BUTTON_Y:.6f},{z:.6f})"
            )

            node_blocks.append(
                f'\n[node name="{name}" parent="." instance=ExtResource("{BUTTON_RES_ID}")]\n'
                f"transform = {transform}\n"
                f'text = "START"\n'
                f"button_color = Color(0.39, 1.0, 0.098, 1)\n"
                f"enable_comms = true\n"
                f'pushbutton_tag_name = "{name}_PB_OIP"\n'
                f'lamp_tag_name = "{name}_LT_OIP"\n'
            )
            
            # Add to existing buttons list to prevent overlap with other new buttons
            existing_buttons.append({'name': name, 'x': x, 'z': z})

# -------- SS STATIONS --------
if place_ss:
    for key in sorted(ss_by_conveyor.keys()):
        devs = ss_by_conveyor[key]

        # sort SS1, SS2...
        def ss_index(d):
            name = (d.get("dev_name") or "")
            m = re.search(r"_SS(\d+)$", name)
            return int(m.group(1)) if m else 999

        devs = sorted(devs, key=ss_index)

        if key not in conveyors:
            print(f"⚠ Conveyor not found for SS station: {key}")
            continue

        info = conveyors[key]
        xx, xz, zx, zz = info["basis"]
        width = info["width"]

        yaw = yaw_from_x_axis(xx, xz) if info["has_size"] else yaw_from_z_axis(zx, zz)

        for i, d in enumerate(devs):
            station = f"{key}_SS{i+1}"

            csv_x = float(d["dev_x"]) * SCALE
            csv_z = -float(d["dev_y"]) * SCALE

            # Get conveyor center position
            cx, _, cz = info["position"]
            
            # Check if SS is to the left of conveyor (lower X coordinate)
            if csv_x < cx:
                # Place at START edge CENTER of conveyor, facing opposite to flow
                if not info["size"]:
                    x = cx
                    z = cz
                else:
                    length = info["size"][0]
                    # Position at start edge center (negative direction along conveyor)
                    x = cx - math.cos(yaw) * (length / 2)
                    z = cz - math.sin(yaw) * (length / 2)
                
                # Rotation: 270 degrees (-90), perpendicular to flow, pointing left
                # Calculate perpendicular based on conveyor's actual basis vectors
                # Right perpendicular uses the conveyor's Z-axis direction
                perp_yaw = math.atan2(zx, zz)  # perpendicular right
                base_yaw = perp_yaw + math.pi / 2  # rotate 90° to get 270° (-90°)
                
                # Spread buttons along the perpendicular (left-right across conveyor width)
                fx = math.cos(yaw - math.pi / 2)  # perpendicular to conveyor flow
                fz = math.sin(yaw - math.pi / 2)
                STATION_SPREAD = 0.45
            else:
                # Original logic: project to conveyor and place on side
                x, z = project_to_conveyor(info, yaw, csv_x, csv_z)
                side = side_of_conveyor(info, yaw, csv_x, csv_z)

                if abs(side) < 0.05:
                    side_sign = 0.0
                else:
                    side_sign = 1.0 if side > 0 else -1.0

                px = math.cos(yaw - math.pi / 2)
                pz = math.sin(yaw - math.pi / 2)

                BUTTON_CLEARANCE = 0.1
                d_off = (width / 2) + BUTTON_CLEARANCE
                x += px * d_off * side_sign
                z += pz * d_off * side_sign

                if side_sign >= 0:
                    base_yaw = yaw + math.pi
                else:
                    base_yaw = yaw

                # Spread buttons along conveyor direction
                fx = math.cos(yaw)
                fz = math.sin(yaw)
                STATION_SPREAD = 0.45

            # Check for collisions with existing buttons
            x, z = check_collision_and_adjust(x, z, yaw, existing_buttons)

            # -------- START (SPB) --------
            start_node = f"{station}_SPB"
            start_x = x + fx * (STATION_SPREAD / 2)
            start_z = z + fz * (STATION_SPREAD / 2)

            c, s = math.cos(base_yaw), math.sin(base_yaw)
            start_transform = (
                f"Transform3D({c:.6f},0,{-s:.6f},0,1,0,{s:.6f},0,{c:.6f},"
                f"{start_x:.6f},{BUTTON_Y:.6f},{start_z:.6f})"
            )

            node_blocks.append(
                f'\n[node name="{start_node}" parent="." instance=ExtResource("{BUTTON_RES_ID}")]\n'
                f"transform = {start_transform}\n"
                f'text = "START"\n'
                f"button_color = Color(0.39, 1.0, 0.098, 1)\n"
                f"enable_comms = true\n"
                f'pushbutton_tag_name = "{station}_SPB_OIP"\n'
                f'lamp_tag_name = "{station}_SPB_LT_OIP"\n'
            )
            
            # Add to existing buttons list
            existing_buttons.append({'name': start_node, 'x': start_x, 'z': start_z})

            # -------- STOP (STPB) --------
            stop_node = f"{station}_STPB"
            stop_x = x - fx * (STATION_SPREAD / 2)
            stop_z = z - fz * (STATION_SPREAD / 2)

            stop_transform = (
                f"Transform3D({c:.6f},0,{-s:.6f},0,1,0,{s:.6f},0,{c:.6f},"
                f"{stop_x:.6f},{BUTTON_Y:.6f},{stop_z:.6f})"
            )

            node_blocks.append(
                f'\n[node name="{stop_node}" parent="." instance=ExtResource("{BUTTON_RES_ID}")]\n'
                f"transform = {stop_transform}\n"
                f'text = "STOP"\n'
                f"button_color = Color(1.0, 0.15, 0.15, 1)\n"
                f"normally_closed = true\n"
                f"enable_comms = true\n"
                f'pushbutton_tag_name = "{station}_STPB_OIP"\n'
            )
            
            # Add to existing buttons list
            existing_buttons.append({'name': stop_node, 'x': stop_x, 'z': stop_z})

# -------- EPC DEVICES --------
if place_epc:
    for key in sorted(epc_by_conveyor.keys()):
        devs = epc_by_conveyor[key]

        if key not in conveyors:
            print(f"⚠ Conveyor not found for EPC: {key}")
            continue

        info = conveyors[key]
        xx, xz, zx, zz = info["basis"]
        width = info["width"]

        yaw = yaw_from_x_axis(xx, xz) if info["has_size"] else yaw_from_z_axis(zx, zz)

        for i, d in enumerate(devs):
            name = d.get("dev_name") or f"{key}_EPC{i+1}"

            csv_x = float(d["dev_x"]) * SCALE
            csv_z = -float(d["dev_y"]) * SCALE

            # Same side logic as S buttons
            x, z = project_to_conveyor(info, yaw, csv_x, csv_z)
            side = side_of_conveyor(info, yaw, csv_x, csv_z)

            side_sign = 1.0 if side >= 0 else -1.0

            px = math.cos(yaw - math.pi / 2)
            pz = math.sin(yaw - math.pi / 2)

            EPC_CLEARANCE = 0.1
            d_off = (width / 2) + EPC_CLEARANCE
            x += px * d_off * side_sign
            z += pz * d_off * side_sign

            if side_sign >= 0:
                base_yaw = yaw + math.pi
            else:
                base_yaw = yaw

            # Spread along conveyor direction (like SS buttons do)
            EPC_SPREAD = 0.35
            fx = math.cos(yaw)
            fz = math.sin(yaw)

            # Check for collisions with existing buttons
            x, z = check_collision_and_adjust(x, z, yaw, existing_buttons)

            c, s = math.cos(base_yaw), math.sin(base_yaw)

            # -------- CH1 --------
            ch1_name = f"{name}_CH1"
            ch1_x = x + fx * (EPC_SPREAD / 2)
            ch1_z = z + fz * (EPC_SPREAD / 2)

            ch1_transform = (
                f"Transform3D({c:.6f},0,{-s:.6f},0,1,0,{s:.6f},0,{c:.6f},"
                f"{ch1_x:.6f},{BUTTON_Y:.6f},{ch1_z:.6f})"
            )

            node_blocks.append(
                f'\n[node name="{ch1_name}" parent="." instance=ExtResource("{BUTTON_RES_ID}")]\n'
                f"transform = {ch1_transform}\n"
                f'text = "EPC"\n'
                f"button_color = Color(1, 0.15, 0.15, 1)\n"
                f"normally_closed = true\n"
                f"enable_comms = true\n"
                f'pushbutton_tag_name = "{name}_CH1_OIP"\n'
            )
            
            # Add to existing buttons list
            existing_buttons.append({'name': ch1_name, 'x': ch1_x, 'z': ch1_z})

            # -------- CH2 --------
            ch2_name = f"{name}_CH2"
            ch2_x = x - fx * (EPC_SPREAD / 2)
            ch2_z = z - fz * (EPC_SPREAD / 2)

            ch2_transform = (
                f"Transform3D({c:.6f},0,{-s:.6f},0,1,0,{s:.6f},0,{c:.6f},"
                f"{ch2_x:.6f},{BUTTON_Y:.6f},{ch2_z:.6f})"
            )

            node_blocks.append(
                f'\n[node name="{ch2_name}" parent="." instance=ExtResource("{BUTTON_RES_ID}")]\n'
                f"transform = {ch2_transform}\n"
                f'text = "EPC"\n'
                f"button_color = Color(1, 0.15, 0.15, 1)\n"
                f"normally_closed = true\n"
                f"enable_comms = true\n"
                f'pushbutton_tag_name = "{name}_CH2_OIP"\n'
            )
            
            # Add to existing buttons list
            existing_buttons.append({'name': ch2_name, 'x': ch2_x, 'z': ch2_z})


scene_text += "".join(node_blocks)

# -----------------------
# WRITE OUTPUT
# -----------------------
out_path = OUTPUT_DIR / f"{SCENE_PATH.stem}_devices{SCENE_PATH.suffix}"
out_path.write_text(scene_text, encoding="utf-8")

print(f"\n✔ {len(node_blocks)} devices added → {out_path}")