# onnxruntime_genai is not lint-friendly yet and PyLint can't
# find any of the class methods
# pylint: disable=no-member
#
# Model builder constraints:
#   11/10/24 Need transformers <4.45.0 OR onnxruntime-genai 0.5.0 (which must be built from source)
#   (transformers v4.45 changes the format of the tokenizer.json file which will be supported in
#   onnxruntime-genai 0.5)
#

import argparse
import os
import time
import json
import shutil
import logging
from fnmatch import fnmatch
from queue import Queue
import subprocess
from packaging.version import Version
from huggingface_hub import snapshot_download
import onnxruntime_genai as og
import onnxruntime_genai.models.builder as model_builder
from transformers import AutoTokenizer
from lemonade.state import State
from lemonade.tools import FirstTool
import lemonade.common.status as status
import lemonade.common.printing as printing
from lemonade.tools.huggingface_load import get_base_model, is_offline
from lemonade.tools.adapter import (
    ModelAdapter,
    TokenizerAdapter,
    PassthroughTokenizerResult,
)
from lemonade.cache import Keys
from lemonade_install.install import (
    get_ryzen_ai_version_info,
    get_oga_npu_dir,
    get_oga_hybrid_dir,
    SUPPORTED_RYZEN_AI_SERIES,
)


# ONNX Runtime GenAI models will be cached in this subfolder of the lemonade cache folder
oga_models_path = "oga_models"

# ONNX Runtime GenAI model builder tool uses this subfolder of the lemonade cache as its cache
oga_model_builder_cache_path = "model_builder"

# Mapping from processor to execution provider, used in pathnames and by model_builder
execution_providers = {
    "cpu": "cpu",
    "npu": "npu",
    "igpu": "dml",
    "hybrid": "hybrid",
    "cuda": "cuda",
}


class OrtGenaiTokenizer(TokenizerAdapter):
    def __init__(self, model: og.Model, hf_tokenizer: AutoTokenizer):
        super().__init__(hf_tokenizer)
        # Initialize OGA tokenizer
        self.tokenizer = og.Tokenizer(model)

        # Placeholder value since some code will try to query it
        # If we actually need this to return a proper value, then
        # og.GeneratorParams.eos_token_id has it
        self.eos_token_id = None

    def __call__(self, prompt: str, return_tensors="np"):
        tokens = self.tokenizer.encode(prompt)
        return PassthroughTokenizerResult(tokens)

    # pylint: disable=unused-argument
    def decode(self, response, skip_special_tokens=True) -> str:
        return self.tokenizer.decode(response)


class OrtGenaiStreamer:
    def __init__(self, tokenizer: OrtGenaiTokenizer, timeout=None):
        self.tokenizer = tokenizer
        self.text_queue = Queue()
        self.stop_signal = None
        self.timeout = timeout

    def add_text(self, text: str):
        self.text_queue.put(text, timeout=self.timeout)

    def done(self):
        self.text_queue.put(self.stop_signal, timeout=self.timeout)

    def __iter__(self):
        return self

    def __next__(self):
        value = self.text_queue.get(timeout=self.timeout)
        if value == self.stop_signal:
            raise StopIteration()
        else:
            return value


