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

from .utils import (
    show_onnx, truncate_onnx, 
    modify_config, build_axmodel, create_dataset,
    count_images_in_tar
)


class AxclConvert:
    def __init__(self, 
                 weights_path: Path, 
                 model_name: str, 
                 img_size: int, 
                 output_dir: Path, 
                 target_nodes: List[str],
                 new_output_names: Optional[List[str]] = 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 AxclConvert 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
            target_nodes: List of target node names to extract.
            new_output_names: Optional list of new output names. If None, uses "1", "2", "3"...
            config_path: Path to the config JSON file. If None, will use output_dir/config/{model_name}_config.json
            dataset_images_path: Path to images for creating calibration dataset. If None, dataset creation is skipped.
            dataset_tar_path: Path to calibration dataset tar file. If None, will use output_dir/dataset/images.tar.
                              If provided and file exists, will use existing file instead of creating new one.
            calibration_size: Number of images in calibration dataset. If None, will be auto-detected from tar file.
            target_hardware: Target hardware (default: AX650)
        """
        weights_path = Path(weights_path)
        if not weights_path.exists():
            raise ValueError(f"Weights file {weights_path} does not exist")
        
        self.weights_path = weights_path.absolute()
        self.model_format = weights_path.suffix
        self.model_name = model_name
        self.img_size = img_size
        self.output_dir = Path(output_dir).absolute()
        self.new_output_names = new_output_names
        self.target_hardware = target_hardware
        self.target_nodes = target_nodes
        
        # Set up paths
        self.output_dir.mkdir(parents=True, exist_ok=True)
        self.model_dir = self.output_dir / "model"
        self.model_dir.mkdir(parents=True, exist_ok=True)
        self.config_dir = self.output_dir / "config"
        self.config_dir.mkdir(parents=True, exist_ok=True)
        
        # ONNX paths
        if self.model_format == ".pt":
            self.onnx_path = self.model_dir / f"{model_name}.onnx"
        else:
            self.onnx_path = weights_path
        
        self.truncated_onnx_path = self.model_dir / f"{model_name}_truncated.onnx"
        
        # Config path
        if config_path is None:
            self.config_path = self.config_dir / f"{model_name}_config.json"
        else:
            self.config_path = Path(config_path).absolute()
        
        # Dataset path
        self.dataset_images_path = Path(dataset_images_path).absolute() if dataset_images_path else None
        
        # Dataset tar path - use user-specified location only, no duplicate creation
        if dataset_tar_path is not None:
            # User specified custom location - use it exclusively
            self.dataset_tar_path = Path(dataset_tar_path).absolute()
        else:
            # Default location: output_dir/dataset/images.tar
            self.dataset_tar_path = self.output_dir / "dataset" / "images.tar"
        
        # Store calibration size
        self.calibration_size = calibration_size

    def convert_pt_to_onnx(self):
        """
        Convert PyTorch model to ONNX format.
        Base class implementation - should be overridden for specific frameworks.
        """
        raise NotImplementedError("convert_pt_to_onnx must be implemented in a subclass for .pt models")

    def create_calibration_dataset(self, max_num: int = 1000):
        """
        Create calibration dataset from images.
        
        Args:
            max_num: Maximum number of images to include in dataset
        """
        # Check if dataset tar file already exists
        if self.dataset_tar_path.exists():
            print(f"✅ Dataset tar file already exists: {self.dataset_tar_path}")
            print(f"   Skipping dataset creation. To recreate, delete the file first.")
            return
        
        if self.dataset_images_path is None:
            print("⚠️  No dataset images path provided, skipping dataset creation")
            return
        
        if not self.dataset_images_path.exists():
            raise ValueError(f"Dataset images path {self.dataset_images_path} does not exist")
        
        print(f"Creating calibration dataset from {self.dataset_images_path}...")
        print(f"Output: {self.dataset_tar_path}")
        actual_count = create_dataset(
            self.dataset_images_path, 
            self.dataset_tar_path, 
            img_size=self.img_size,
            max_num=max_num
        )
        # Store the actual number of images created
        self.calibration_size = actual_count
        print(f"✅ Dataset created: {self.dataset_tar_path}")
        print(f"   Calibration size: {self.calibration_size} images")

    def truncate_onnx(self):
        """Truncate ONNX model and rename output nodes."""
        print(f"Truncating 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)

    def modify_config(self, calibration_size: Optional[int] = None, config_dir: Optional[Path] = None):
        """
        Modify config file to update output_processors and calibration_size.
        """
        if not self.config_path.exists():
            raise ValueError(f"Config file {self.config_path} does not exist. Please create it first.")
        
        print(f"Modifying config file: {self.config_path}")
        modify_config(self.config_path, self.target_nodes, self.new_output_names, self.config_dir)
        
        # Update calibration dataset path and size if dataset tar file exists
        if self.dataset_tar_path.exists():
            import json
            with open(self.config_path, 'r', encoding='utf-8') as f:
                config = json.load(f)
            
            # Update calibration dataset path and size
            if 'quant' in config and 'input_configs' in config['quant']:
                for input_config in config['quant']['input_configs']:
                    if input_config.get('tensor_name') == 'images':
                        # Update calibration dataset path
                        config_dir = self.config_path.parent
                        try:
                            rel_path = self.dataset_tar_path.relative_to(config_dir)
                            input_config['calibration_dataset'] = f"./{rel_path}"
                        except ValueError:
                            # If not in same directory tree, use absolute path
                            input_config['calibration_dataset'] = str(self.dataset_tar_path)
                        
                        # Update calibration_size
                        if calibration_size is not None:
                            input_config['calibration_size'] = calibration_size
                        elif self.calibration_size is not None:
                            input_config['calibration_size'] = self.calibration_size
                        else:
                            count = count_images_in_tar(self.dataset_tar_path)
                            if count > 0:
                                input_config['calibration_size'] = count
                                print(f"   Auto-detected calibration_size: {count} images")
                        break
            
            print(f"✅ Updated calibration dataset path in config: {self.dataset_tar_path}")
        else:
            print(f"⚠️  Dataset tar file not found: {self.dataset_tar_path}")

    def build_axmodel(self):
        """Build Axera model using pulsar2."""
        print(f"Building Axera model...")
        build_axmodel(
            self.output_dir, 
            self.model_name, 
            self.truncated_onnx_path,
            self.config_path,
            self.target_hardware
        )

    def convert(self, create_dataset_flag: bool = True, max_dataset_images: int = 1000):
        """Complete conversion pipeline."""
        print(f"Starting conversion pipeline for {self.model_name}...")
        
        if self.model_format == ".pt":
            self.convert_pt_to_onnx()
        
        if create_dataset_flag and self.dataset_images_path:
            self.create_calibration_dataset(max_num=max_dataset_images)
        
        self.truncate_onnx()
        self.modify_config()
        self.build_axmodel()
        
        print(f"✅ Conversion completed! Output directory: {self.output_dir}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Axera 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("--target-nodes", nargs="+", required=True, help="List of target node names to extract")
    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 = AxclConvert(
        weights_path=Path(args.weights_path),
        model_name=args.model_name,
        img_size=args.img_size,
        output_dir=Path(args.output_dir),
        target_nodes=args.target_nodes,
        new_output_names=args.new_output_names,
        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
    )