from __future__ import annotations

import io
import logging
from typing import Any, List, NamedTuple, Optional, Sequence

import tensorrt as trt
import torch
from torch_tensorrt._enums import dtype
from torch_tensorrt._features import ENABLED_FEATURES
from torch_tensorrt._Input import Input
from torch_tensorrt.dynamo._engine_cache import BaseEngineCache
from torch_tensorrt.dynamo._settings import CompilationSettings, settings_are_compatible
from torch_tensorrt.dynamo.conversion._TRTInterpreter import (
    TRTInterpreter,
    TRTInterpreterResult,
)
from torch_tensorrt.dynamo.runtime import PythonTorchTensorRTModule, TorchTensorRTModule
from torch_tensorrt.dynamo.utils import (
    get_cpu_memory_usage,
    get_output_dtypes,
    release_host_and_device_memory,
)
from torch_tensorrt.logging import TRT_LOGGER

logger = logging.getLogger(__name__)


class SerializedInterpreterResult(NamedTuple):
    serialized_engine: bytes
    input_names: Sequence[str]
    output_names: Sequence[str]
    weight_name_map: Optional[dict[Any, Any]]
    requires_output_allocator: bool


def infer_module_output_dtypes(
    module: torch.fx.GraphModule,
    truncate_double: bool = False,
) -> List[dtype]:
    """
    This function get the output dtypes from node.meta['val'] which was set during dynamo compile_module step
    and truncates them accordingly.
    """
    outputs = [node for node in module.graph.nodes if node.op == "output"]
    outputs = outputs[0].args
    return get_output_dtypes(outputs, truncate_double)  # type: ignore[no-any-return]


def insert_engine_to_cache(
    hash_val: Optional[str],
    interpreter_result: TRTInterpreterResult,
    engine_cache: BaseEngineCache,
    settings: CompilationSettings,
    inputs: Sequence[Input],
) -> bool:
    if hash_val is None:
        logger.warning("Hash value is not provided, so the engine will not be cached")
        return False

    if not ENABLED_FEATURES.refit:
        logger.warning(
            "Refit feature is not available, so the engine cache will not be used"
        )
        return False

    # Cache the weight-stripped engine regardless of the `strip_engine_weights` setting
    if engine_cache.check(hash_val) is not None:
        logger.info(
            f"Detected that the engine with hash: {hash_val} exists in cache. It will be refreshed"
        )

    serialization_config = interpreter_result.engine.create_serialization_config()
    # for TensorRT >= 10.14, we save weight-stripped engine in cache
    if hasattr(trt.SerializationFlag, "INCLUDE_REFIT"):
        serialization_config.set_flag(trt.SerializationFlag.EXCLUDE_WEIGHTS)
    serialized_engine_in_cache = interpreter_result.engine.serialize_with_config(
        serialization_config
    )

    # Insert weight-stripped engine to cache
    engine_cache.insert(
        hash_val,
        (
            serialized_engine_in_cache,
            interpreter_result.input_names,
            interpreter_result.output_names,
            inputs,
            settings,
            interpreter_result.weight_name_map,
            interpreter_result.requires_output_allocator,
        ),
    )
    logger.info(f"Engine with hash: {hash_val} was successfully inserted into cache")
    return True