class OrtGenaiModel(ModelAdapter):

    def __init__(self, input_folder):
        super().__init__()
        self.model = og.Model(input_folder)
        self.type = "ort-genai"
        self.config = self.load_config(input_folder)

    def load_config(self, input_folder):
        rai_config_path = os.path.join(input_folder, "rai_config.json")
        if os.path.exists(rai_config_path):
            with open(rai_config_path, "r", encoding="utf-8") as f:
                max_prompt_length = json.load(f)["max_prompt_length"]["1.4.1"]
        else:
            max_prompt_length = None

        config_path = os.path.join(input_folder, "genai_config.json")
        if os.path.exists(config_path):
            with open(config_path, "r", encoding="utf-8") as f:
                config_dict = json.load(f)
                if max_prompt_length:
                    config_dict["max_prompt_length"] = max_prompt_length
                return config_dict
        return None

    def generate(
        self,
        input_ids,
        max_new_tokens=512,
        min_new_tokens=0,
        do_sample=True,
        top_k=50,
        top_p=1.0,
        temperature=0.7,
        streamer: OrtGenaiStreamer = None,
        pad_token_id=None,
        stopping_criteria=None,
        max_length=None,
        random_seed=1,
    ):
        params = og.GeneratorParams(self.model)

        prompt_length = len(input_ids)
        max_prompt_length = self.config.get("max_prompt_length")
        if max_prompt_length and prompt_length > max_prompt_length:
            raise ValueError(
                f"This prompt (length {prompt_length}) exceeds the model's "
                f"maximum allowed prompt length ({max_prompt_length})."
            )

        # There is a breaking API change in OGA 0.6.0
        # Determine whether we should use the old or new APIs
        # This also supports 0.6.0.dev0, which evaluates to less than 0.6.0 in Version
        use_oga_post_6_api = (
            Version(og.__version__) >= Version("0.6.0") or "0.6.0" in og.__version__
        )
        use_oga_pre_6_api = not use_oga_post_6_api

        if pad_token_id:
            params.pad_token_id = pad_token_id

        # Handle max_length and max_new_tokens
        if max_length and max_new_tokens:
            logging.warning(
                "Both max_length and max_new_tokens were provided. "
                "max_length will take precedence. "
                "When setting max_length, please explicitly set max_new_tokens to None."
            )
        max_length_to_use = None
        if max_length:
            max_length_to_use = max_length
        elif max_new_tokens:
            max_length_to_use = prompt_length + max_new_tokens

        min_length = prompt_length + min_new_tokens

        if use_oga_pre_6_api:
            params.input_ids = input_ids

        if random_seed is None:
            random_seed = -1  # In og.Generator, -1 = seed with random device

        if self.config and "search" in self.config:
            search_config = self.config["search"]
            params.set_search_options(
                do_sample=search_config.get("do_sample", do_sample),
                top_k=search_config.get("top_k", top_k),
                top_p=search_config.get("top_p", top_p),
                temperature=search_config.get("temperature", temperature),
                max_length=max_length_to_use,
                min_length=min_length,
                early_stopping=search_config.get("early_stopping", False),
                length_penalty=search_config.get("length_penalty", 1.0),
                num_beams=search_config.get("num_beams", 1),
                num_return_sequences=search_config.get("num_return_sequences", 1),
                repetition_penalty=search_config.get("repetition_penalty", 1.0),
                past_present_share_buffer=search_config.get(
                    "past_present_share_buffer", True
                ),
                random_seed=random_seed,
                # Not currently supported by OGA
                # diversity_penalty=search_config.get('diversity_penalty', 0.0),
                # no_repeat_ngram_size=search_config.get('no_repeat_ngram_size', 0),
            )
        else:
            params.set_search_options(
                do_sample=do_sample,
                top_k=top_k,
                top_p=top_p,
                temperature=temperature,
                max_length=max_length_to_use,
                min_length=min_length,
                random_seed=random_seed,
            )
        params.try_graph_capture_with_max_batch_size(1)

        generator = og.Generator(self.model, params)

        if streamer is None:
            prompt_start_time = time.perf_counter()
            if use_oga_post_6_api:
                generator.append_tokens(input_ids)
            if use_oga_pre_6_api:
                generator.compute_logits()
            generator.generate_next_token()
            prompt_end_time = time.perf_counter()

            self.time_to_first_token = prompt_end_time - prompt_start_time

            if max_new_tokens > 1:

                token_gen_times = []
                while not generator.is_done():
                    token_gen_start_time = time.perf_counter()
                    if use_oga_pre_6_api:
                        generator.compute_logits()
                    generator.generate_next_token()
                    token_gen_end_time = time.perf_counter()

                    token_gen_times.append(token_gen_end_time - token_gen_start_time)

                if token_gen_times:
                    # List will be empty if we generated 1 or 0 tokens, and we don't
                    # want a divide-by-zero error in those cases
                    avg_token_gen_latency_s = sum(token_gen_times) / len(
                        token_gen_times
                    )
                    self.tokens_per_second = 1 / avg_token_gen_latency_s

            return [generator.get_sequence(0)]
        else:
            if use_oga_post_6_api:
                generator.append_tokens(input_ids)
            tokenizer_stream = streamer.tokenizer.tokenizer.create_stream()

            stop_early = False

            while not generator.is_done() and not stop_early:
                if use_oga_pre_6_api:
                    generator.compute_logits()
                generator.generate_next_token()

                new_token = generator.get_next_tokens()[0]
                new_text = tokenizer_stream.decode(new_token)

                streamer.add_text(new_text)

                if stopping_criteria is not None:
                    if stopping_criteria[0].stop_event.is_set():
                        stop_early = True

            streamer.done()

    def _model_call(self, input_ids):
        """
        Run the model on input_ids and get logits.

        This method directly accesses model logits rather than using the full generate pipeline for
        several important reasons:
        1. Purpose: We need raw logits from a single forward pass, while generate() is optimized for
           producing multiple tokens through iterative inference
        2. Efficiency: Direct access is more efficient for logprob calculations with no
           sampling overhead
        3. Precision: Logprob calculations require exact control over input-to-output mapping
        4. Consistency: Similar approach used in both HF and OGA implementations

        Args:
            input_ids: Input token IDs

        Returns:
            Logits for each token in the sequence
        """
        import torch

        # Setup generator params
        params = og.GeneratorParams(self.model)

        # Configure for a simple forward pass
        params.set_search_options(
            do_sample=False,
            temperature=0.0,
            max_length=len(input_ids),
        )

        # Initialize generator
        generator = og.Generator(self.model, params)

        # Feed tokens to model based on API version
        generator.append_tokens(input_ids)

        # Extract logits - this returns a list of logits tensors
        logits = generator.get_output("logits")

        # Convert to torch tensor for easier processing
        return torch.tensor(logits[0])

    def _select_cont_toks(self, logits, context_len, continuation_tokens):
        """
        Select and process logits for continuation tokens.

        Args:
            logits: Full sequence logits
            context_len: Length of context tokens
            continuation_tokens: List or tensor of continuation token IDs

        Returns:
            Log probabilities for continuation tokens
        """
        import torch

        # Extract relevant logits for continuation prediction (shift by one)
        cont_logits = logits[
            context_len - 1 : context_len - 1 + len(continuation_tokens)
        ]

        # Convert to torch tensors if needed
        if not isinstance(continuation_tokens, torch.Tensor):
            continuation_tokens = torch.tensor(continuation_tokens, dtype=torch.long)

        # Apply log softmax to get log probabilities
        log_probs = torch.log_softmax(cont_logits, dim=-1)

        # Get log probs for the specific continuation tokens
        token_log_probs = torch.gather(
            log_probs, 1, continuation_tokens.unsqueeze(-1)
        ).squeeze(-1)

        return token_log_probs

    def compute_logprobs(
        self, text, tokenizer, prompt_length=None, logprobs=None, echo=False
    ):
        """
        Compute log probabilities for all tokens in the given text.

        Args:
            text: The full text to analyze (e.g., prompt + completion)
            prompt_length: Number of tokens in the prompt. If provided and echo=False,
                only completion tokens after this position will be returned.
            logprobs: If not None, return log probabilities. Value indicates how many top
                alternatives to return. If True but not an integer, defaults to 5 alternatives.
            echo: If True, include logprobs for prompt tokens. If False, only return logprobs
                for completion tokens.

        Returns:
            - text_offset: Character offsets for each token in the text
            - token_logprobs: Log probability for each token
            - tokens: The actual tokens used
            - top_logprobs: Top alternative log probabilities for each position
        """
        import torch

        if tokenizer is None:
            raise ValueError("Tokenizer is required for logprob calculation")

        # Encode the full text
        tokens = tokenizer(text).input_ids  # pylint: disable=E1102

        # Track character offsets for each token
        text_offset = []
        start_idx = 0

        token_strings = []
        for token_id in tokens:
            token_str = tokenizer.decode([token_id])
            token_strings.append(token_str)

            # Calculate character offsets for tokens - handles cases where tokens
            # may not directly match in the original text due to encoding differences,
            # special characters, or tokenization artifacts
            try:
                pos = text[start_idx:].find(token_str)
                if pos != -1:
                    text_offset.append(start_idx + pos)
                    start_idx += pos + len(token_str)
                else:
                    text_offset.append(start_idx)
            except (TypeError, ValueError, UnicodeError):
                # Fallback to current position when matching fails due to encoding issues
                text_offset.append(start_idx)

        # Get logits from model
        logits = self._model_call(tokens)

        # Calculate log probabilities for each token
        all_log_probs = torch.log_softmax(logits, dim=-1)

        # The first token doesn't have a conditional probability
        # For tokens after the first, get the predicted probability
        token_log_probs = []
        top_logprobs_list = []

        # For each position, get the actual token probability and top alternatives
        for i in range(len(tokens)):
            # Get previous token position logits
            if i > 0:  # First token has no preceding context
                prev_logits = all_log_probs[i - 1]
                curr_token_id = tokens[i]
                # Get probability of the actual token that appeared
                token_logprob = prev_logits[curr_token_id].item()
                token_log_probs.append(token_logprob)

                # Get top-k alternatives if requested
                if logprobs is not None:
                    num_alternatives = logprobs if isinstance(logprobs, int) else 5
                    topk_values, topk_indices = torch.topk(
                        prev_logits, min(num_alternatives, prev_logits.size(-1))
                    )

                    # Create dictionary of token: logprob
                    position_logprobs = {}
                    for val, idx in zip(topk_values.tolist(), topk_indices.tolist()):
                        token_str = tokenizer.decode([idx])
                        position_logprobs[token_str] = val

                    top_logprobs_list.append(position_logprobs)
            else:
                # For the first token, we don't have a conditional probability
                token_log_probs.append(None)
                top_logprobs_list.append({})

        # If we don't want to echo prompt tokens, filter them out
        if not echo and prompt_length is not None:
            # Ensure prompt_length is within bounds
            prompt_length = min(prompt_length, len(tokens))

            # Filter results to only include completion tokens
            if prompt_length < len(tokens):
                filtered_text_offset = text_offset[prompt_length:]
                filtered_token_logprobs = token_log_probs[prompt_length:]
                filtered_tokens = token_strings[prompt_length:]
                filtered_top_logprobs = top_logprobs_list[prompt_length:]

                return (
                    filtered_text_offset,
                    filtered_token_logprobs,
                    filtered_tokens,
                    filtered_top_logprobs,
                )
            else:
                # No completion tokens
                return [], [], [], []

        return text_offset, token_log_probs, token_strings, top_logprobs_list


