import argparse
from pathlib import Path
from typing import List, Optional

from .axcl_convert import AxclConvert
from .yolo_utils import get_yolo_end_nodes, convert_pt_to_onnx


class AxclYoloConvert(AxclConvert):
    def __init__(self, 
                 weights_path: Path, 
                 model_name: str, 
                 img_size: int, 
                 output_dir: Path, 
                 yolo_version: str,
                 task: str,
                 target_nodes: Optional[List[str]] = None,
                 new_output_names: Optional[List[str]] = None,
                 yaml_path: Optional[Path] = None,
                 config_path: Optional[Path] = None,
                 dataset_images_path: Optional[Path] = None,
                 dataset_tar_path: Optional[Path] = None,
                 calibration_size: Optional[int] = None,
                 target_hardware: str = "AX650"):
        """
        Initialize AxclYoloConvert class.
        
        Args:
            weights_path: Path to the weights file (.pt or .onnx)
            model_name: Name of the model
            img_size: Image size for input
            output_dir: Output directory for all generated files
            yolo_version: YOLO version (e.g., 'YOLOV5', 'YOLO11')
            task: Task type (e.g., 'detect', 'pose', 'segment', 'classify')
            target_nodes: Optional list of target node names. If None, will get from YAML.
            new_output_names: Optional list of new output names.
            yaml_path: Path to YOLO model end nodes YAML file.
            config_path: Path to the config JSON file.
            dataset_images_path: Path to images for creating calibration dataset.
            dataset_tar_path: Path to calibration dataset tar file.
            calibration_size: Number of images in calibration dataset.
            target_hardware: Target hardware (default: AX650)
        """
        self.yolo_version_yaml = yolo_version
        self.yolo_version = yolo_version.lower()
        self.task = task.lower()
        
        # Determine target_nodes if not provided
        if target_nodes is None:
            if yaml_path is None:
                yaml_path = Path(__file__).parent / "yolo_model_end_nodes.yaml"
            else:
                yaml_path = Path(yaml_path).absolute()
            
            target_nodes = get_yolo_end_nodes(yaml_path, self.yolo_version_yaml, self.task)
            print(f"✅ Auto-loaded {len(target_nodes)} end nodes from YAML for {self.yolo_version_yaml}/{self.task}")
        
        super().__init__(
            weights_path=weights_path,
            model_name=model_name,
            img_size=img_size,
            output_dir=output_dir,
            target_nodes=target_nodes,
            new_output_names=new_output_names,
            config_path=config_path,
            dataset_images_path=dataset_images_path,
            dataset_tar_path=dataset_tar_path,
            calibration_size=calibration_size,
            target_hardware=target_hardware
        )

    def convert_pt_to_onnx(self):
        """Convert PyTorch model to ONNX format using YOLO toolbox."""
        print(f"Converting YOLO PyTorch model to ONNX...")
        
        self.onnx_path = convert_pt_to_onnx(
            self.weights_path, 
            self.model_name, 
            self.img_size, 
            self.model_dir, 
            self.target_nodes, 
            self.new_output_names,
            yolo_version=self.yolo_version,
            task=self.task
        )
        print(f"✅ YOLO ONNX model saved: {self.onnx_path}")

    def truncate_onnx(self):
        """Truncate ONNX model with YOLO-specific parameters."""
        from .utils import truncate_onnx, show_onnx
        print(f"Truncating YOLO ONNX model...")
        print("Before truncation:")
        show_onnx(self.onnx_path)
        
        truncate_onnx(
            self.onnx_path, 
            self.truncated_onnx_path, 
            target_nodes=self.target_nodes,
            new_output_names=self.new_output_names
        )
        
        print("After truncation:")
        show_onnx(self.truncated_onnx_path)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Axera YOLO model conversion")
    parser.add_argument("--weights-path", type=str, required=True, help="Path to the weights file (.pt or .onnx)")
    parser.add_argument("--model-name", type=str, required=True, help="Name of the model")
    parser.add_argument("--img-size", type=int, required=True, help="Image size")
    parser.add_argument("--output-dir", type=str, required=True, help="Path to the output directory")
    parser.add_argument("--yolo-version", type=str, required=True, help="YOLO version (e.g., 'YOLOV5', 'YOLO11')")
    parser.add_argument("--task", type=str, required=True, help="Task type (e.g., 'detect', 'pose')")
    parser.add_argument("--target-nodes", nargs="+", default=None, help="Optional list of target node names")
    parser.add_argument("--yaml-path", type=str, default=None, help="Path to YOLO model end nodes YAML file")
    parser.add_argument("--new-output-names", nargs="+", default=None, help="Optional list of new output names")
    parser.add_argument("--config-path", type=str, default=None, help="Path to the config JSON file")
    parser.add_argument("--dataset-images-path", type=str, default=None, help="Path to images for calibration")
    parser.add_argument("--dataset-tar-path", type=str, default=None, help="Path to calibration dataset tar file")
    parser.add_argument("--calibration-size", type=int, default=None, help="Number of images in calibration dataset")
    parser.add_argument("--target-hardware", type=str, default="AX650", help="Target hardware")
    parser.add_argument("--skip-dataset", action="store_true", help="Skip dataset creation")
    parser.add_argument("--max-dataset-images", type=int, default=1000, help="Maximum number of images for dataset")
    
    args = parser.parse_args()
    
    converter = AxclYoloConvert(
        weights_path=Path(args.weights_path),
        model_name=args.model_name,
        img_size=args.img_size,
        output_dir=Path(args.output_dir),
        yolo_version=args.yolo_version,
        task=args.task,
        target_nodes=args.target_nodes,
        new_output_names=args.new_output_names,
        yaml_path=Path(args.yaml_path) if args.yaml_path else None,
        config_path=Path(args.config_path) if args.config_path else None,
        dataset_images_path=Path(args.dataset_images_path) if args.dataset_images_path else None,
        dataset_tar_path=Path(args.dataset_tar_path) if args.dataset_tar_path else None,
        calibration_size=args.calibration_size,
        target_hardware=args.target_hardware
    )
    
    converter.convert(
        create_dataset_flag=not args.skip_dataset,
        max_dataset_images=args.max_dataset_images
    )
