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

def convert_pt_to_onnx(weights_path: Path, model_name: str, img_size: int, 
                      output_dir: Path, target_nodes: List[str], 
                      new_output_names: Optional[List[str]] = None,
                      yolo_version: str = "yolov5",
                      task: str = "detect",
                      nc: int = 80) -> Path:
    """
    Convert PyTorch model to ONNX format using YOLO toolbox.
    """
    import sys
    from pathlib import Path as PathLib
    
    # Get the directory where this script is located
    script_dir = PathLib(__file__).parent.absolute()
    
    # Add toolbox directory to path if not already there
    toolbox_dir = script_dir
    if str(toolbox_dir) not in sys.path:
        sys.path.insert(0, str(toolbox_dir))
    
    # Also add parent directory in case nxva is at a different level
    parent_dir = toolbox_dir.parent
    if str(parent_dir) not in sys.path:
        sys.path.insert(0, str(parent_dir))
    
    from nxva.yolo.yolo import YOLOOptions, YOLOToolBox
    
    weights_path = Path(weights_path).absolute()
    if not weights_path.exists():
        raise ValueError(f"Weights file {weights_path} does not exist")
    
    print(f"Weights file {weights_path} exists")
    print(f"YOLO version: {yolo_version}, task: {task}")
    
    options = YOLOOptions(
        weights=str(weights_path),
        version=yolo_version,
        task=task,
        size=img_size,
        device='cpu',
        nc=nc
    )
    
    toolbox = YOLOToolBox(options)
    onnx_output_path = output_dir / f"{model_name}.onnx"
    output_dir.mkdir(parents=True, exist_ok=True)
    
    toolbox.export(
        shape=(1, 3, img_size, img_size),
        dynamic=False,
        output_path=str(onnx_output_path),
        task='torch2onnx'
    )
    return onnx_output_path


def get_yolo_end_nodes(yaml_path: Path, yolo_version: str, task: str) -> List[str]:
    """
    Get YOLO end nodes from YAML file based on version and task.
    """
    yaml_path = Path(yaml_path)
    if not yaml_path.exists():
        raise ValueError(f"YAML file {yaml_path} does not exist")
    
    with open(yaml_path, 'r', encoding='utf-8') as f:
        data = yaml.safe_load(f)
    
    yolo_version = yolo_version.upper()
    task = task.lower()
    
    if yolo_version not in data:
        raise ValueError(f"YOLO version '{yolo_version}' not found in YAML file. Available: {list(data.keys())}")
    
    if task not in data[yolo_version]:
        raise ValueError(f"Task '{task}' not found for '{yolo_version}' in YAML file. Available: {list(data[yolo_version].keys())}")
    
    nodes = data[yolo_version][task]
    
    if not nodes:
        raise ValueError(f"No end nodes defined for {yolo_version}/{task}")
    
    cleaned_nodes = []
    for node in nodes:
        if isinstance(node, str):
            node_clean = node.split('#')[0].strip()
            if node_clean and not node_clean.startswith('#'):
                cleaned_nodes.append(node_clean)
    
    if not cleaned_nodes:
        raise ValueError(f"No valid end nodes found for {yolo_version}/{task} after filtering comments")
    
    return cleaned_nodes
