"""
TensorRT export functionality for BADAS models.

Provides functions to convert ONNX models to TensorRT engines
for maximum inference performance.
"""

import subprocess
import time
from pathlib import Path
from typing import Optional
from dataclasses import dataclass

from badas.export.utils import get_trtexec_command


@dataclass
class TensorRTConfig:
    """Configuration for TensorRT conversion.

    Attributes:
        fp16: Enable FP16 precision (default: True).
        workspace_mb: Workspace size in MB (default: 4096).
        verbose: Print detailed conversion info (default: False).
        trtexec_path: Path to trtexec binary (default: auto-detect).
    """
    fp16: bool = True
    workspace_mb: int = 4096
    verbose: bool = False
    trtexec_path: Optional[str] = None


@dataclass
class TensorRTResult:
    """Result of TensorRT conversion.

    Attributes:
        success: Whether conversion succeeded.
        engine_path: Path to TensorRT engine file.
        file_size_mb: Size of engine file in MB.
        build_time_s: Build time in seconds.
        error: Error message (if failed).
    """
    success: bool
    engine_path: Optional[str] = None
    file_size_mb: Optional[float] = None
    build_time_s: Optional[float] = None
    error: Optional[str] = None


def find_trtexec() -> Optional[str]:
    """
    Find the trtexec binary on the system.

    Returns:
        Path to trtexec or None if not found.
    """
    import shutil

    # Common locations
    locations = [
        "/usr/src/tensorrt/bin/trtexec",
        "/usr/local/tensorrt/bin/trtexec",
        "trtexec",  # In PATH
    ]

    for loc in locations:
        if loc == "trtexec":
            path = shutil.which(loc)
            if path:
                return path
        elif Path(loc).exists():
            return loc

    return None


def convert_to_tensorrt(
    onnx_path: str,
    output_path: str,
    config: Optional[TensorRTConfig] = None,
) -> TensorRTResult:
    """
    Convert an ONNX model to TensorRT engine using trtexec.

    This function uses the trtexec command-line tool for conversion,
    which is more reliable than the Python API for complex models.

    Args:
        onnx_path: Path to input ONNX model.
        output_path: Path for output TensorRT engine.
        config: Conversion configuration.

    Returns:
        TensorRTResult with conversion status.

    Example:
        >>> from badas.export import convert_to_tensorrt, TensorRTConfig
        >>>
        >>> # Default FP16 conversion
        >>> result = convert_to_tensorrt("model.onnx", "model.trt")
        >>> print(f"Engine: {result.engine_path} ({result.file_size_mb:.1f} MB)")
        >>>
        >>> # Custom config
        >>> config = TensorRTConfig(fp16=True, workspace_mb=8192)
        >>> result = convert_to_tensorrt("model.onnx", "model.trt", config=config)
    """
    config = config or TensorRTConfig()

    # Find trtexec
    trtexec = config.trtexec_path or find_trtexec()
    if trtexec is None:
        return TensorRTResult(
            success=False,
            error="trtexec not found. Install TensorRT or specify trtexec_path.",
        )

    # Verify ONNX exists
    onnx_path = Path(onnx_path)
    if not onnx_path.exists():
        return TensorRTResult(
            success=False,
            error=f"ONNX file not found: {onnx_path}",
        )

    # Setup output path
    output_path = Path(output_path)
    output_path.parent.mkdir(parents=True, exist_ok=True)

    # Build command
    cmd = [
        trtexec,
        f"--onnx={onnx_path}",
        f"--saveEngine={output_path}",
        f"--memPoolSize=workspace:{config.workspace_mb}MiB",
    ]
    if config.fp16:
        cmd.append("--fp16")

    if config.verbose:
        print(f"   Running: {' '.join(cmd[:3])}...")

    # Run conversion
    start_time = time.time()
    try:
        result = subprocess.run(
            cmd,
            capture_output=True,
            text=True,
            timeout=1800,  # 30 minute timeout
        )
        build_time = time.time() - start_time

        if result.returncode != 0:
            # Extract error message
            error_lines = [
                line for line in result.stderr.split('\n')
                if '[E]' in line or 'error' in line.lower()
            ]
            error_msg = '\n'.join(error_lines[:5]) if error_lines else result.stderr[:500]
            return TensorRTResult(
                success=False,
                error=f"trtexec failed: {error_msg}",
            )

        if not output_path.exists():
            return TensorRTResult(
                success=False,
                error="trtexec completed but engine file not created",
            )

        file_size_mb = output_path.stat().st_size / 1024 / 1024

        return TensorRTResult(
            success=True,
            engine_path=str(output_path),
            file_size_mb=file_size_mb,
            build_time_s=build_time,
        )

    except subprocess.TimeoutExpired:
        return TensorRTResult(
            success=False,
            error="TensorRT conversion timed out after 30 minutes",
        )
    except Exception as e:
        return TensorRTResult(
            success=False,
            error=str(e),
        )


