#!/usr/bin/env -S uv run --script
#
# /// script
# requires-python = ">=3.12"
# dependencies = [
#     "pyyaml",
#     "scipy",
# ]
# ///

import json
import os
import re
import sys
import traceback
from scipy.spatial.transform import Rotation
from pathlib import Path

import yaml
#
#    [ TD2 TRACK STRUCTURE PREFAB PARSER ]
#    by dkgl
#    v1.0
#
#   The assets are exported using https://github.com/AssetRipper/AssetRipper in the "Export Unity Project" mode.
#   Pass the path to the "(...)/ExportedProject/Assets/Resources/track structures" directory as a command line argument to this script.
#
#   This script uses stdout for the output and stderr for log and error messages.
#

def eprint(*args, **kwargs):
    print(*args, file=sys.stderr, **kwargs)


def parse_prefab(file, prefab_name):
    main_component = parse_prefab_components(file)
    tracks = find_tracks(main_component)
    isolation_label_pos = get_isolation_label_pos(prefab_name)
    return {
        "name": prefab_name,
        "tracks": tracks,
        "isolation_label_pos": isolation_label_pos,
    }


prefab_header_pattern = re.compile(r"^--- !u!(\d+) &(\d+)$")


def parse_prefab_components(file):
    component_list = []
    component_map = {}
    current_yaml = ""
    component_type = None
    component_id = None

    for line in file:
        match = prefab_header_pattern.match(line)
        if match:
            if component_type is not None:
                component = parse_prefab_component(current_yaml, component_type, component_id)
                if component is not None:
                    component_list.append(component)
                    component_map[component_id] = component
            current_yaml = ""
            component_type = match.group(1)
            component_id = match.group(2)
        else:
            current_yaml += line

    if component_type is not None:
        component = parse_prefab_component(current_yaml, component_type, component_id)
        if component is not None:
            component_list.append(component)
            component_map[component_id] = component

    for component in component_list:
        if not isinstance(component, dict):
            continue
        for value in component.values():
            if "m_GameObject" not in value:
                continue
            component_map[str(value["m_GameObject"]["fileID"])]["__children"].append(component)

    visited_ids = set()
    for component in component_list:
        resolve_references(component, component_map, visited_ids)

    return component_list[0]


def try_parse_switch_component(prefab):
    if "component" not in prefab:
        return None
    component = prefab["component"]
    if "MonoBehaviour" not in component:
        return None
    mono = component["MonoBehaviour"]
    if "tracks" not in mono or "namedSwitch" not in mono:
        return None
    return mono


def try_parse_track_shape(component):
    if "MonoBehaviour" not in component:
        return None
    mono = component["MonoBehaviour"]
    if "slope1" not in mono:
        return None
    return {
        "radius": mono["radius"],
        "length": mono["length"],
        "rz1": mono["rz1"],
        "rz2": mono["rz2"],
        "slope1": mono["slope1"],
        "slope2": mono["slope2"],
    }


def mirror_rotation_x(rotation):
    rot_vet = rotation.as_rotvec()
    # Flip the x-axis to apply the mirroring and then negate the whole vector to change the rotation direction
    return Rotation.from_rotvec([rot_vet[0], -rot_vet[1], -rot_vet[2]])


def try_parse_track_transform(component):
    if "Transform" not in component:
        return None
    transform = component["Transform"]
    if "Transform" in transform["m_Father"]:
        parent = try_parse_track_transform(transform["m_Father"])
    else:
        parent = {
            "position": [0, 0, 0],
            "rotation": Rotation.identity(),
            "mirror_x": False,
            "negate_radius": False,
        }

    if transform["m_LocalScale"]["x"] not in {-1, 1} or transform["m_LocalScale"]["y"] != 1 or transform["m_LocalScale"]["z"] != 1:
        eprint("Unexpected scale in track transform")

    position = [transform["m_LocalPosition"][key] for key in ["x", "y", "z"]]
    rotation = Rotation.from_quat([transform["m_LocalRotation"][key] for key in ["x", "y", "z", "w"]])
    mirror_x = transform["m_LocalScale"]["x"] == -1
    negate_radius = mirror_x

    # Apply own scaleX transform to rotation, then the parent's
    if mirror_x ^ parent["mirror_x"]:
        rotation = mirror_rotation_x(rotation)

    # Apply remaining parent transforms
    if parent["mirror_x"]:
        position[0] = -position[0]
    position = parent["position"] + parent["rotation"].apply(position)
    rotation = parent["rotation"] * rotation
    if parent["negate_radius"]:
        negate_radius = not negate_radius

    return {
        "position": position,
        "rotation": rotation,
        "mirror_x": mirror_x,
        "negate_radius": negate_radius,
    }


