import argparse
import json
import io
import glob
import subprocess
import tarfile
from pathlib import Path
from typing import List, Optional

import onnx
from onnx.utils import extract_model
import numpy as np
from PIL import Image
import cv2


def show_onnx(input_path: Path):
    """Show ONNX model input and output information."""
    m = onnx.load(str(input_path))
    print("Inputs:")
    for i in m.graph.input:
        t = i.type.tensor_type
        shape = [d.dim_value for d in t.shape.dim]
        elem_type = t.elem_type  # Usually 1=float32
        print(f"  {i.name}: shape={shape}, dtype={elem_type}")
    
    print("\nOutputs:")
    for o in m.graph.output:
        t = o.type.tensor_type
        shape = [d.dim_value for d in t.shape.dim]
        elem_type = t.elem_type
        print(f"  {o.name}: shape={shape}, dtype={elem_type}")


def modify_config(config_path: Path, target_nodes: List[str], 
                  new_output_names: Optional[List[str]] = None, config_dir: Optional[Path] = None):
    """
    Modify config file to update output_processors based on number of outputs.
    """
    config_path = Path(config_path)
    if not config_path.exists():
        raise ValueError(f"Config file {config_path} does not exist")
    
    # Determine output names
    if new_output_names is None:
        output_names = [str(i+1) for i in range(len(target_nodes))]
    else:
        if len(new_output_names) != len(target_nodes):
            raise ValueError(f"Number of new_output_names ({len(new_output_names)}) must match number of target_nodes ({len(target_nodes)})")
        output_names = new_output_names
    
    # Load config
    with open(config_path, 'r', encoding='utf-8') as f:
        config = json.load(f)
    
    # Update output_processors
    config['output_processors'] = []
    for output_name in output_names:
        config['output_processors'].append({
            "tensor_name": output_name
        })
    
    # Save config
    if config_dir:
        config_path = config_dir / config_path.name
    
    with open(config_path, 'w', encoding='utf-8') as f:
        json.dump(config, f, indent=2, ensure_ascii=False)
    
    print(f"✅ Config file updated: {config_path}")
    print(f"   Output processors: {output_names}")


def build_axmodel(output_dir: Path, model_name: str, onnx_path: Path, 
                  config_path: Path, target_hardware: str = "AX650"):
    """
    Build Axera model using pulsar2 build command.
    """
    output_dir = Path(output_dir)
    onnx_path = Path(onnx_path)
    config_path = Path(config_path)
    
    if not onnx_path.exists():
        raise ValueError(f"ONNX file {onnx_path} does not exist")
    if not config_path.exists():
        raise ValueError(f"Config file {config_path} does not exist")
    
    output_dir.mkdir(parents=True, exist_ok=True)
    
    # Build pulsar2 command
    cmd = [
        "pulsar2", "build",
        "--target_hardware", target_hardware,
        "--input", str(onnx_path),
        "--output_dir", str(output_dir),
        "--output_name", f"{model_name}.axmodel",
        "--config", str(config_path)
    ]
    
    print(f"Running: {' '.join(cmd)}")
    result = subprocess.run(cmd, capture_output=True, text=True)
    
    if result.returncode != 0:
        print(f"Error output: {result.stderr}")
        raise RuntimeError(f"pulsar2 build failed with return code {result.returncode}")
    
    print(f"✅ Axera model built successfully: {output_dir / f'{model_name}.axmodel'}")


def letterbox(img, new_shape=(640, 640), color=(114, 114, 114)):
    """Resize image with letterbox padding."""
    shape = img.shape[:2]  # current shape [height, width]
    if isinstance(new_shape, int):
        new_shape = (new_shape, new_shape)
    
    # Scale ratio (new / old)
    r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
    
    # Compute padding
    new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))
    dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1]
    dw /= 2
    dh /= 2
    
    if shape[::-1] != new_unpad:  # resize
        img = cv2.resize(img, new_unpad, interpolation=cv2.INTER_LINEAR)
    
    top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))
    left, right = int(round(dw - 0.1)), int(round(dw + 0.1))
    img = cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color)
    
    return img