class OgaLoad(FirstTool):
    """
    Tool that loads an LLM in OnnxRuntime-GenAI for use with CPU or DirectML execution providers.

    Input: path to a checkpoint.
        Supported choices for cpu and igpu from HF model repository:
            LLM models on Huggingface supported by model_builder.  See documentation
            (https://github.com/lemonade-sdk/lemonade/blob/main/docs/ort_genai_igpu.md)
            for supported models.
        Supported choices for npu from HF model repository:
            Models on Hugging Face that follow the "amd/**-onnx-ryzen-strix" pattern
        Local models for cpu, igpu, or npu:
            The specified checkpoint is converted to a local path, via mapping to lower case
            and replacing '/' with '_'.  If this model already exists in the 'models' folder
            of the lemonade cache and if it has a subfolder <device>-<dtype>, then this model
            will be used.  If the --force flag is used and the model is built with model_builder,
            then it will be rebuilt.



    Output:
        state.model: handle to a Huggingface-style LLM loaded on DirectML device
        state.tokenizer = Huggingface-style LLM tokenizer instance
        state.dtype = data type of the model on DirectML device
        state.checkpoint = name of the checkpoint used to load state.model

    Note: This tool expects the onnxruntime-genai-directml library to be pre-installed.
            If that library is not installed, this tool will not load.
    """

    unique_name = "oga-load"

    def __init__(self):
        super().__init__(monitor_message="Loading OnnxRuntime-GenAI model")

        self.status_stats = [
            Keys.DTYPE,
            Keys.DEVICE,
            Keys.LOCAL_MODEL_FOLDER,
        ]

    @staticmethod
    def parser(add_help: bool = True) -> argparse.ArgumentParser:
        parser = __class__.helpful_parser(
            short_description="Load model in onnxruntime-genai (OGA)",
            add_help=add_help,
        )

        parser.add_argument(
            "-ip",
            "--input_path",
            default="",
            help="the local huggingface model in your disk",
        )

        parser.add_argument(
            "-d",
            "--device",
            choices=["igpu", "npu", "cpu", "hybrid", "cuda"],
            default="igpu",
            help="Which device to load the model on to (default: igpu)",
        )

        parser.add_argument(
            "--dtype",
            choices=["int4", "fp16", "fp32"],
            required=True,
            help="Data type to load the model in",
        )

        parser.add_argument(
            "--int4-block-size",
            default=None,
            help="Specify the block_size for int4 quantization.",
            choices=[16, 32, 64, 128, 256],
            type=int,
        )

        parser.add_argument(
            "--force",
            action="store_true",
            help="Forces downloading of Hugging-Face model again (if changed).  Additionally for"
            " cpu and igpu devices only, forces model_builder to run again on the HF model"
            " (changed or not).",
        )

        parser.add_argument(
            "--download-only",
            action="store_true",
            help="Download the model if needed, but don't load it",
        )

        parser.add_argument(
            "--trust-remote-code",
            action="store_true",
            help="Set this flag to use models whose code is on the Hugging Face hub rather "
            "than natively in the OnnxRuntime Gen AI libraries.  Please review the model code "
            "in advance as this is a security risk.",
        )

        parser.add_argument(
            "--subfolder",
            default=None,
            help="Subfolder where model is located <LEMONADE CACHE>/oga_models/<MODELNAME>"
            "/<SUBFOLDER>, default is <EP for device>-<dtype>.  The EPs are: "
            f'{", ".join([value + " for " + key for key, value in execution_providers.items()])}.',
        )

        return parser

    @staticmethod
    def _validate_model_configuration(device, dtype, checkpoint):
        """
        Validate if the device, dtype, platform and checkpoint combination are consistent with
        HuggingFace checkpoint naming conventions and specifically for AMD models for NPU
        and hybrid flows.

        Returns True if device, dtype, and model are consistent.
        """

        hf_supported_models = {
            "cpu": {"int4": "*/*", "fp32": "*/*"},
            "igpu": {"int4": "*/*", "fp16": "*/*"},
            "npu": {"int4": "*/*"},
            "hybrid": {"int4": "*/*"},
            "cuda": {"int4": "*/*", "fp16": "*/*"},
        }

        hf_supported = (
            device in hf_supported_models
            and dtype in hf_supported_models[device]
            and fnmatch(checkpoint, hf_supported_models[device][dtype])
        )
        return hf_supported

    @staticmethod
    def _setup_model_paths(
        state, checkpoint, device, dtype, subfolder, int4_block_size
    ):
        """
        Determines and returns the following model path information for models produced by OGA
        model builder:

           full_model_path - Full path to where the OGA model files are stored.
           oga_models_subfolder - The subfolder of the oga_models folder where the model files
                are stored.  (<full_model_path> = <oga_models>/<oga_models_subfolder>)
                This subfolder is usually
                  <checkpoint_string>/<device>-<dtype>[-block-<int4_block_size]>
                but the if the argument subfolder is not None it will override the latter part
                of this path.
           model_exists_locally - True if full_model_path is a folder that contains files

        Note: Model files already in ONNX format on Hugging Face will be stored in the
            Hugging Face cache, not this folder.  The <oga_models> folder contains model
            files that have locally been quantized/converted to OGA format and any other
            models that have been manually added by the user.
        """
        if subfolder is None:
            subfolder = f"{execution_providers[device]}-{dtype}"
            subfolder += (
                f"-block-{int4_block_size}"
                if dtype == "int4" and int4_block_size is not None
                else ""
            )

        # First, check in the lemonade oga_models cache
        oga_models_subfolder = os.path.join(
            checkpoint.replace("/", "_").lower(), subfolder
        )
        full_model_path = os.path.join(
            state.cache_dir, oga_models_path, oga_models_subfolder
        )
        model_exists_locally = os.path.isdir(full_model_path) and os.listdir(
            full_model_path
        )

        # If not found in lemonade cache, check in Hugging Face cache
        if not model_exists_locally:
            try:
                snapshot_path = snapshot_download(
                    repo_id=checkpoint,
                    local_files_only=True,
                )

                # Check if the snapshot contains ONNX files
                if os.path.isdir(snapshot_path) and os.listdir(snapshot_path):
                    is_onnx_model = any(
                        filename.endswith(".onnx")
                        for filename in os.listdir(snapshot_path)
                    )

                    if is_onnx_model:
                        # If the model is in HF cache and has ONNX files, use it
                        full_model_path = snapshot_path
                        model_exists_locally = True
                        printing.log_info(
                            f"Found ONNX model in Hugging Face cache: {full_model_path}"
                        )
            except Exception as e:  # pylint: disable=broad-exception-caught
                # Log any errors but continue with the original path
                printing.log_info(f"Error checking Hugging Face cache: {e}")

        return full_model_path, model_exists_locally

    @staticmethod
    def _update_hybrid_custom_ops_library_path(full_model_path):
        """
        Modifies the genai_config.json file in the hybrid model folder to set the custom_ops_library
        path to the location of the onnx_custom_ops.dll in the current environment.
        This is needed for hybrid inference.
        """
        oga_path, version = get_oga_hybrid_dir()

        if "1.3.0" in version:
            custom_ops_path = os.path.join(
                oga_path,
                "onnx_utils",
                "bin",
                "onnx_custom_ops.dll",
            )
        else:
            custom_ops_path = os.path.join(oga_path, "libs", "onnx_custom_ops.dll")

        # Insert the custom_ops_path into the model config file
        config_path = os.path.join(full_model_path, "genai_config.json")
        if os.path.exists(config_path):
            with open(config_path, "r", encoding="utf-8") as f:
                config = json.load(f)

            if (
                "model" in config
                and "decoder" in config["model"]
                and "session_options" in config["model"]["decoder"]
            ):
                config["model"]["decoder"]["session_options"][
                    "custom_ops_library"
                ] = custom_ops_path

            with open(config_path, "w", encoding="utf-8") as f:
                json.dump(config, f, indent=4)

        else:
            printing.log_info(
                f"Model's `genai_config.json` not found in {full_model_path}"
            )

    @staticmethod
    def _is_preoptimized_model(input_model_path):
        """
        Checks if the 'custom_ops_library' field exists in the genai_config.json file
        to determine if this is a pre-optimized model for hybrid as well
        as NPU only.

        Args:
            input_model_path (str): Path to the input model directory.

        Returns:
            bool: True if 'custom_ops_library' exists, False otherwise.
        """
        config_path = os.path.join(input_model_path, "genai_config.json")
        if not os.path.exists(config_path):
            printing.log_info(f"Model's `genai_config.json` not found in {config_path}")
            return False

        with open(config_path, "r", encoding="utf-8") as f:
            config = json.load(f)
        if (
            "model" in config
            and "decoder" in config["model"]
            and "session_options" in config["model"]["decoder"]
        ):
            return "custom_ops_library" in config["model"]["decoder"]["session_options"]
        return False

    @staticmethod
    def _download_and_build_safetensors_model(
        checkpoint, device, dtype, full_model_path, int4_block_size, input_path, state
    ):
        """
        Uses OGA model builder to quantize safetensors format model and convert to ONNX
        format.  The model files are saved to the full_model_path folder.
        """
        printing.log_info(f"Building {checkpoint} for {device} using {dtype}")
        extra_options = {}
        if int4_block_size is not None:
            extra_options["int4-block-size"] = int4_block_size
        try:
            model_builder.create_model(
                checkpoint,
                input_path,
                full_model_path,
                dtype,
                execution_providers[device],
                os.path.join(state.cache_dir, oga_model_builder_cache_path),
                **extra_options,
            )
        except NotImplementedError as e:
            raise NotImplementedError("[Model builder] " + str(e)) from e
        except OSError as e:
            raise ValueError("[Model builder] " + str(e)) from e

        return full_model_path

    @staticmethod
    def _setup_npu_environment():
        """
        Sets up environment for NPU flow of ONNX model and returns saved state to be restored
        later in cleanup.
        """
        oga_path, version = get_oga_npu_dir()

        if not os.path.exists(os.path.join(oga_path, "libs", "onnxruntime.dll")):
            raise RuntimeError(
                f"Cannot find libs/onnxruntime.dll in lib folder: {oga_path}"
            )

        # Save current state so they can be restored after inference.
        saved_state = {"cwd": os.getcwd(), "path": os.environ["PATH"]}

        # Setup NPU environment (cwd and path will be restored later)
        os.chdir(oga_path)
        os.environ["PATH"] = (
            os.path.join(oga_path, "libs") + os.pathsep + os.environ["PATH"]
        )
        if "1.3.0" in version:
            os.environ["DD_ROOT"] = ".\\bins"
            os.environ["DEVICE"] = "stx"
            os.environ["XLNX_ENABLE_CACHE"] = "0"

        return saved_state

    @staticmethod
    def _setup_hybrid_environment():
        """
        Sets up the environment for the Hybrid flow and returns saved state to be restored later
        in cleanup.
        """
        # Determine the Ryzen AI OGA version and hybrid artifacts path
        oga_path, version = get_oga_hybrid_dir()

        if "1.3.0" in version:
            dst_dll = os.path.join(
                oga_path,
                "onnx_utils",
                "bin",
                "DirectML.dll",
            )
            if not os.path.isfile(dst_dll):
                # Artifacts 1.3.0 has DirectML.dll in different subfolder, so copy it to the
                # correct place.  This should not be needed in later RAI release artifacts.
                src_dll = os.path.join(
                    oga_path,
                    "onnxruntime_genai",
                    "lib",
                    "DirectML.dll",
                )
                os.makedirs(os.path.dirname(dst_dll), exist_ok=True)
                shutil.copy2(src_dll, dst_dll)

        saved_state = None
        return saved_state

    @staticmethod
    def _load_model_and_setup_state(
        state, full_model_path, checkpoint, trust_remote_code
    ):
        """
        Loads the OGA model from local folder and then loads the tokenizer.
        Will auto-detect if we're offline.
        """
        try:
            state.model = OrtGenaiModel(full_model_path)
        except Exception as e:
            if "invalid unordered_map<K, T>" in str(e):
                raise ValueError(
                    "Error initializing model: Invalid configuration detected.\n"
                    "Please check the following:\n"
                    f"1. Please check your model's config file in {full_model_path} "
                    "and ensure custom_ops_library points to the valid "
                    "onnx_custom_ops.dll path.\n"
                    "2. Make sure the NPU driver is loaded.\n"
                    "3. Make sure hybrid has been installed on a Ryzen AI "
                    f"{'or '.join(SUPPORTED_RYZEN_AI_SERIES)}-series processor."
                ) from e
            raise

        # Auto-detect offline mode
        offline = is_offline()

        try:
            # Always try to use local files first
            local_files_only = True

            hf_tokenizer = AutoTokenizer.from_pretrained(
                full_model_path,
                local_files_only=local_files_only,
                trust_remote_code=trust_remote_code,
            )
        except ValueError as e:
            if "trust_remote_code" in str(e):
                raise ValueError(
                    "This model requires you to execute code from the repo.  Please review it "
                    "and if you trust it, then use the `--trust-remote-code` flag with oga-load."
                )

            if offline and "Can't load tokenizer for" in str(e):
                raise ValueError(
                    f"Cannot load tokenizer for {checkpoint} in offline mode. "
                    f"The tokenizer files may not be available locally in {full_model_path}."
                )
            raise

        state.tokenizer = OrtGenaiTokenizer(
            state.model.model,
            hf_tokenizer,
        )

        status.add_to_state(state=state, name=checkpoint, model=checkpoint)

    @staticmethod
    def _cleanup_environment(saved_state):
        """
        Restores environment to its original state after inference is complete.
        """
        if saved_state:
            os.chdir(saved_state["cwd"])
            os.environ["PATH"] = saved_state["path"]

    def _generate_model_for_hybrid_or_npu(
        self, output_model_path, device, input_model_path
    ):
        """
        Uses a subprocess to run the 'model_generate' command for hybrid or npu devices.
        """

        # Determine the appropriate flag based on the device type
        if device == "hybrid":
            device_flag = "--hybrid"
        elif device == "npu":
            device_flag = "--npu"
        else:
            raise ValueError(f"Unsupported device type for model generation: {device}")

        command = [
            "model_generate",
            device_flag,
            output_model_path,  # Output model directory
            input_model_path,  # Input model directory
        ]

        printing.log_info(f"Running command: {' '.join(command)}")
        try:
            with open(self.logfile_path, "w", encoding="utf-8") as log_file:
                subprocess.run(
                    command, check=True, text=True, stdout=log_file, stderr=log_file
                )
        except FileNotFoundError as e:
            error_message = (
                "The 'model_generate' package is missing from your system. "
                "Ensure all required packages are installed. "
                "To install it, run the following command:\n\n"
                "    lemonade-install --ryzenai <target> --build-model\n"
            )
            raise RuntimeError(error_message) from e

    def run(
        self,
        state: State,
        input: str,
        input_path: str = "",
        device: str = "igpu",
        dtype: str = "int4",
        int4_block_size: int = None,
        force: bool = False,
        download_only: bool = False,
        trust_remote_code=False,
        subfolder: str = None,
    ) -> State:
        # Auto-detect offline status
        offline = is_offline()
        if offline:
            printing.log_warning(
                "Network connectivity to huggingface.co not detected. Running in offline mode."
            )

        state.device = device
        state.dtype = dtype

        # Log initial stats
        state.save_stat(Keys.DTYPE, dtype)
        state.save_stat(Keys.DEVICE, device)
        if device in ["hybrid", "npu"]:
            ryzen_ai_version_info = get_ryzen_ai_version_info()
            state.save_stat(Keys.RYZEN_AI_VERSION_INFO, ryzen_ai_version_info)

        # Check if input is a local folder
        if os.path.isdir(input):
            # input is a local folder
            full_model_path = os.path.abspath(input)
            checkpoint = "local_model"
            state.checkpoint = checkpoint
            state.save_stat(Keys.CHECKPOINT, checkpoint)
            state.save_stat(Keys.LOCAL_MODEL_FOLDER, full_model_path)
            # See if there is a file ending in ".onnx" in this folder
            dir = os.listdir(input)
            has_onnx_file = any([filename.endswith(".onnx") for filename in dir])
            if not has_onnx_file:
                raise ValueError(
                    f"The folder {input} does not contain an ONNX model file."
                )
            if force:
                raise ValueError(
                    "Your input (-i, --input) points to a local folder, which is not "
                    "compatible with the force argument."
                )

        else:
            # input is a model checkpoint
            checkpoint = input
            state.checkpoint = checkpoint
            state.save_stat(Keys.CHECKPOINT, checkpoint)

            # Get base model information
            if not offline:
                base_model = get_base_model(checkpoint)
                if base_model is not None:
                    state.save_stat("base_model", base_model)

            # Setup paths
            full_model_path, model_exists_locally = self._setup_model_paths(
                state, checkpoint, device, dtype, subfolder, int4_block_size
            )

            # If in offline mode, we can only use locally available models
            if offline and not model_exists_locally:
                raise ValueError(
                    f"Model {checkpoint} is not available locally for {device} with {dtype}. "
                    f"Cannot download in offline mode. Check {full_model_path}"
                )

            # Handle download/build if needed
            if (not model_exists_locally) or force:
                if offline:
                    raise ValueError(
                        f"Cannot download or build model {checkpoint} in offline mode"
                    )

                # Validate configuration
                hf_supported = self._validate_model_configuration(
                    device, dtype, checkpoint
                )

                if not hf_supported:
                    raise ValueError(
                        "The (device, dtype, checkpoint) combination is not supported: "
                        f"({device}, {dtype}, {checkpoint})"
                    )
                input_model_path = snapshot_download(
                    repo_id=checkpoint,
                    ignore_patterns=["*.md", "*.txt"],
                    local_files_only=offline,
                )
                # Check if model is ONNX or safetensors
                is_onnx_model = any(
                    [
                        filename.endswith(".onnx")
                        for filename in os.listdir(input_model_path)
                    ]
                )
                is_preoptimized_onnx = is_onnx_model and self._is_preoptimized_model(
                    input_model_path
                )
                is_safetensors_model = any(
                    [
                        filename.endswith(".safetensors")
                        for filename in os.listdir(input_model_path)
                    ]
                )
                if not (is_onnx_model or is_safetensors_model):
                    raise ValueError(
                        f"The model {checkpoint} is not supported. "
                        "It does not contain ONNX or safetensors files."
                    )
                if device in ["npu", "hybrid"]:
                    if is_onnx_model:
                        if is_preoptimized_onnx:
                            # Use HuggingFace cache path as it is
                            full_model_path = input_model_path
                        else:
                            # If ONNX but not modified yet for Hybrid or NPU,
                            # needs further optimization
                            self._generate_model_for_hybrid_or_npu(
                                full_model_path,
                                device,
                                input_model_path,
                            )
                    elif is_safetensors_model:
                        config_path = os.path.join(input_model_path, "config.json")
                        if os.path.exists(config_path):
                            with open(config_path, "r", encoding="utf-8") as f:
                                config = json.load(f)
                            if "quantization_config" in config:
                                # If quantized, use subprocess to generate the model
                                self._generate_model_for_hybrid_or_npu(
                                    full_model_path, device, input_model_path
                                )
                            else:
                                raise ValueError(
                                    f"The safetensors model {checkpoint} is not quantized. "
                                    "Only quantized safetensors models are supported"
                                    " on npu or hybrid targets."
                                )
                        else:
                            raise ValueError(
                                f"config.json not found for safetensors model: {checkpoint}"
                            )
                    else:
                        raise ValueError(
                            f"Unsupported model type for checkpoint: {checkpoint}"
                        )
                else:
                    if is_onnx_model:
                        # Use HuggingFace cache path as it is
                        full_model_path = input_model_path
                    else:
                        self._download_and_build_safetensors_model(
                            checkpoint,
                            device,
                            dtype,
                            full_model_path,
                            int4_block_size,
                            input_path,
                            state,
                        )
                    state.save_stat(Keys.LOCAL_MODEL_FOLDER, full_model_path)

        # Load model if download-only argument is not set
        if not download_only:

            saved_env_state = None
            try:
                if device == "npu":
                    saved_env_state = self._setup_npu_environment()
                    # Set USE_AIE_RoPE based on model type
                    os.environ["USE_AIE_RoPE"] = (
                        "0" if "phi-" in checkpoint.lower() else "1"
                    )
                elif device == "hybrid":
                    saved_env_state = self._setup_hybrid_environment()
                    self._update_hybrid_custom_ops_library_path(full_model_path)

                self._load_model_and_setup_state(
                    state, full_model_path, checkpoint, trust_remote_code
                )
            finally:
                self._cleanup_environment(saved_env_state)

        return state


# This file was originally licensed under Apache 2.0. It has been modified.
# Modifications Copyright (c) 2025 AMD
