using System.Collections.Generic;
using System.IO;
using UnityEngine;
#if UNITY_EDITOR
using UnityEditor;
#endif

namespace URDF
{
    /// <summary>
    /// Main URDF importer that creates Unity GameObjects from URDF data
    /// </summary>
    public static class URDFImporter
    {
        /// <summary>
        /// Import a URDF file and create a GameObject hierarchy in the scene
        /// </summary>
        /// <param name="urdfFilePath">Path to the URDF file</param>
        /// <param name="parentTransform">Parent transform to attach the robot to (null = scene root)</param>
        /// <param name="ignoreURDFScale">If true, ignores URDF scale values (useful if models are already scaled to meters)</param>
        /// <param name="jointPositionScale">Scale factor to apply to joint positions (default 1.0). Use 1000 if models were scaled from mm to m.</param>
        /// <returns>Root GameObject of the imported robot</returns>
        public static GameObject ImportURDF(string urdfFilePath, Transform parentTransform = null, bool ignoreURDFScale = false, float jointPositionScale = 1.0f)
        {
            // Parse URDF
            URDFRobot robot = URDFParser.ParseURDF(urdfFilePath);
            if (robot == null)
            {
                Debug.LogError("Failed to parse URDF file");
                return null;
            }

            // Create root GameObject
            GameObject rootObject = new GameObject(robot.name);
            if (parentTransform != null)
            {
                rootObject.transform.SetParent(parentTransform);
            }

            // Create a dictionary to store link GameObjects
            Dictionary<string, GameObject> linkGameObjects = new Dictionary<string, GameObject>();

            // First pass: Create all link GameObjects
            foreach (var linkPair in robot.links)
            {
                string linkName = linkPair.Key;
                URDFLink link = linkPair.Value;

                GameObject linkObject = new GameObject(linkName);
                linkGameObjects[linkName] = linkObject;

                // Set as child of root for now (will be reparented based on joints)
                linkObject.transform.SetParent(rootObject.transform);
            }

            // Second pass: Set up joint hierarchy and apply transforms
            foreach (var jointPair in robot.joints)
            {
                URDFJoint joint = jointPair.Value;

                if (!linkGameObjects.ContainsKey(joint.parentLink) || 
                    !linkGameObjects.ContainsKey(joint.childLink))
                {
                    Debug.LogWarning($"Joint {joint.name} references missing links: parent={joint.parentLink}, child={joint.childLink}");
                    continue;
                }

                GameObject parentObject = linkGameObjects[joint.parentLink];
                GameObject childObject = linkGameObjects[joint.childLink];

                // Parent child to parent
                childObject.transform.SetParent(parentObject.transform);

                // Apply joint origin as local transform
                Vector3 unityPosition = URDFCoordinateConverter.ConvertPosition(joint.origin.xyz);
                Quaternion unityRotation = URDFCoordinateConverter.ConvertRotation(joint.origin.rpy);

                // Scale joint positions if needed (e.g., if models were scaled from mm to m)
                unityPosition *= jointPositionScale;

                childObject.transform.localPosition = unityPosition;
                childObject.transform.localRotation = unityRotation;
            }

            // Third pass: Add visual meshes to links
            foreach (var linkPair in robot.links)
            {
                string linkName = linkPair.Key;
                URDFLink link = linkPair.Value;
                GameObject linkObject = linkGameObjects[linkName];

                // Add visual meshes
                foreach (URDFVisual visual in link.visuals)
                {
                    if (visual.geometry != null && visual.geometry.type == URDFGeometry.GeometryType.Mesh)
                    {
                        string meshPath = URDFMeshResolver.ResolveMeshPath(visual.geometry.meshFilename, urdfFilePath);
                        if (!string.IsNullOrEmpty(meshPath))
                        {
                            #if UNITY_EDITOR
                            Debug.Log($"Loading mesh for link {linkName}: {meshPath}");
                            
                            // Ensure path uses forward slashes for Unity
                            meshPath = meshPath.Replace('\\', '/');
                            
                            // Load the mesh/model - OBJ files are imported as GameObjects in Unity
                            GameObject meshPrefab = AssetDatabase.LoadAssetAtPath<GameObject>(meshPath);
                            
                            // Fallback: Try to find by filename if direct path fails
                            if (meshPrefab == null)
                            {
                                string fileName = Path.GetFileNameWithoutExtension(meshPath);
                                string[] guids = AssetDatabase.FindAssets(fileName + " t:GameObject");
                                foreach (string guid in guids)
                                {
                                    string foundPath = AssetDatabase.GUIDToAssetPath(guid);
                                    if (foundPath.Contains(fileName) && (foundPath.EndsWith(".obj") || foundPath.EndsWith(".fbx") || foundPath.EndsWith(".glb")))
                                    {
                                        Debug.Log($"Found mesh by search: {foundPath}");
                                        meshPrefab = AssetDatabase.LoadAssetAtPath<GameObject>(foundPath);
                                        if (meshPrefab != null)
                                            break;
                                    }
                                }
                            }
                            
                            if (meshPrefab != null)
                            {
                                Debug.Log($"Successfully loaded mesh prefab: {meshPath}");
                                
                                // Instantiate the mesh as a child
                                GameObject meshInstance;
                                
                                // Try to instantiate as prefab first (preserves prefab connection)
                                if (PrefabUtility.IsPartOfPrefabAsset(meshPrefab))
                                {
                                    meshInstance = PrefabUtility.InstantiatePrefab(meshPrefab) as GameObject;
                                }
                                else
                                {
                                    meshInstance = Object.Instantiate(meshPrefab);
                                }
                                
                                if (meshInstance != null)
                                {
                                    meshInstance.name = Path.GetFileNameWithoutExtension(meshPath);
                                    meshInstance.transform.SetParent(linkObject.transform);
                                    
                                    // Apply visual origin transform
                                    Vector3 visualPos = URDFCoordinateConverter.ConvertPosition(visual.origin.xyz);
                                    Quaternion visualRot = URDFCoordinateConverter.ConvertRotation(visual.origin.rpy);
                                    meshInstance.transform.localPosition = visualPos;
                                    meshInstance.transform.localRotation = visualRot;
                                    
                                    // Apply mesh scale from URDF
                                    // Note: URDF scale="0.001" typically means models are in millimeters
                                    // If models are already scaled to meters, ignoreURDFScale should be true
                                    if (!ignoreURDFScale && visual.geometry.scale != Vector3.one)
                                    {
                                        meshInstance.transform.localScale = visual.geometry.scale;
                                    }
                                    else if (ignoreURDFScale)
                                    {
                                        // Models are already in meters, use unity scale
                                        meshInstance.transform.localScale = Vector3.one;
                                    }
                                    
                                    // Apply materials to all renderers
                                    MeshRenderer[] renderers = meshInstance.GetComponentsInChildren<MeshRenderer>(true);
                                    if (renderers != null && renderers.Length > 0)
                                    {
                                        Debug.Log($"Found {renderers.Length} renderer(s) on mesh {meshInstance.name}");
                                        foreach (MeshRenderer renderer in renderers)
                                        {
                                            if (renderer == null) continue;
                                            
                                            // Check if the renderer already has materials (from .mtl file import)
                                            bool hasValidMaterials = renderer.sharedMaterials != null && 
                                                                     renderer.sharedMaterials.Length > 0 && 
                                                                     renderer.sharedMaterials[0] != null;
                                            
                                            if (hasValidMaterials)
                                            {
                                                Debug.Log($"Renderer on {renderer.gameObject.name} already has materials from .mtl file, preserving them");
                                                
                                                // Only apply URDF color override if specified and different from default
                                                if (visual.color != Color.white && visual.color.a > 0)
                                                {
                                                    // Create a copy of the existing material and apply color tint
                                                    Material[] materials = new Material[renderer.sharedMaterials.Length];
                                                    for (int i = 0; i < renderer.sharedMaterials.Length; i++)
                                                    {
                                                        if (renderer.sharedMaterials[i] != null)
                                                        {
                                                            materials[i] = new Material(renderer.sharedMaterials[i]);
                                                            materials[i].color = visual.color;
                                                        }
                                                        else
                                                        {
                                                            materials[i] = renderer.sharedMaterials[i];
                                                        }
                                                    }
                                                    renderer.materials = materials;
                                                    Debug.Log($"Applied URDF color override to existing materials");
                                                }
                                                // Otherwise, keep the imported materials as-is
                                            }
                                            else
                                            {
                                                // No materials from .mtl file, create new ones
                                                Debug.Log($"No materials found on {renderer.gameObject.name}, creating new material");
                                                
                                                Material mat;
                                                
                                                // Try URP shaders first, then fallback to built-in shaders
                                                Shader shader = Shader.Find("Universal Render Pipeline/Lit");
                                                if (shader == null)
                                                {
                                                    shader = Shader.Find("Universal Render Pipeline/Simple Lit");
                                                }
                                                if (shader == null)
                                                {
                                                    shader = Shader.Find("Standard");
                                                }
                                                if (shader == null)
                                                {
                                                    shader = Shader.Find("Diffuse");
                                                }
                                                
                                                if (shader == null)
                                                {
                                                    Debug.LogError("Could not find any suitable shader for material!");
                                                    continue;
                                                }
                                                
                                                mat = new Material(shader);
                                                
                                                // Use color from URDF if specified, otherwise use default grey
                                                if (visual.color != Color.white && visual.color.a > 0)
                                                {
                                                    mat.color = visual.color;
                                                }
                                                else
                                                {
                                                    // Default grey material
                                                    mat.color = new Color(0.7f, 0.7f, 0.7f, 1f); // Light grey
                                                }
                                                
                                                if (mat != null)
                                                {
                                                    renderer.material = mat;
                                                    Debug.Log($"Applied new material to renderer on {renderer.gameObject.name}");
                                                }
                                                else
                                                {
                                                    Debug.LogWarning($"Failed to create material for renderer on {renderer.gameObject.name}");
                                                }
                                            }
                                        }
                                    }
                                    else
                                    {
                                        Debug.LogWarning($"No renderers found on mesh instance: {meshInstance.name}");
                                    }
                                    
                                    Debug.Log($"Successfully added mesh to link {linkName}: {meshInstance.name}");
                                }
                                else
                                {
                                    Debug.LogError($"Failed to instantiate mesh prefab: {meshPath}");
                                }
                            }
                            else
                            {
                                Debug.LogWarning($"Could not load mesh prefab from: {meshPath}. File may not be imported in Unity yet.");
                                
                                // Try to refresh the asset database and reload
                                AssetDatabase.Refresh();
                                meshPrefab = AssetDatabase.LoadAssetAtPath<GameObject>(meshPath);
                                if (meshPrefab != null)
                                {
                                    Debug.Log($"Mesh loaded after refresh: {meshPath}");
                                    GameObject meshInstance = Object.Instantiate(meshPrefab);
                                    meshInstance.name = Path.GetFileNameWithoutExtension(meshPath);
                                    meshInstance.transform.SetParent(linkObject.transform);
                                    
                                    Vector3 visualPos = URDFCoordinateConverter.ConvertPosition(visual.origin.xyz);
                                    Quaternion visualRot = URDFCoordinateConverter.ConvertRotation(visual.origin.rpy);
                                    meshInstance.transform.localPosition = visualPos;
                                    meshInstance.transform.localRotation = visualRot;
                                    
                                    if (visual.geometry.scale != Vector3.one)
                                    {
                                        meshInstance.transform.localScale = visual.geometry.scale;
                                    }
                                }
                            }
                            #endif
                        }
                        else
                        {
                            Debug.LogWarning($"Could not resolve mesh path for link {linkName}: {visual.geometry.meshFilename}");
                        }
                    }
                }
            }

            Debug.Log($"Successfully imported URDF robot: {robot.name}");
            return rootObject;
        }
    }
}