# -*- encoding: utf-8 -*-
# @Author: SWHL
# @Contact: liekkaskono@163.com
import os
import traceback
from pathlib import Path
from typing import Any, Dict, List, Optional

import numpy as np
import tensorrt as trt
from cuda.bindings import runtime as cudart

from ...utils.download_file import DownloadFile, DownloadFileInput
from ...utils.log import logger
from ...utils.model_resolver import normalize_lang, resolve_model_key
from ...utils.typings import EngineType
from ...utils.utils import mkdir
from ..base import FileInfo, InferSession
from .engine_builder import TRTEngineBuilder
from .memory_utils import allocate_buffers, free_buffers


class TRTInferSession(InferSession):
    def __init__(self, cfg: Dict[str, Any]):
        self.cfg = cfg
        self.engine_cfg = cfg.get("engine_cfg", {})
        self.model_root_dir = None
        self._closed = False
        self.device_id = self._setup_cuda_device()

        self.trt_logger = trt.Logger(trt.Logger.WARNING)

        engine_path = self._get_engine_path(cfg)
        self.engine = self._load_or_build_engine(cfg, engine_path)

        self.context = self._create_execution_context(engine_path)

        # Allocate memory buffers (pre-allocated with max shape)
        self.inputs, self.outputs, self.bindings, self.stream = allocate_buffers(
            self.engine, self.context
        )

        logger.info(f"TensorRT engine loaded: {engine_path}")

        self._max_square_size = 2048

    def __enter__(self) -> "TRTInferSession":
        return self

    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
        self.close()

    def close(self) -> None:
        if self._closed:
            return

        self._closed = True

        try:
            # Synchronize stream before cleanup
            if hasattr(self, "stream") and self.stream is not None:
                try:
                    cudart.cudaStreamSynchronize(self.stream)
                except Exception as e:
                    logger.debug(f"Stream sync error during close: {e}")

            # Free GPU memory buffers
            if hasattr(self, "inputs") and hasattr(self, "outputs"):
                try:
                    free_buffers(self.inputs, self.outputs, self.stream)
                except Exception as e:
                    logger.debug(f"Buffer free error during close: {e}")

            # Destroy CUDA stream
            if hasattr(self, "stream") and self.stream is not None:
                try:
                    cudart.cudaStreamDestroy(self.stream)
                except Exception as e:
                    logger.debug(f"Stream destroy error during close: {e}")

        except Exception as e:
            logger.debug(f"Error during session close: {e}")

    def __del__(self):
        self.close()

    def __call__(self, input_content: np.ndarray) -> np.ndarray:
        """Run inference on input data.

        This method executes the TensorRT engine on the provided input.
        It uses pre-allocated buffers to avoid memory allocation overhead.

        Args:
            input_content: Input numpy array with shape matching the model.
                          For detection: (N, C, H, W)
                          For recognition: (N, C, H, W)

        Returns:
            Output numpy array from the model.

        Raises:
            TensorRTError: If inference fails.

        Example:
            >>> input_data = np.random.randn(1, 3, 224, 224).astype(np.float32)
            >>> output = session(input_data)
        """
        try:
            # Step 1: Set input shape for dynamic dimensions
            output_shape = self._set_input_shape(input_content)

            # Step 2: Copy input data to GPU
            self._copy_input_to_device(input_content)

            # Step 3: Execute inference
            self._execute_inference()

            # Step 4: Copy output back to CPU
            output = self._copy_output_to_host(output_shape)

            return output

        except Exception as e:
            error_info = traceback.format_exc()
            raise TensorRTError(
                f"Inference failed for input shape {input_content.shape}:\n{error_info}"
            ) from e

    def _set_input_shape(self, input_content: np.ndarray) -> tuple:
        input_name = self.engine.get_tensor_name(0)
        self.context.set_input_shape(input_name, input_content.shape)

        output_name = self.engine.get_tensor_name(1)
        return self.context.get_tensor_shape(output_name)

    def _copy_input_to_device(self, input_content: np.ndarray) -> None:
        input_flat = input_content.ravel()
        self.inputs[0].host[: input_flat.size] = input_flat

        cudart.cudaMemcpyAsync(
            self.inputs[0].device,
            self.inputs[0].host.ctypes.data,
            input_flat.nbytes,
            cudart.cudaMemcpyKind.cudaMemcpyHostToDevice,
            self.stream,
        )

    def _execute_inference(self) -> None:
        """Execute TensorRT inference asynchronously."""
        self.context.execute_async_v3(stream_handle=self.stream)

    def _copy_output_to_host(self, output_shape: tuple) -> np.ndarray:
        output_size = int(np.prod(output_shape))
        output_nbytes = output_size * self.outputs[0].host.itemsize

        cudart.cudaMemcpyAsync(
            self.outputs[0].host.ctypes.data,
            self.outputs[0].device,
            output_nbytes,
            cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost,
            self.stream,
        )

        # Wait for all async operations to complete
        cudart.cudaStreamSynchronize(self.stream)

        return self.outputs[0].host[:output_size].reshape(output_shape)

    def _get_max_profile_size(self) -> int:
        try:
            # Try to get from config first
            profile_cfg = self.engine_cfg.get("det_profile", {})
            max_shape = profile_cfg.get("max_shape")
            if max_shape and len(max_shape) >= 4:
                return max(max_shape[2], max_shape[3])  # max(H, W)

            # Fallback: query from engine
            input_name = self.engine.get_tensor_name(0)
            profile_shape = self.engine.get_tensor_profile_shape(input_name, 0)
            if profile_shape and len(profile_shape) >= 3:
                max_shape = profile_shape[2]  # Index 2 = max shape
                return max(max_shape[2], max_shape[3])
        except Exception as e:
            logger.debug(f"Could not get max profile size: {e}")

        # Default fallback
        return 2048

    def _setup_cuda_device(self) -> int:
        device_id = self.engine_cfg.get("device_id", 0)
        status_tuple = cudart.cudaSetDevice(device_id)
        status = status_tuple[0]
        assert status.value == 0, (
            f"Failed to set CUDA device {device_id}: {status}. "
            f"Ensure CUDA is properly installed and GPU is available."
        )
        return device_id

    def _get_engine_path(self, cfg: Dict[str, Any]) -> Path:
        cache_dir = self.engine_cfg.get("cache_dir")
        if cache_dir is None:
            # Check model_root_dir only if cache_dir is None
            if self.model_root_dir is None:
                model_root_dir = cfg.get("model_root_dir", None)
                if model_root_dir is None:
                    raise ValueError(
                        "Either model_path or model_root_dir must be provided in the configuration."
                    )
                model_root_dir = Path(model_root_dir)
                mkdir(model_root_dir)

                self.model_root_dir = model_root_dir
                if not self.model_root_dir.exists():
                    raise FileNotFoundError(
                        f"model_root_dir {self.model_root_dir} does not exist"
                    )

            cache_dir = self.model_root_dir / "models"

        cache_dir = Path(cache_dir)
        cache_dir.mkdir(parents=True, exist_ok=True)

        model_name = self._get_model_name(cfg)
        gpu_arch = self._get_gpu_arch()
        precision = "fp16" if self.engine_cfg.get("use_fp16", True) else "fp32"
        tf32_override = os.environ.get("NVIDIA_TF32_OVERRIDE", "unset")

        return cache_dir / (
            f"{model_name}_{gpu_arch}_{precision}_tf32{tf32_override}.engine"
        )

    def _get_model_name(self, cfg: Dict[str, Any]) -> str:
        """Extract model name from config for engine filename."""
        if cfg.get("model_path"):
            return Path(cfg["model_path"]).stem

        model_key = resolve_model_key(
            cfg.task_type,
            cfg.ocr_version,
            cfg.lang_type,
            cfg.model_type,
        )
        if model_key is not None:
            return model_key

        task_type = cfg.task_type.value
        lang_type = normalize_lang(cfg.lang_type)
        ocr_version = cfg.ocr_version.value
        model_type = cfg.model_type.value
        return f"{lang_type}_{ocr_version}_{task_type}_{model_type}"

    def _get_gpu_arch(self) -> str:
        """Get GPU architecture string for cache key (e.g., 'sm87')."""
        status, major = cudart.cudaDeviceGetAttribute(
            cudart.cudaDeviceAttr.cudaDevAttrComputeCapabilityMajor, self.device_id
        )
        assert status.value == 0, f"Failed to get compute capability: {status}"

        status, minor = cudart.cudaDeviceGetAttribute(
            cudart.cudaDeviceAttr.cudaDevAttrComputeCapabilityMinor, self.device_id
        )
        assert status.value == 0, f"Failed to get compute capability: {status}"

        return f"sm{major}{minor}"

    def _load_or_build_engine(
        self, cfg: Dict[str, Any], engine_path: Path
    ) -> trt.ICudaEngine:
        force_rebuild = self.engine_cfg.get("force_rebuild", False)

        # Try to load cached engine
        if engine_path.exists() and not force_rebuild:
            try:
                return self._load_engine(engine_path)
            except Exception as e:
                logger.warning(
                    f"Failed to load cached engine {engine_path}: {e}. "
                    f"Will rebuild from ONNX."
                )

        # Build new engine from ONNX
        onnx_path = self._get_onnx_path(cfg)
        logger.info(f"Building TensorRT engine from {onnx_path}")

        builder = TRTEngineBuilder(
            onnx_path=onnx_path,
            engine_path=engine_path,
            cfg=self.engine_cfg,
            task_type=cfg.task_type.value,
            trt_logger=self.trt_logger,
            ocr_version=cfg.ocr_version,
        )
        return builder.build()

    def _get_onnx_path(self, cfg: Dict[str, Any]) -> Path:
        model_path = cfg.get("model_path")

        if model_path is None:
            original_engine_type = cfg.engine_type
            cfg.engine_type = EngineType.ONNXRUNTIME

            model_info = self.get_model_url(
                FileInfo(
                    engine_type=EngineType.ONNXRUNTIME,
                    ocr_version=cfg.ocr_version,
                    task_type=cfg.task_type,
                    lang_type=cfg.lang_type,
                    model_type=cfg.model_type,
                )
            )

            cfg.engine_type = original_engine_type

            # Check model_root_dir only if model_path is None
            if self.model_root_dir is None:
                self.model_root_dir = Path(cfg.get("model_root_dir"))
                if not self.model_root_dir.exists():
                    raise FileNotFoundError(
                        f"model_root_dir {self.model_root_dir} does not exist"
                    )

            model_path = self.model_root_dir / Path(model_info["model_dir"]).name
            download_params = DownloadFileInput(
                file_url=model_info["model_dir"],
                sha256=model_info["SHA256"],
                save_path=model_path,
                logger=logger,
            )
            DownloadFile.run(download_params)

        model_path = Path(model_path)
        self._verify_model(model_path)
        return model_path

    def _load_engine(self, engine_path: Path) -> trt.ICudaEngine:
        runtime = trt.Runtime(self.trt_logger)
        with open(engine_path, "rb") as f:
            engine_data = f.read()

        engine = runtime.deserialize_cuda_engine(engine_data)
        if engine is None:
            raise TensorRTError(f"Failed to deserialize TensorRT engine: {engine_path}")
        return engine

    def _create_execution_context(self, engine_path: Path) -> trt.IExecutionContext:
        context = self.engine.create_execution_context()
        if context is None:
            tf32_override = os.environ.get("NVIDIA_TF32_OVERRIDE", "unset")
            raise TensorRTError(
                "Failed to create TensorRT execution context for "
                f"{engine_path}. NVIDIA_TF32_OVERRIDE={tf32_override}. "
                "Remove the cached engine or rebuild it under the current "
                "TensorRT/CUDA environment."
            )
        return context

    def have_key(self, key: str = "character") -> bool:
        return False

    def get_character_list(self, key: str = "character") -> List[str]:
        return []

    @classmethod
    def get_dict_key_url(cls, file_info: FileInfo) -> Optional[str]:
        # Try Paddle first (usually has dict_url)
        for engine_type in [EngineType.PADDLE, EngineType.ONNXRUNTIME]:
            try:
                fallback_info = FileInfo(
                    engine_type=engine_type,
                    ocr_version=file_info.ocr_version,
                    task_type=file_info.task_type,
                    lang_type=file_info.lang_type,
                    model_type=file_info.model_type,
                )
                model_dict = cls.get_model_url(fallback_info)
                if model_dict and "dict_url" in model_dict:
                    return model_dict["dict_url"]
            except Exception as e:
                logger.debug(f"Failed to get dict URL from {engine_type.value}: {e}")

        return None


class TensorRTError(Exception):
    pass
