# -*- encoding: utf-8 -*-
# @Author: SWHL
# @Contact: liekkaskono@163.com
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.typings import EngineType
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.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"
            )

        self.cfg = cfg
        self.engine_cfg = cfg.get("engine_cfg", {})
        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.engine.create_execution_context()

        # 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}")

        # Detect MULTI model for square input requirement
        self._requires_square_input = self._check_multi_model(cfg)
        if self._requires_square_input:
            self._max_square_size = self._get_max_profile_size()
            logger.debug(
                f"MULTI det model: requires square input, max_size={self._max_square_size}"
            )
        else:
            self._requires_square_input = False
            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: Handle square input for MULTI model
            original_hw = None
            if self._requires_square_input:
                input_content, original_hw = self._pad_to_square(input_content)

            # Step 2: Set input shape for dynamic dimensions
            output_shape = self._set_input_shape(input_content)

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

            # Step 4: Execute inference
            self._execute_inference()

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

            # Step 6: Crop output back to original shape
            if original_hw is not None:
                output = self._crop_output(output, original_hw)

            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 _check_multi_model(self, cfg: Dict[str, Any]) -> bool:
        try:
            from ...utils.typings import LangDet, TaskType

            is_det = getattr(cfg, "task_type", None) == TaskType.DET
            is_multi = getattr(cfg, "lang_type", None) == LangDet.MULTI
            return is_det and is_multi
        except Exception:
            return False

    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 _pad_to_square(self, input_content: np.ndarray) -> tuple:
        N, C, H, W = input_content.shape

        if H == W:
            return input_content, None  # Already square

        # Calculate square size
        square_size = max(H, W)

        # Limit to max profile size if needed
        if square_size > self._max_square_size:
            logger.warning(
                f"Square size {square_size} exceeds max profile size {self._max_square_size}. "
                f"Limiting to {self._max_square_size}. This may cause accuracy loss."
            )
            square_size = self._max_square_size

        # Ensure divisible by 32 (TensorRT requirement)
        square_size = int(round(square_size / 32) * 32)

        # Create padded array (zero padding)
        padded = np.zeros((N, C, square_size, square_size), dtype=input_content.dtype)

        # Copy original content to top-left
        copy_h = min(H, square_size)
        copy_w = min(W, square_size)
        padded[:, :, :copy_h, :copy_w] = input_content[:, :, :copy_h, :copy_w]

        return padded, (H, W)

    def _crop_output(self, output: np.ndarray, original_hw: tuple) -> np.ndarray:
        """Crop output back to original shape after square inference.

        Args:
            output: Output array with shape (N, C, S, S) from square inference.
            original_hw: Original (H, W) before padding.

        Returns:
            Cropped output array with shape (N, C, H, W).
        """
        if original_hw is None:
            return output

        orig_h, orig_w = original_hw
        out_h, out_w = output.shape[2:4]

        if out_h >= orig_h and out_w >= orig_w:
            # Normal case: crop to original size
            return output[:, :, :orig_h, :orig_w]
        else:
            # Output is smaller - need to scale
            # This happens when square_size was limited by max_profile_size
            scale_h = out_h / max(orig_h, orig_w)
            scale_w = out_w / max(orig_h, orig_w)
            crop_h = int(orig_h * scale_h)
            crop_w = int(orig_w * scale_w)
            return output[:, :, :crop_h, :crop_w]

    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:
            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"

        return cache_dir / f"{model_name}_{gpu_arch}_{precision}.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

        task_type = cfg.task_type.value
        lang_type = cfg.lang_type.value
        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,
        )
        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

            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()
        return runtime.deserialize_cuda_engine(engine_data)

    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