def pull_cached_engine(
    hash_val: Optional[str],
    module: torch.fx.GraphModule,
    engine_cache: BaseEngineCache,
    settings: CompilationSettings,
    inputs: Sequence[Input],
) -> Optional[SerializedInterpreterResult]:
    if hash_val is None:
        logger.warning(
            "Hash value is not provided, so the engine cache will not be used"
        )
        return None

    if not ENABLED_FEATURES.refit:
        logger.warning(
            "Refit feature is not available, so the engine cache will not be used"
        )
        return None

    # query the cached TRT engine
    cached_data = engine_cache.check(hash_val)
    if cached_data is not None:  # hit the cache
        (
            serialized_engine,  # weight-stripped engine
            input_names,
            output_names,
            cached_engine_inputs,
            cached_engine_compilation_settings,
            weight_name_map,
            requires_output_allocator,
        ) = cached_data

        setting_compatiblity, incompattible_settings = settings_are_compatible(
            settings, cached_engine_compilation_settings
        )
        assert (
            setting_compatiblity
        ), f"Attempted to refit a cached engine with incompatible settings: {incompattible_settings}, (old_settings: {cached_engine_compilation_settings}, new_settings: {settings})"

        for i, e in enumerate(
            [Input.equivalent_spec(c, i) for c, i in zip(cached_engine_inputs, inputs)]
        ):
            assert (
                e
            ), f"Attempted to refit a cached engine built for a different input size (input: {i}, cached size: {cached_engine_inputs[i]}, new size: {inputs[i]}"

        logger.info(
            f"Found the cached engine with hash {hash_val} that corresponds to this graph. It is directly loaded."
        )

        # refit the cached engine with the new graph module
        if not settings.strip_engine_weights:
            runtime = trt.Runtime(TRT_LOGGER)
            engine = runtime.deserialize_cuda_engine(
                serialized_engine
            )  # weight-stripped engine

            from torch_tensorrt.dynamo._refit import (
                _refit_single_trt_engine_with_gm,
            )

            # weight-stripped engine --in place--> weight-included engine
            _refit_single_trt_engine_with_gm(
                new_gm=module,
                old_engine=engine,
                input_list=inputs,
                settings=settings,
                weight_name_map=weight_name_map,
            )

            serialization_config = engine.create_serialization_config()
            serialization_config.clear_flag(trt.SerializationFlag.EXCLUDE_WEIGHTS)
            # for TensorRT >= 10.14, we set INCLUDE_REFIT flag to make the engine refittable
            if hasattr(trt.SerializationFlag, "INCLUDE_REFIT"):
                serialization_config.set_flag(trt.SerializationFlag.INCLUDE_REFIT)
            serialized_engine = engine.serialize_with_config(serialization_config)

            del engine
            with io.BytesIO() as engine_bytes:
                engine_bytes.write(serialized_engine)
                serialized_engine = engine_bytes.getvalue()

        return SerializedInterpreterResult(
            serialized_engine=serialized_engine,
            input_names=input_names,
            output_names=output_names,
            weight_name_map=weight_name_map,
            requires_output_allocator=requires_output_allocator,
        )
    return None