def convert_to_tensorrt_python(
    onnx_path: str,
    output_path: str,
    config: Optional[TensorRTConfig] = None,
) -> TensorRTResult:
    """
    Convert ONNX to TensorRT using Python API.

    This method uses the TensorRT Python bindings directly.
    Falls back to trtexec if Python bindings fail.

    Args:
        onnx_path: Path to input ONNX model.
        output_path: Path for output TensorRT engine.
        config: Conversion configuration.

    Returns:
        TensorRTResult with conversion status.
    """
    config = config or TensorRTConfig()

    try:
        import tensorrt as trt
    except ImportError:
        if config.verbose:
            print("   TensorRT Python not available, using trtexec")
        return convert_to_tensorrt(onnx_path, output_path, config)

    onnx_path = Path(onnx_path)
    output_path = Path(output_path)

    if not onnx_path.exists():
        return TensorRTResult(
            success=False,
            error=f"ONNX file not found: {onnx_path}",
        )

    output_path.parent.mkdir(parents=True, exist_ok=True)

    try:
        logger = trt.Logger(trt.Logger.VERBOSE if config.verbose else trt.Logger.WARNING)

        builder = trt.Builder(logger)
        network = builder.create_network(
            1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
        )
        parser = trt.OnnxParser(network, logger)

        # Parse ONNX
        with open(onnx_path, 'rb') as f:
            if not parser.parse(f.read()):
                errors = [parser.get_error(i) for i in range(parser.num_errors)]
                return TensorRTResult(
                    success=False,
                    error=f"ONNX parsing failed: {errors}",
                )

        # Configure builder
        builder_config = builder.create_builder_config()
        workspace_bytes = config.workspace_mb * 1024 * 1024
        builder_config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, workspace_bytes)

        if config.fp16 and builder.platform_has_fast_fp16:
            builder_config.set_flag(trt.BuilderFlag.FP16)

        if hasattr(trt.BuilderFlag, 'TF32'):
            builder_config.set_flag(trt.BuilderFlag.TF32)

        # Build engine
        start_time = time.time()
        serialized_engine = builder.build_serialized_network(network, builder_config)
        build_time = time.time() - start_time

        if serialized_engine is None:
            return TensorRTResult(
                success=False,
                error="Failed to build TensorRT engine",
            )

        # Save engine
        with open(output_path, 'wb') as f:
            f.write(serialized_engine)

        file_size_mb = output_path.stat().st_size / 1024 / 1024

        return TensorRTResult(
            success=True,
            engine_path=str(output_path),
            file_size_mb=file_size_mb,
            build_time_s=build_time,
        )

    except Exception as e:
        # Fall back to trtexec
        if config.verbose:
            print(f"   Python API failed ({e}), trying trtexec")
        return convert_to_tensorrt(onnx_path, output_path, config)


__all__ = [
    'TensorRTConfig',
    'TensorRTResult',
    'find_trtexec',
    'convert_to_tensorrt',
    'convert_to_tensorrt_python',
]