def create_dataset(
    images_path: Path,
    output_path: Path,
    img_size: int = 640,
    max_num: int = 1000,
    image_extensions: List[str] = None,
):
    if image_extensions is None:
        image_extensions = ['.jpg', '.jpeg', '.png', '.bmp']

    images_path = Path(images_path)
    output_path = Path(output_path).with_suffix('.tar')
    output_path.parent.mkdir(parents=True, exist_ok=True)

    # collect images
    image_files = []
    if images_path.is_dir():
        for ext in image_extensions:
            image_files += glob.glob(str(images_path / f"*{ext}"))
            image_files += glob.glob(str(images_path / f"*{ext.upper()}"))
    else:
        image_files = glob.glob(str(images_path))

    if not image_files:
        raise ValueError(f"No images found in {images_path}")
    import random
    random.shuffle(image_files)
    image_files = image_files[:max_num]

    print(f"Processing {len(image_files)} images...")

    processed = 0
    errors = 0

    with tarfile.open(output_path, mode="w") as tar:
        for idx, img_path in enumerate(image_files):
            try:
                img = cv2.imread(img_path)
                if img is None or img.ndim != 3 or img.shape[2] != 3:
                    errors += 1
                    continue

                img = letterbox(img, new_shape=(img_size, img_size))
                img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

                buf = io.BytesIO()
                Image.fromarray(img).save(buf, format="JPEG", quality=95)
                buf.seek(0)

                info = tarfile.TarInfo(name=f"{idx:06d}.jpg")
                info.size = buf.getbuffer().nbytes

                tar.addfile(info, buf)
                processed += 1

            except Exception as e:
                print(f"Error processing {img_path}: {e}")
                errors += 1

    if processed == 0:
        raise RuntimeError("No valid images processed, calibration dataset is empty")

    print(f"✅ Dataset created: {output_path}")
    print(f"   Total images: {processed}, errors: {errors}")

    return processed


def count_images_in_tar(tar_path: Path) -> int:
    """Count the number of images in a tar file."""
    tar_path = Path(tar_path)
    if not tar_path.exists():
        return 0
    
    count = 0
    try:
        with tarfile.open(tar_path, 'r') as tar:
            for member in tar.getmembers():
                if member.isfile() and any(member.name.lower().endswith(ext) for ext in ['.jpg', '.jpeg', '.png', '.bmp']):
                    count += 1
    except Exception as e:
        print(f"⚠️  Warning: Could not count images in tar file {tar_path}: {e}")
        return 0
    
    return count


def truncate_onnx(input_onnx: Path, output_onnx: Path, 
                  target_nodes: List[str],
                  new_output_names: Optional[List[str]] = None):
    """Extract model and rename output tensors."""
    input_onnx = Path(input_onnx)
    output_onnx = Path(output_onnx)
    
    if not input_onnx.exists():
        raise ValueError(f"Input ONNX file {input_onnx} does not exist")
    
    inputs = ["images"]  # Must be the actual input tensor name in the graph
    extract_model(str(input_onnx), str(output_onnx), input_names=inputs, output_names=target_nodes)
    onnx.checker.check_model(onnx.load(str(output_onnx)))
    model = onnx.load(str(output_onnx))

    if new_output_names is None:
        output_names = [str(i+1) for i in range(len(target_nodes))]
    else:
        if len(new_output_names) != len(target_nodes):
            raise ValueError(f"Number of new_output_names must match number of target_nodes")
        output_names = new_output_names

    name_mapping = {old: new for old, new in zip(target_nodes, output_names)}

    for node in model.graph.node:
        for i, out in enumerate(node.output):
            if out in name_mapping:
                node.output[i] = name_mapping[out]

    for i, out in enumerate(model.graph.output):
        if out.name in name_mapping:
            model.graph.output[i].name = name_mapping[out.name]

    output_onnx.parent.mkdir(parents=True, exist_ok=True)
    onnx.save(model, str(output_onnx))
    onnx.checker.check_model(str(output_onnx))
    print(f"✅ Model truncated successfully: {output_onnx}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="ONNX toolbox utilities")
    subparsers = parser.add_subparsers(dest="command", help="Available commands")
    
    show_parser = subparsers.add_parser("show", help="Show ONNX model information")
    show_parser.add_argument("input_path", type=str)
    
    extract_parser = subparsers.add_parser("extract", help="Extract and rename ONNX model")
    extract_parser.add_argument("input_onnx", type=str)
    extract_parser.add_argument("output_onnx", type=str)
    extract_parser.add_argument("--target-nodes", nargs="+", required=True)
    extract_parser.add_argument("--new-output-names", nargs="+", default=None)
    
    args = parser.parse_args()
    if args.command == "show":
        show_onnx(Path(args.input_path))
    elif args.command == "extract":
        truncate_onnx(Path(args.input_onnx), Path(args.output_onnx), args.target_nodes, args.new_output_names)
