"""
Pinocchio model loader for robots.

This module provides PinocchioModel class for loading and managing Pinocchio models
and collision geometry, making them available for multiple tasks (kinematics, dynamics,
collision checking, visualization, etc.).
"""

import logging
import pinocchio as pin
import numpy as np
from pathlib import Path
from feathersdk.utils.feathertypes import List, Optional
from feathersdk.robot.controls.planner import EasyIKPlanner, quaternion_to_rotation_matrix, rotation_matrix_to_quaternion
from typing import Tuple, Union
from feathersdk.utils.logger import info


logger = logging.getLogger(__name__)


class PinocchioModel:
    """
    Reusable Pinocchio model and geometry loader.
    
    This class loads and manages the Pinocchio model and collision geometry,
    making them available for multiple tasks (kinematics, dynamics, collision checking, etc.).
    """
    
    def __init__(
        self, 
        urdf_path: str, 
        package_dir: Optional[str] = None,
        load_collision: bool = True
    ):
        """
        Initialize Pinocchio model from URDF.
        
        Args:
            urdf_path: Path to the URDF file
            package_dir: Directory containing mesh files (for package:// paths)
            load_collision: If True, load collision geometry model
        """
        self.urdf_path = Path(urdf_path)
        
        # Determine package directory
        if package_dir is None:
            # Try to auto-detect: look for package.xml to find the package root
            current = self.urdf_path.parent
            while current != current.parent:  # Stop at filesystem root
                if (current / "package.xml").exists():
                    # Found package root (Nov11_description/)
                    # package_dir should be parent of this
                    self.package_dir = current.parent
                    logger.info(f"Auto-detected package_dir: {self.package_dir}")
                    break
                current = current.parent
            else:
                # Fallback: assume package is 2 levels up from URDF
                self.package_dir = self.urdf_path.parent.parent.parent
                logger.info(f"Using inferred package_dir: {self.package_dir}")
        else:
            self.package_dir = Path(package_dir)
        
        # Load the kinematic model
        self.model = pin.buildModelFromUrdf(str(self.urdf_path))
        self.data = self.model.createData()
        self.last_q = pin.neutral(self.model)
        self.ik_planner = EasyIKPlanner(self.model)
        
        # Load collision geometry if requested
        self.geom_model = None
        self.geom_data = None
        if load_collision:
            self._load_collision_geometry()
        
        logger.info(f"✓ Model loaded: {self.model.name}")
        self.n_frames = getattr(self.model, 'nframes', len(self.model.names))
        self.link_names = [name for name in self.model.names]
        logger.info(f"✓ Number of frames: {self.n_frames}")
        logger.info(f"✓ Number of joints: {self.model.njoints}")
        if self.geom_model:
            logger.info(f"✓ Number of collision geometries: {len(self.geom_model.geometryObjects)}")

    def update_q(self, q: np.ndarray):
        self.last_q = q
        self.ik_planner.update_q(q)

    def get_joint_pose(self, joint_id: Union[int, str], is_dict: bool = False) -> np.ndarray:
        """Get the pose of the given joint as a pose."""
        if isinstance(joint_id, str):
            joint_id = self.model.getJointId(joint_id)
        if is_dict:
            return {
                "position": self.data.oMi[joint_id].translation,
                "rotation": self.data.oMi[joint_id].rotation
            }
        else:
            return np.concatenate([self.data.oMi[joint_id].translation, rotation_matrix_to_quaternion(self.data.oMi[joint_id].rotation)])
    
    def _load_collision_geometry(self):
        """Load collision geometry from URDF."""
        logger.info(f"Loading collision geometry with package_dirs: [{self.package_dir}]")
        
        try:
            self.geom_model = pin.buildGeomFromUrdf(
                self.model,
                str(self.urdf_path),
                pin.GeometryType.COLLISION,
                None,  # geometry_model=None
                [str(self.package_dir)]  # package_dirs as list
            )
            # 2. Add all possible collision pairs (Self-Collision)
            # This includes every link against every other link.
            self.geom_model.addAllCollisionPairs()
            # 3. Optimization: Remove adjacent links 
            # (Links connected by a joint usually shouldn't collide with each other)
            # pin.removeCollisionPairs(self.model, self.geom_model, str(self.urdf_path))
            self.geom_data = self.geom_model.createData()
            self.collision_data = pin.GeometryData(self.geom_model)

            # Needed to remove collisions of neighboring links and when the torso is at the bottom.
            self.check_collision({
                "EpL": -1.57,
                "EpR": 1.57,
                "TvC": -0.55
            })
            self.remove_all_current_collisions()
            logger.info(f"✓ Successfully loaded collision geometry")
        except Exception as e:
            # Fallback: try without package_dirs
            logger.warning(f"⚠ Loading with package_dirs failed, trying without...")
            try:
                self.geom_model = pin.buildGeomFromUrdf(
                    self.model,
                    str(self.urdf_path),
                    pin.GeometryType.COLLISION
                )
                self.geom_data = self.geom_model.createData()
                self.geom_model.addAllCollisionPairs()
                # 3. Optimization: Remove adjacent links 
                # (Links connected by a joint usually shouldn't collide with each other)
                # pin.removeCollisionPairs(self.model, self.geom_model, "/home/hoa/Workspace/featheros/jetson_kernel/on_filesys/feathersys/user_config/conf/simple_URDF_description/meshes/robot.srdf")
                
                self.geom_data = self.geom_model.createData()
                self.collision_data = pin.GeometryData(self.geom_model)

                # Needed to remove collisions of neighboring links and when the torso is at the bottom.
                self.check_collision({
                    "EpL": -1.57,
                    "EpR": 1.57,
                    "TvC": -0.55
                })
                self.remove_all_current_collisions()

                logger.info(f"✓ Successfully loaded collision geometry (without package_dirs)")
            except Exception as e2:
                logger.error(f"\n❌ Failed to load collision geometry")
                logger.error(f"   package_dirs tried: [{self.package_dir}]")
                logger.error(f"   Errors: {e}, {e2}")
                raise RuntimeError(
                    f"Failed to load collision geometry. "
                    f"package_dirs: [{self.package_dir}]. "
                    f"Last error: {e2}"
                )
        
        # Initialize geometry data with neutral configuration
        # try:
        #     q_neutral = pin.neutral(self.model)
        #     pin.forwardKinematics(self.model, self.data, q_neutral)
        #     pin.updateGeometryPlacements(
        #         self.model, self.data, self.geom_model, self.geom_data, q_neutral
        #     )
        #     logger.info("✓ Geometry data initialized successfully")
        # except Exception as e:
        #     logger.warning(f"⚠ Warning: Could not initialize geometry data: {e}")
    
    def check_collision_path(self, q_start: np.ndarray, q_end: np.ndarray) -> bool:
        # Implement me
        pass

    def check_collision(self, q_updates: Union[np.ndarray, dict], stop_at_first_collision: bool = False) -> bool:
        q = self.q_from_q_updates(q_updates)
        return pin.computeCollisions(
            self.model, 
            self.data, 
            self.geom_model, 
            self.geom_data, 
            q,
            stop_at_first_collision
        )

    def get_collision_pairs_from_last_collision_check(self) -> List[Tuple[str, str]]:
        # Find which pairs are actually in collision
        collisions = []
        for i, result in enumerate(self.geom_data.collisionResults):
            if result.isCollision():
                # Get the pair definition
                pair = self.geom_model.collisionPairs[i]
                
                # Get the names of the geometry objects
                name1 = self.geom_model.geometryObjects[pair.first].name
                name2 = self.geom_model.geometryObjects[pair.second].name
                
                collisions.append((name1, name2))
        return collisions
    
    def remove_all_current_collisions(self):
        collisions = self.get_collision_pairs_from_last_collision_check()
        for link_a, link_b in collisions:
            # 1. Get the IDs of the collision objects
            # Note: These names must match the names in your geometry_model
            if self.geom_model.existGeometryName(link_a) and self.geom_model.existGeometryName(link_b):
                id_a = self.geom_model.getGeometryId(link_a)
                id_b = self.geom_model.getGeometryId(link_b)
                
                # 2. Create a CollisionPair object
                pair = pin.CollisionPair(id_a, id_b)
                
                # 3. Remove the pair from the model
                if self.geom_model.existCollisionPair(pair):
                    self.geom_model.removeCollisionPair(pair)
                    info(f"Removed: {link_a} <-> {link_b}")
            else:
                info(f"Warning: One of these links was not found: {link_a}, {link_b}")

        # 4. Critical: Re-generate the geometry data to apply changes
        self.geom_data = pin.GeometryData(self.geom_model)

    def q_from_q_updates(self, q_updates: Union[np.ndarray, dict]) -> np.ndarray:
        q = q_updates
        if isinstance(q, dict):
            q = self.last_q
        for name in q_updates.keys():
            q[self.model.joints[self.model.getJointId(name)].idx_q] = q_updates[name]
        if q.ndim > 1:
            q = q.flatten()
        q = np.asarray(q, dtype=np.float64)
        return q

    def update_kinematics(self, q_updates: Union[np.ndarray, dict]):
        """
        Update forward kinematics for the given configuration.
        
        Args:
            q: Joint configuration vector (nq x 1) or a map of the joint names to the joint values
        """
        q = self.q_from_q_updates(q_updates)
        self.update_q(q)
        
        if len(q) != self.model.nq:
            raise ValueError(
                f"Configuration size mismatch: got {len(q)}, expected {self.model.nq}"
            )
        
        pin.forwardKinematics(self.model, self.data, q)
    
    def update_geometry(self, q: np.ndarray):
        """
        Update geometry placements based on joint configuration.
        
        Args:
            q: Joint configuration vector (nq x 1)
        """
        if self.geom_model is None:
            raise RuntimeError("Collision geometry not loaded. Set load_collision=True in __init__")
        
        self.update_kinematics(q)
        pin.updateGeometryPlacements(
            self.model, self.data, self.geom_model, self.geom_data, q
        )
    
    def get_link_names(self) -> List[str]:
        return self.link_names

    def get_geometry_names(self) -> List[str]:
        """Get list of all collision geometry names."""
        if self.geom_model is None:
            return []
        return [geom.name for geom in self.geom_model.geometryObjects]
