# Install ROS Packages for Unity Project
# 为Unity项目安装ROS包

param(
    [string]$ProjectPath = "D:\ros2\DockerRos2Arm-niryo_arm\Unity_MoveIt2\NewUnityProject\ROS2_Niryo_Control",
    [string]$UnityPath = "C:\Program Files\Unity\Hub\Editor\2021.3.33f1c1\Editor\Unity.exe",
    [switch]$InstallTCPConnector,
    [switch]$InstallURDFImporter,
    [switch]$InstallAll
)

Write-Host "=== Installing ROS Packages for Unity Project ===" -ForegroundColor Green

# Check if Unity project exists
if (-not (Test-Path $ProjectPath)) {
    Write-Host "Error: Unity project not found at $ProjectPath" -ForegroundColor Red
    exit 1
}

# Check if Unity exists
if (-not (Test-Path $UnityPath)) {
    Write-Host "Error: Unity not found at $UnityPath" -ForegroundColor Red
    exit 1
}

# Function to add package to manifest.json
function Add-PackageToManifest {
    param(
        [string]$ManifestPath,
        [string]$PackageName,
        [string]$PackageVersion,
        [string]$PackageUrl = $null
    )
    
    if (-not (Test-Path $ManifestPath)) {
        Write-Host "Error: manifest.json not found at $ManifestPath" -ForegroundColor Red
        return $false
    }
    
    try {
        $manifest = Get-Content $ManifestPath -Raw | ConvertFrom-Json
        
        if ($PackageUrl) {
            # Add Git URL package
            $manifest.dependencies | Add-Member -MemberType NoteProperty -Name $PackageName -Value $PackageUrl -Force
        } else {
            # Add registry package
            $manifest.dependencies | Add-Member -MemberType NoteProperty -Name $PackageName -Value $PackageVersion -Force
        }
        
        $manifest | ConvertTo-Json -Depth 10 | Set-Content $ManifestPath -Encoding UTF8
        Write-Host "Added package: $PackageName" -ForegroundColor Green
        return $true
    }
    catch {
        Write-Host "Error updating manifest.json: $($_.Exception.Message)" -ForegroundColor Red
        return $false
    }
}

$manifestPath = Join-Path $ProjectPath "Packages\manifest.json"

# Install ROS TCP Connector
if ($InstallTCPConnector -or $InstallAll) {
    Write-Host "Installing ROS TCP Connector..." -ForegroundColor Yellow
    $success = Add-PackageToManifest -ManifestPath $manifestPath -PackageName "com.unity.robotics.ros-tcp-connector" -PackageUrl "https://github.com/Unity-Technologies/ROS-TCP-Connector.git?path=/com.unity.robotics.ros-tcp-connector"
    
    if ($success) {
        Write-Host "ROS TCP Connector added to manifest.json" -ForegroundColor Green
    }
}

# Install URDF Importer
if ($InstallURDFImporter -or $InstallAll) {
    Write-Host "Installing URDF Importer..." -ForegroundColor Yellow
    $success = Add-PackageToManifest -ManifestPath $manifestPath -PackageName "com.unity.robotics.urdf-importer" -PackageUrl "https://github.com/Unity-Technologies/URDF-Importer.git?path=/com.unity.robotics.urdf-importer"
    
    if ($success) {
        Write-Host "URDF Importer added to manifest.json" -ForegroundColor Green
    }
}

# Create ROS Messages directory structure
Write-Host "Creating ROS Messages structure..." -ForegroundColor Yellow
$rosMessagesPath = Join-Path $ProjectPath "Assets\RosMessages"
if (-not (Test-Path $rosMessagesPath)) {
    New-Item -ItemType Directory -Path $rosMessagesPath -Force | Out-Null
}

# Create basic ROS connection script
$rosConnectionScript = @"
using System.Collections;
using UnityEngine;
using Unity.Robotics.ROSTCPConnector;
using RosMessageTypes.Std;

namespace ROS2NiryoControl
{
    public class ROSConnection : MonoBehaviour
    {
        [Header("ROS Settings")]
        public string rosIPAddress = "localhost";
        public int rosPort = 10000;
        
        [Header("Connection Status")]
        public bool isConnected = false;
        
        private ROSConnection rosConnection;
        
        void Start()
        {
            // Get ROS connection
            rosConnection = ROSConnection.GetOrCreateInstance();
            rosConnection.ConnectOnStart = true;
            
            // Set ROS IP and port
            rosConnection.RosIPAddress = rosIPAddress;
            rosConnection.RosPort = rosPort;
            
            StartCoroutine(CheckConnection());
        }
        
        IEnumerator CheckConnection()
        {
            yield return new WaitForSeconds(2f);
            
            while (true)
            {
                isConnected = rosConnection.HasConnectionThread && rosConnection.IsConnected.Value;
                
                if (isConnected)
                {
                    Debug.Log("ROS2 Connection established successfully!");
                }
                else
                {
                    Debug.LogWarning("ROS2 Connection not established. Check Docker container status.");
                }
                
                yield return new WaitForSeconds(5f);
            }
        }
        
        void OnGUI()
        {
            GUI.Label(new Rect(10, 10, 300, 20), $"ROS2 Status: {(isConnected ? "Connected" : "Disconnected")}");
            GUI.Label(new Rect(10, 30, 300, 20), $"ROS IP: {rosIPAddress}:{rosPort}");
        }
    }
}
"@

$scriptsPath = Join-Path $ProjectPath "Assets\Scripts\ROS"
$connectionScriptPath = Join-Path $scriptsPath "ROSConnection.cs"
Write-Output $rosConnectionScript | Out-File -FilePath $connectionScriptPath -Encoding UTF8

Write-Host "Created ROS connection script at: $connectionScriptPath" -ForegroundColor Green

# Create assembly definition for ROS scripts
$asmdefContent = @"
{
    "name": "ROS2NiryoControl",
    "rootNamespace": "ROS2NiryoControl",
    "references": [
        "Unity.Robotics.ROSTCPConnector",
        "Unity.Robotics.ROSTCPConnector.Messages"
    ],
    "includePlatforms": [],
    "excludePlatforms": [],
    "allowUnsafeCode": false,
    "overrideReferences": false,
    "precompiledReferences": [],
    "autoReferenced": true,
    "defineConstraints": [],
    "versionDefines": [],
    "noEngineReferences": false
}
"@

$asmdefPath = Join-Path $scriptsPath "ROS2NiryoControl.asmdef"
Write-Output $asmdefContent | Out-File -FilePath $asmdefPath -Encoding UTF8

Write-Host "Created assembly definition at: $asmdefPath" -ForegroundColor Green

Write-Host ""
Write-Host "ROS packages installation completed!" -ForegroundColor Green
Write-Host ""
Write-Host "Next steps:" -ForegroundColor Yellow
Write-Host "1. Open Unity project and wait for package compilation" -ForegroundColor White
Write-Host "2. Check Package Manager for installed ROS packages" -ForegroundColor White
Write-Host "3. Configure ROS TCP Endpoint settings" -ForegroundColor White
Write-Host "4. Test connection with Docker ROS2 container" -ForegroundColor White
Write-Host ""
Write-Host "Note: Unity will automatically download and compile the packages when you open the project." -ForegroundColor Cyan