#!/usr/bin/env python3
import os
import yaml
import sys
import xml.etree.ElementTree as ET

SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
PROJECT_DIR = os.path.dirname(SCRIPT_DIR)
CONFIG_DIR = os.path.join(PROJECT_DIR, "config")
WORLD_DIR = os.path.join(PROJECT_DIR, "worlds")

def load_config(name):
    path = os.path.join(CONFIG_DIR, name)
    if not os.path.exists(path):
        return {}
    with open(path, 'r') as f:
        return yaml.safe_load(f) or {}

def generate_world(base_filename="uav_ugv_search_base.sdf", output_filename="uav_ugv_search.sdf"):
    ugv_cfg = load_config("ugv.yaml")
    search_cfg = load_config("search.yaml")

    ugv_x = ugv_cfg.get("position", {}).get("x", 0.0)
    ugv_y = ugv_cfg.get("position", {}).get("y", 0.0)

    uav_x = ugv_x
    uav_y = ugv_y
    uav_z = 0.40

    marker = search_cfg.get("marker", {})
    target_pos = marker.get("target_position", [8.0, -6.0])
    target_x = target_pos[0]
    target_y = target_pos[1]

    base_path = os.path.join(WORLD_DIR, base_filename)
    output_path = os.path.join(WORLD_DIR, output_filename)

    if not os.path.exists(base_path):
        print(f"[ERROR] Base world file not found: {base_path}")
        sys.exit(1)

    try:
        tree = ET.parse(base_path)
        root = tree.getroot()

        world = root.find('world')
        if world is None:
            print("[ERROR] No <world> tag found in SDF.")
            sys.exit(1)

        for include in world.findall('include'):
            uri = include.find('uri')
            if uri is not None and uri.text == 'model://iris_with_gimbal':
                pose = include.find('pose')
                if pose is not None:
                    pose.text = f"{uav_x} {uav_y} {uav_z} 0 0 90"

            name = include.find('name')
            if name is not None and name.text == 'ugv':
                pose = include.find('pose')
                if pose is not None:
                    pose.text = f"{ugv_x} {ugv_y} 0 0 0 0"

        for model in world.findall('model'):
            name_attr = model.get('name')
            if name_attr == 'target_tag_1':
                pose = model.find('pose')
                if pose is not None:
                    pose.text = f"{target_x} {target_y} 0.005 0 0 0"

        geofence_cfg = search_cfg.get("geofence", {})
        if geofence_cfg.get("enabled", False):
            points = geofence_cfg.get("points", [])
            if len(points) >= 3:

                for old_model in world.findall('model'):
                    if old_model.get('name') == 'geofence_visual':
                        world.remove(old_model)

                gf_model = ET.SubElement(world, "model", name="geofence_visual")
                ET.SubElement(gf_model, "static").text = "true"
                link = ET.SubElement(gf_model, "link", name="link")

                import math
                for i in range(len(points)):
                    p1 = points[i]
                    p2 = points[(i + 1) % len(points)]
                    x1, y1 = float(p1[0]), float(p1[1])
                    x2, y2 = float(p2[0]), float(p2[1])

                    dx = x2 - x1
                    dy = y2 - y1
                    length = math.sqrt(dx*dx + dy*dy)
                    cx = x1 + dx / 2.0
                    cy = y1 + dy / 2.0
                    yaw = math.atan2(dy, dx)

                    visual = ET.SubElement(link, "visual", name=f"edge_{i}")
                    ET.SubElement(visual, "pose").text = f"{cx} {cy} 0.01 0 0 {yaw}"

                    geometry = ET.SubElement(visual, "geometry")
                    box = ET.SubElement(geometry, "box")

                    ET.SubElement(box, "size").text = f"{length} 0.2 0.02"

                    material = ET.SubElement(visual, "material")
                    ET.SubElement(material, "ambient").text = "1 0 0 1"
                    ET.SubElement(material, "diffuse").text = "1 0 0 1"
                    ET.SubElement(material, "emissive").text = "0.8 0 0 0.5"

        tree.write(output_path, encoding="utf-8", xml_declaration=True)
        print(f"[INFO] Generated world file: {output_path}")
        print(f"[INFO] UGV set to ({ugv_x}, {ugv_y})")
        print(f"[INFO] Target Marker set to ({target_x}, {target_y})")

    except Exception as e:
        print(f"[ERROR] Failed to parse or write XML: {e}")
        sys.exit(1)

if __name__ == "__main__":
    generate_world()