def parse_track(track, data_index):
    track_id = track["__component_id"]
    track_shape = None
    track_transform = None

    game_object = track["MonoBehaviour"]["m_GameObject"]["GameObject"]

    for child in game_object["m_Component"]:
        if "component" not in child:
            continue
        child = child["component"]
        maybe_track_shape = try_parse_track_shape(child)
        if maybe_track_shape is not None:
            track_shape = maybe_track_shape
        maybe_track_transform = try_parse_track_transform(child)
        if maybe_track_transform is not None:
            track_transform = maybe_track_transform

    connections = []

    def append_connection(end, field):
        if "__component_id" in track["MonoBehaviour"][field]:
            other_id = track["MonoBehaviour"][field]["__component_id"]
            if other_id == track_id:
                # Ignore connections to self, TD2 has them in Rkp switches
                return
            connections.append({
                "type": "<INTERNAL>",
                "end": end,
                "internalId": other_id
            })
        else:
            connections.append({
                "type": "<EXTERNAL>",
                "end": end,
            })

    append_connection("<START>", "prevTrack")
    append_connection("<END>", "nextTrack")

    if track_transform["negate_radius"]:
        track_shape["radius"] = -track_shape["radius"]

    return (track_id, {
        "dataIndex": data_index,
        **track_shape,
        "pos": f"<new Vector3({", ".join([str(x) for x in track_transform["position"].tolist()])})>",
        "rot": f"<new Quaternion({", ".join([str(x) for x in track_transform["rotation"].as_quat().tolist()])})>",
        "connections": connections,
    })


def find_tracks(prefab):
    tracks = []
    for component in prefab["GameObject"]["m_Component"]:
        switch = try_parse_switch_component(component)
        if switch is None:
            continue
        for data_index, track in enumerate(switch["tracks"]):
            tracks.append(parse_track(track, data_index))
    add_missing_connections(tracks)
    tracks = change_ids(tracks)
    return tracks

def add_missing_connections(tracks):
    track_map = dict(tracks)
    for (track_id, track) in tracks:
        for connection in track["connections"]:
            if connection["type"] == "<EXTERNAL>":
                continue
            other_track = track_map.get(connection["internalId"], None)
            if other_track is None:
                eprint(f"Warning: Track {track_id} has a connection to an unknown track {connection['internalId']}")
                continue
            reverse_connections = [conn for conn in other_track["connections"] if conn["type"] == "<INTERNAL>" and conn["internalId"] == track_id]
            if len(reverse_connections) > 1:
                eprint(f"Warning: Track {other_track['id']} has multiple connections to {track_id}")
            if len(reverse_connections) == 0:
                other_track["connections"].append({
                    "type": "<INTERNAL>",
                    "end": "<START>" if connection["end"] == "<END>" else "<END>", # Heuristic
                    "internalId": track_id
                })
    return tracks

def change_ids(tracks):
    new_id_map = dict()
    new_tracks = []
    for (track_id, track) in tracks:
        new_id = f"I{track["dataIndex"]}"
        new_id_map[track_id] = new_id
        new_tracks.append((new_id, track))
    for (new_id, track) in new_tracks:
        for connection in track["connections"]:
            if connection["type"] == "<INTERNAL>":
                connection["internalId"] = new_id_map[connection["internalId"]]
    return new_tracks

def parse_prefab_component(component_yaml, component_type, component_id):
    loaded = yaml.load(component_yaml, yaml.SafeLoader)
    loaded["__component_id"] = component_id
    loaded["__component_type"] = component_type
    loaded["__children"] = []
    return loaded


def resolve_references(value, component_map, visited_ids):
    if id(value) in visited_ids:
        return
    visited_ids.add(id(value))
    if isinstance(value, dict):
        keys = value.keys()
    elif isinstance(value, list):
        keys = range(len(value))
    else:
        return
    for key in keys:
        if isinstance(value[key], dict) and "fileID" in value[key]:
            file_id = str(value[key]["fileID"])
            if file_id in component_map:
                component = component_map[file_id]
                value[key] = component

        resolve_references(value[key], component_map, visited_ids)


def get_isolation_label_pos(prefab_name):
    if prefab_name.startswith("Rkp") or prefab_name.startswith("Crossing"):
        return "new Vector3(0, 0, 0)"
    return "new Vector3(0, 0, 6)"


def format_prefab(prefab):
    print(f"'{prefab["name"]}': {{")
    print("    tracks: {")
    for (track_id, track) in prefab["tracks"]:
        track_object = json.dumps(track)
        # Replace "<FOO>" with FOO
        track_object = re.sub(r'(?<!\\)"<([^<>"]+)>"', r'\1', track_object)
        # Remove quotes around keys
        track_object = re.sub(r'(?<!\\)"([a-zA-Z][a-zA-Z0-9]*)":', r'\1:', track_object)
        print(f"        '{track_id}': {track_object},")
    print("    },")
    print(f"    isolationLabelPos: {prefab["isolation_label_pos"]},")
    print("},")


def process_file(file_path):
    if file_path.endswith(".prefab.meta"):
        return

    if not file_path.endswith(".prefab"):
        eprint(f"Skipping non-.prefab file: {file_path}", file=sys.stderr)
        return

    eprint(f"Processing file: {file_path}")
    try:
        with open(file_path, encoding="utf8") as infile:
            prefab_name = Path(file_path).stem
            prefab = parse_prefab(infile, prefab_name)
            format_prefab(prefab)
    except Exception:
        eprint(f"Error processing {file_path}:", file=sys.stderr)
        traceback.print_exc(file=sys.stderr)


def process_directory(directory: str):
    eprint(f"Processing directory: {directory}")
    for entry in sorted(os.listdir(directory)):
        full_path = os.path.join(directory, entry)
        if os.path.isfile(full_path):
            process_file(full_path)


def main():
    if len(sys.argv) < 2:
        eprint("Usage: python prase-prefabs.py <file_or_directory> [<file_or_directory> ...]")
        return

    for path in sys.argv[1:]:
        if os.path.isdir(path):
            process_directory(path)
        elif os.path.isfile(path):
            process_file(path)
        else:
            eprint(f"Invalid path: {path}")

    eprint("Done!")


if __name__ == "__main__":
    main()