def interpret_module_to_result(
    module: torch.fx.GraphModule,
    inputs: Sequence[Input],
    settings: CompilationSettings = CompilationSettings(),
    engine_cache: Optional[BaseEngineCache] = None,
) -> SerializedInterpreterResult:
    """Interpret an FX module to a TRTInterpreterResult
    Args:
        module: FX GraphModule to interpret
        inputs: It requires a sequence of FLATTENED Inputs representing inputs to the module. It should include both arg_inputs and kwarg_inputs, if applicable.
        settings: Compilation settings
        engine_cache: Engine cache instance
    Returns:
        SerializedInterpreterResult
    """

    # engine_cache could be None if:
    # 1) engine_cache is not passed in when calling this function like convert_exported_program_to_serialized_trt_engine etc., or
    # 2) both cache_built_engines and reuse_cached_engines are False

    is_engine_caching_supported = (
        engine_cache is not None
        and ENABLED_FEATURES.refit
        and not settings.immutable_weights
    )
    # calculate the hash only once. It will be used in pulling and inserting the engine.
    hash_val = (
        engine_cache.get_hash(module, inputs, settings)  # type: ignore
        if is_engine_caching_supported
        and (settings.cache_built_engines or settings.reuse_cached_engines)
        else None
    )

    if settings.reuse_cached_engines:
        if engine_cache is None:
            logger.warning(
                "Engine cache is not provided, so the engine will not be reused from cache"
            )
        elif not ENABLED_FEATURES.refit:
            logger.warning(
                "Refit feature is not available, so the engine will not be reused from cache"
            )
        elif settings.immutable_weights:
            logger.warning(
                "The engine weights are immutable, so the engine will not be reused from cache"
            )
        else:
            serialized_interpreter_result = pull_cached_engine(
                hash_val, module, engine_cache, settings, inputs
            )
            if serialized_interpreter_result is not None:  # hit the cache
                return serialized_interpreter_result

    output_dtypes = infer_module_output_dtypes(
        module, truncate_double=settings.truncate_double
    )

    interpreter = TRTInterpreter(
        module,
        inputs,
        output_dtypes=output_dtypes,
        compilation_settings=settings,
        engine_cache=engine_cache,
    )

    interpreter_result = interpreter.run()
    # Delete the frozen parameters from the module to release CPU memory
    del interpreter
    for attr in dir(module):
        if attr.startswith("_frozen_param"):
            delattr(module, attr)
    release_host_and_device_memory()
    logger.debug(
        f"CPU memory usage after clearing frozen parameters and building memory in conversion: {get_cpu_memory_usage()} MB"
    )

    if settings.cache_built_engines:
        if engine_cache is None:
            logger.warning(
                "Engine cache is not provided, so the engine will not be cached"
            )
        elif not ENABLED_FEATURES.refit:
            logger.warning(
                "Refit feature is not available, so the engine will not be cached"
            )
        elif settings.immutable_weights:
            logger.warning(
                "The engine weights are immutable, so the engine will not be cached"
            )
        else:
            _ = insert_engine_to_cache(
                hash_val, interpreter_result, engine_cache, settings, inputs
            )

    serialized_engine = interpreter_result.engine.serialize()
    with io.BytesIO() as engine_bytes:
        engine_bytes.write(serialized_engine)
        serialized_engine = engine_bytes.getvalue()
        logger.debug(
            f"CPU memory usage after serializing engine: {get_cpu_memory_usage()} MB"
        )

    serialized_interpreter_result = SerializedInterpreterResult(
        serialized_engine=serialized_engine,
        input_names=interpreter_result.input_names,
        output_names=interpreter_result.output_names,
        weight_name_map=interpreter_result.weight_name_map,
        requires_output_allocator=interpreter_result.requires_output_allocator,
    )

    return serialized_interpreter_result


def convert_module(
    module: torch.fx.GraphModule,
    inputs: Sequence[Input],
    settings: CompilationSettings = CompilationSettings(),
    name: str = "",
    engine_cache: Optional[BaseEngineCache] = None,
) -> PythonTorchTensorRTModule | TorchTensorRTModule:
    """Convert an FX module to a TRT module
    Args:
        module: FX GraphModule to convert
        inputs: Sequence of Tensors representing inputs to the module
        settings: Compilation settings
        name: TRT engine name
        engine_cache: Engine cache instance
    Returns:
        PythonTorchTensorRTModule or TorchTensorRTModule
    """
    serialized_interpreter_result = interpret_module_to_result(
        module, inputs, settings, engine_cache=engine_cache
    )

    rt_cls = PythonTorchTensorRTModule

    if ENABLED_FEATURES.torch_tensorrt_runtime and not settings.use_python_runtime:
        from torch_tensorrt.dynamo.runtime import TorchTensorRTModule

        rt_cls = TorchTensorRTModule

    elif (
        not ENABLED_FEATURES.torch_tensorrt_runtime and not settings.use_python_runtime
    ):
        logger.info(
            "Since Torch-TensorRT runtime is not available, using Python Runtime, some features may not be available"
        )

    return rt_cls(
        serialized_engine=serialized_interpreter_result.serialized_engine,
        input_binding_names=list(serialized_interpreter_result.input_names),
        output_binding_names=list(serialized_interpreter_result.output_names),
        name=name,
        settings=settings,
        weight_name_map=serialized_interpreter_result.weight_name_map,
        requires_output_allocator=serialized_interpreter_result.requires_output_allocator,
    )