import ctypes
import os
import sys

import numpy
import ailia

import os
import platform

#### dependency check
if sys.platform == "win32":
    import ctypes
    if platform.machine() != "ARM64":
        try:
            for library in ["vcruntime140.dll", "vcruntime140_1.dll", "msvcp140.dll"]:
                ctypes.windll.LoadLibrary(library)
        except:
            print("  WARNING Please install MSVC 2015-2019 runtime from https://docs.microsoft.com/ja-jp/cpp/windows/latest-supported-vc-redist")


#### loading DLL / DYLIB / SO  ####
if sys.platform == "win32":
    if platform.machine() == "ARM64":
        dll_platform = "windows/arm64"
    else:
        dll_platform = "windows/x64"
    dll_name = "ailia_llm.dll"
    dll_name_fallback = "ailia_llm_fallback.dll"
    load_fn = ctypes.WinDLL
elif sys.platform == "darwin":
    dll_platform = "mac"
    dll_name = "libailia_llm.dylib"
    dll_name_fallback = None
    load_fn = ctypes.CDLL
else:
    is_arm = "arm" in platform.machine() or platform.machine() == "aarch64"
    if is_arm:
        if platform.architecture()[0] == "32bit":
            dll_platform = "linux/armeabi-v7a"
        else:
            dll_platform = "linux/arm64-v8a"
    else:
        dll_platform = "linux/x64"
    dll_name = "libailia_llm.so"
    dll_name_fallback = None
    load_fn = ctypes.CDLL

dll_found = False
candidate = ["", str(os.path.dirname(os.path.abspath(__file__))) + str(os.sep), str(os.path.dirname(os.path.abspath(__file__))) + str(os.sep) + dll_platform + str(os.sep)]

for dir in candidate:
    try:
        dll = load_fn(dir + dll_name) # Load gpu library first
        dll_found = True
    except:
        pass

if not dll_found:
    for dir in candidate:
        try:
            dll = load_fn(dir + dll_name_fallback) # Fallback cpu library second
            dll_found = True
        except:
            pass

if not dll_found:
    msg = "DLL load failed : \'" + dll_name + "\' is not found"
    raise ImportError(msg)

# ==============================================================================

from ctypes import *

AILIA_LLM_STATUS_SUCCESS = ( 0 )
AILIA_LLM_STATUS_CONTEXT_FULL = ( -8 )

class AILIALLMChatMessage(ctypes.Structure):
    """Structure representing a single chat message for text-only prompts.

    Attributes
    ----------
    role : bytes
        The role of the message sender (e.g., b"system", b"user", b"assistant").
    content : bytes
        The text content of the message in UTF-8 encoding.
    """
    _fields_ = [
        ("role", ctypes.c_char_p),
        ("content", ctypes.c_char_p)
    ]

class AILIALLMMediaData(ctypes.Structure):
    """Structure representing media data for multimodal prompts.

    Attributes
    ----------
    media_type : bytes
        The type of media (e.g., b"image").
    file_path : bytes
        File path to the media resource in UTF-8 encoding.
    data : ctypes.POINTER(ctypes.c_ubyte)
        Pointer to raw media data buffer. Can be None if file_path is used.
    data_size : int
        Size of raw media data in bytes.
    width : int
        Width of the media in pixels (0 if not applicable).
    height : int
        Height of the media in pixels (0 if not applicable).
    """
    _fields_ = [
        ("media_type", ctypes.c_char_p),
        ("file_path", ctypes.c_char_p),
        ("data", ctypes.POINTER(ctypes.c_ubyte)),
        ("data_size", ctypes.c_uint),
        ("width", ctypes.c_uint),
        ("height", ctypes.c_uint)
    ]

class AILIALLMMultimodalChatMessage(ctypes.Structure):
    """Structure representing a single chat message with optional media attachments.

    Attributes
    ----------
    role : bytes
        The role of the message sender (e.g., b"system", b"user", b"assistant").
    content : bytes
        The text content of the message in UTF-8 encoding.
    media_data : ctypes.POINTER(AILIALLMMediaData)
        Pointer to an array of media data attached to this message.
        Can be None if no media is attached.
    media_count : int
        Number of media data entries attached to this message.
    """
    _fields_ = [
        ("role", ctypes.c_char_p),
        ("content", ctypes.c_char_p),
        ("media_data", ctypes.POINTER(AILIALLMMediaData)),
        ("media_count", ctypes.c_uint)
    ]

# ==============================================================================

dll.ailiaLLMCreate.restype = c_int
dll.ailiaLLMCreate.argtypes = (POINTER(c_void_p), )

dll.ailiaLLMDestroy.restype = None
dll.ailiaLLMDestroy.argtypes = (c_void_p, )

dll.ailiaLLMOpenModelFileA.restype = c_int
dll.ailiaLLMOpenModelFileA.argtypes = (c_void_p, c_char_p, c_uint32)

dll.ailiaLLMOpenModelFileW.restype = c_int
dll.ailiaLLMOpenModelFileW.argtypes = (c_void_p, c_wchar_p, c_uint32)

dll.ailiaLLMGetDeltaTextSize.restype = c_int
dll.ailiaLLMGetDeltaTextSize.argtypes = (c_void_p, POINTER(c_uint32))

dll.ailiaLLMGetDeltaText.restype = c_int
dll.ailiaLLMGetDeltaText.argtypes = (c_void_p, numpy.ctypeslib.ndpointer(
                dtype=numpy.uint8, flags='CONTIGUOUS'
            ),                               # text
            ctypes.c_uint)

dll.ailiaLLMSetSamplingParams.restype = c_int
dll.ailiaLLMSetSamplingParams.argtypes = (c_void_p, ctypes.c_uint, ctypes.c_float, ctypes.c_float, ctypes.c_uint)

dll.ailiaLLMSetPrompt.restype = c_int
dll.ailiaLLMSetPrompt.argtypes = (c_void_p, ctypes.POINTER(AILIALLMChatMessage), ctypes.c_uint)

dll.ailiaLLMGenerate.restype = c_int
dll.ailiaLLMGenerate.argtypes = (c_void_p, POINTER(c_uint32))

dll.ailiaLLMGetTokenCount.restype = c_int
dll.ailiaLLMGetTokenCount.argtypes = (c_void_p, POINTER(c_uint32), c_char_p)

dll.ailiaLLMGetPromptTokenCount.restype = c_int
dll.ailiaLLMGetPromptTokenCount.argtypes = (c_void_p, POINTER(c_uint32))

dll.ailiaLLMGetGeneratedTokenCount.restype = c_int
dll.ailiaLLMGetGeneratedTokenCount.argtypes = (c_void_p, POINTER(c_uint32))

dll.ailiaLLMOpenMultimodalProjectorFileA.restype = c_int
dll.ailiaLLMOpenMultimodalProjectorFileA.argtypes = (c_void_p, c_char_p)

dll.ailiaLLMOpenMultimodalProjectorFileW.restype = c_int
dll.ailiaLLMOpenMultimodalProjectorFileW.argtypes = (c_void_p, c_wchar_p)

dll.ailiaLLMGetMultimodalCapabilities.restype = c_int
dll.ailiaLLMGetMultimodalCapabilities.argtypes = (c_void_p, POINTER(c_uint32), POINTER(c_uint32))

dll.ailiaLLMSetMultimodalPrompt.restype = c_int
dll.ailiaLLMSetMultimodalPrompt.argtypes = (c_void_p, ctypes.POINTER(AILIALLMMultimodalChatMessage), ctypes.c_uint)

dll.ailiaLLMSetThinking.restype = c_int
dll.ailiaLLMSetThinking.argtypes = (c_void_p, ctypes.c_uint)


# ==============================================================================
# base model class
# ==============================================================================

class AiliaLLMError(RuntimeError):
    """Exception raised when an ailia LLM operation fails.

    Attributes
    ----------
    code : int
        The native error code returned by the ailia LLM library.
    """
    def __init__(self, message, code):
        """Initialize AiliaLLMError.

        Parameters
        ----------
        message : str
            Human-readable error description.
        code : int
            Native error code from the ailia LLM library.
        """
        super().__init__(f"{message} code:{code}")
        self.code = code

class AiliaLLM():
    """Main class for ailia LLM inference.

    Provides text generation using large language models via the ailia LLM
    native backend. Supports both text-only and multimodal (vision/audio)
    generation.

    Examples
    --------
    >>> import ailia_llm
    >>> llm = ailia_llm.AiliaLLM()
    >>> llm.open("model.gguf")
    >>> for delta in llm.generate([{"role": "user", "content": "Hello!"}]):
    ...     print(delta, end="", flush=True)
    """
    _instance = None
    _context_full = False
    _multimodal_projector_opened = False

    def _check(self, status):
        if status != AILIA_LLM_STATUS_SUCCESS:
            raise AiliaLLMError(f"ailia LLM error", status)

    def _string_buffer_aw(self, path):
        if sys.platform == "win32":
            return ctypes.create_unicode_buffer(path)
        else:
            return ctypes.create_string_buffer(path.encode("utf-8"))

    def _string_buffer(self, path):
        return ctypes.create_string_buffer(path.encode("utf-8"))

    def _string_to_c_char_p(self, s):
        return ctypes.c_char_p(s.encode('utf-8'))

    def __init__(self):
        """Create an ailia LLM instance.

        Allocates a native LLM context. Call :meth:`open` to load a model
        before generating text.

        Raises
        ------
        AiliaLLMError
            If the native context creation fails.
        """
        self._instance = ctypes.c_void_p(None)
        self._check(dll.ailiaLLMCreate(cast(pointer(self._instance), POINTER(c_void_p))))

    def open(self, model_path, n_ctx = 0):
        """Open and load a language model file.

        Parameters
        ----------
        model_path : str
            Path to the GGUF model file.
        n_ctx : int, optional, default=0
            Context window size (number of tokens).
            When 0, the default context size of the model is used.

        Raises
        ------
        AiliaLLMError
            If the model file cannot be opened or loaded.

        Examples
        --------
        >>> import ailia_llm
        >>> llm = ailia_llm.AiliaLLM()
        >>> llm.open("path/to/model.gguf")
        >>> llm.open("path/to/model.gguf", n_ctx=4096)
        """
        if "time_license" in ailia.get_version():
            ailia.check_and_download_license()

        # Reset multimodal projector state when opening a new model
        self._multimodal_projector_opened = False

        p1 = self._string_buffer_aw(model_path)

        if sys.platform == "win32":
            self._check(dll.ailiaLLMOpenModelFileW(self._instance, p1, ctypes.c_uint32(n_ctx)))
        else:
            self._check(dll.ailiaLLMOpenModelFileA(self._instance, p1, ctypes.c_uint32(n_ctx)))

    def open_multimodal_projector(self, mmproj_path):
        """Open a multimodal projector file for vision or audio support.

        Must be called after :meth:`open` to enable multimodal capabilities
        such as image or audio understanding.

        Parameters
        ----------
        mmproj_path : str
            Path to the multimodal projector file (e.g., mmproj-model.gguf).

        Raises
        ------
        AiliaLLMError
            If the projector file cannot be opened or loaded.

        Examples
        --------
        >>> import ailia_llm
        >>> llm = ailia_llm.AiliaLLM()
        >>> llm.open("model.gguf")
        >>> llm.open_multimodal_projector("mmproj-model.gguf")
        """
        p1 = self._string_buffer_aw(mmproj_path)

        if sys.platform == "win32":
            self._check(dll.ailiaLLMOpenMultimodalProjectorFileW(self._instance, p1))
        else:
            self._check(dll.ailiaLLMOpenMultimodalProjectorFileA(self._instance, p1))

        self._multimodal_projector_opened = True

    def get_multimodal_capabilities(self):
        """Query the multimodal capabilities of the loaded model.

        Returns
        -------
        dict
            A dictionary with the following keys:
                - ``"vision"`` (bool): Whether the model supports image input.
                - ``"audio"`` (bool): Whether the model supports audio input.

        Raises
        ------
        AiliaLLMError
            If the capability query fails.

        Examples
        --------
        >>> import ailia_llm
        >>> llm = ailia_llm.AiliaLLM()
        >>> llm.open("model.gguf")
        >>> llm.open_multimodal_projector("mmproj-model.gguf")
        >>> caps = llm.get_multimodal_capabilities()
        >>> caps["vision"]
        True
        """
        vision_support = ctypes.c_uint(0)
        audio_support = ctypes.c_uint(0)
        self._check(dll.ailiaLLMGetMultimodalCapabilities(self._instance, ctypes.byref(vision_support), ctypes.byref(audio_support)))
        return {"vision": bool(vision_support.value), "audio": bool(audio_support.value)}

    def _has_media_data(self, prompts):
        """Check if any message in prompts contains media_data.

        Parameters
        ----------
        prompts : list[dict]
            A list of message dictionaries.

        Returns
        -------
        bool
            True if any message contains a non-empty media_data key.
        """
        for prompt in prompts:
            if "media_data" in prompt and prompt["media_data"]:
                return True
        return False

    def set_thinking(self, enable):
        """Enable or disable thinking (reasoning output) for thinking models.

        Parameters
        ----------
        enable : bool
            True to enable thinking output, False to disable it.

        Raises
        ------
        AiliaLLMError
            If the operation fails.
        """
        self._check(dll.ailiaLLMSetThinking(self._instance, ctypes.c_uint(1 if enable else 0)))

    def generate(self, prompts, top_k = 40, top_p = 0.9, temp = 0.4, dist = 1234):
        """Generate text from chat-style prompts.

        This is a unified generator function that automatically detects whether
        the prompts contain media_data and routes to the appropriate internal
        API. It yields text fragments (delta tokens) as they are produced by
        the model.

        Parameters
        ----------
        prompts : list[dict]
            A list of message dictionaries, each containing:
                - ``"role"`` (str): The role (e.g., "system", "user", "assistant").
                - ``"content"`` (str): The text content of the message.
                - ``"media_data"`` (list[dict], optional): A list of media entries,
                  each containing:
                    - ``"media_type"`` (str): Type of media (e.g., "image").
                    - ``"file_path"`` (str, optional): Path to the media file.
                    - ``"data"`` (bytes, optional): Raw media data.
                    - ``"width"`` (int, optional): Media width in pixels.
                    - ``"height"`` (int, optional): Media height in pixels.
        top_k : int, optional, default=40
            Top-k sampling parameter. Limits the token candidates
            to the top-k most probable tokens.
        top_p : float, optional, default=0.9
            Top-p (nucleus) sampling parameter. Limits the token
            candidates to those within the cumulative probability p.
        temp : float, optional, default=0.4
            Sampling temperature. Higher values produce more random output.
        dist : int, optional, default=1234
            Random seed for reproducible generation.

        Yields
        ------
        str
            Delta text fragments as they are generated by the model.
            Concatenating all yielded strings produces the full response.

        Raises
        ------
        AiliaLLMError
            If prompt setting or generation fails.
        RuntimeError
            If media_data is provided but multimodal projector is not loaded.

        Notes
        -----
        If the context window becomes full during generation,
        the generator stops early. Use :meth:`context_full` to check
        whether the context was exhausted.

        For multimodal generation with media_data, you must first call
        :meth:`open_multimodal_projector` to load the projector file.

        Examples
        --------
        Text-only generation:

        >>> import ailia_llm
        >>> llm = ailia_llm.AiliaLLM()
        >>> llm.open("model.gguf")
        >>> messages = [
        ...     {"role": "system", "content": "You are a helpful assistant."},
        ...     {"role": "user", "content": "What is Python?"}
        ... ]
        >>> for delta in llm.generate(messages):
        ...     print(delta, end="", flush=True)

        Multimodal generation with images:

        >>> import ailia_llm
        >>> llm = ailia_llm.AiliaLLM()
        >>> llm.open("model.gguf")
        >>> llm.open_multimodal_projector("mmproj-model.gguf")
        >>> messages = [
        ...     {"role": "user", "content": "Describe this image: <__media__>",
        ...      "media_data": [{"media_type": "image", "file_path": "photo.jpg"}]}
        ... ]
        >>> for delta in llm.generate(messages):
        ...     print(delta, end="", flush=True)
        """
        has_media = self._has_media_data(prompts)

        # If media_data exists, check that the multimodal projector is loaded
        if has_media:
            if not self._multimodal_projector_opened:
                raise RuntimeError(
                    "media_data was provided but multimodal projector is not loaded. "
                    "Call open_multimodal_projector() first to enable multimodal generation."
                )
            # Use multimodal path
            yield from self._generate_multimodal_internal(prompts, top_k, top_p, temp, dist)
        else:
            # Use text-only path
            yield from self._generate_text_internal(prompts, top_k, top_p, temp, dist)

    def _generate_text_internal(self, prompts, top_k, top_p, temp, dist):
        """Internal text-only generation implementation.

        Parameters
        ----------
        prompts : list[dict]
            A list of message dictionaries with role and content.
        top_k : int
            Top-k sampling parameter.
        top_p : float
            Top-p sampling parameter.
        temp : float
            Sampling temperature.
        dist : int
            Random seed.

        Yields
        ------
        str
            Delta text fragments.
        """
        messages = (AILIALLMChatMessage * len(prompts))()
        for i in range(len(prompts)):
            messages[i].role = self._string_to_c_char_p(prompts[i]["role"])
            messages[i].content = self._string_to_c_char_p(prompts[i]["content"])

        status = dll.ailiaLLMSetSamplingParams(self._instance, top_k, top_p, temp, dist)
        self._check(status)

        self._context_full = False
        status = dll.ailiaLLMSetPrompt(self._instance, messages, ctypes.c_uint32(len(prompts)))
        if status == AILIA_LLM_STATUS_CONTEXT_FULL:
            self._context_full = True
            return
        self._check(status)

        buf = numpy.zeros((0), dtype=numpy.uint8, order='C')
        before_text = ""

        while True:
            done = ctypes.c_uint(0)
            status = dll.ailiaLLMGenerate(self._instance, ctypes.byref(done))

            if status == AILIA_LLM_STATUS_CONTEXT_FULL:
                self._context_full = True
                break

            if done.value == 1:
                break

            self._check(status)

            count = ctypes.c_uint(0)
            self._check(dll.ailiaLLMGetDeltaTextSize(self._instance, ctypes.byref(count)))

            new_buf = numpy.zeros((count.value), dtype=numpy.uint8, order='C')

            self._check(dll.ailiaLLMGetDeltaText(self._instance, new_buf, count))

            buf = numpy.concatenate([buf[:-1], new_buf])

            text = bytes(buf[:- 1]).decode("utf-8", errors="ignore")
            delta_text = text[len(before_text):]

            yield delta_text
            before_text = text

    def _generate_multimodal_internal(self, prompts, top_k, top_p, temp, dist):
        """Internal multimodal generation implementation.

        Parameters
        ----------
        prompts : list[dict]
            A list of message dictionaries with role, content, and optional media_data.
        top_k : int
            Top-k sampling parameter.
        top_p : float
            Top-p sampling parameter.
        temp : float
            Sampling temperature.
        dist : int
            Random seed.

        Yields
        ------
        str
            Delta text fragments.
        """
        messages = (AILIALLMMultimodalChatMessage * len(prompts))()

        # Keep references to prevent garbage collection
        media_data_arrays = []

        for i in range(len(prompts)):
            messages[i].role = self._string_to_c_char_p(prompts[i]["role"])
            messages[i].content = self._string_to_c_char_p(prompts[i]["content"])

            if "media_data" in prompts[i] and prompts[i]["media_data"]:
                media_list = prompts[i]["media_data"]
                media_count = len(media_list)

                media_array = (AILIALLMMediaData * media_count)()
                media_data_arrays.append(media_array)  # Keep reference

                for j, media in enumerate(media_list):
                    media_array[j].media_type = self._string_to_c_char_p(media.get("media_type", "image"))
                    if "file_path" in media:
                        media_array[j].file_path = self._string_to_c_char_p(media["file_path"])
                    elif "data" in media:
                        # Handle raw data (not implemented in this basic version)
                        media_array[j].data = None
                        media_array[j].data_size = 0
                    media_array[j].width = media.get("width", 0)
                    media_array[j].height = media.get("height", 0)

                messages[i].media_data = ctypes.cast(media_array, ctypes.POINTER(AILIALLMMediaData))
                messages[i].media_count = media_count
            else:
                messages[i].media_data = None
                messages[i].media_count = 0

        status = dll.ailiaLLMSetSamplingParams(self._instance, top_k, top_p, temp, dist)
        self._check(status)

        self._context_full = False
        status = dll.ailiaLLMSetMultimodalPrompt(self._instance, messages, ctypes.c_uint32(len(prompts)))
        if status == AILIA_LLM_STATUS_CONTEXT_FULL:
            self._context_full = True
            return
        self._check(status)

        buf = numpy.zeros((0), dtype=numpy.uint8, order='C')
        before_text = ""

        while True:
            done = ctypes.c_uint(0)
            status = dll.ailiaLLMGenerate(self._instance, ctypes.byref(done))

            if status == AILIA_LLM_STATUS_CONTEXT_FULL:
                self._context_full = True
                break

            if done.value == 1:
                break

            self._check(status)

            count = ctypes.c_uint(0)
            self._check(dll.ailiaLLMGetDeltaTextSize(self._instance, ctypes.byref(count)))

            new_buf = numpy.zeros((count.value), dtype=numpy.uint8, order='C')

            self._check(dll.ailiaLLMGetDeltaText(self._instance, new_buf, count))

            buf = numpy.concatenate([buf[:-1], new_buf])

            text = bytes(buf[:- 1]).decode("utf-8", errors="ignore")
            delta_text = text[len(before_text):]

            yield delta_text
            before_text = text

    def context_full(self):
        """Check whether the context window is full.

        Returns
        -------
        bool
            True if the context window was exhausted during the last
            :meth:`generate` or :meth:`generate_multimodal` call.

        Examples
        --------
        >>> for delta in llm.generate(messages):
        ...     print(delta, end="")
        >>> if llm.context_full():
        ...     print("Warning: context window was full.")
        """
        return self._context_full

    def token_count(self, text):
        """Count the number of tokens in the given text.

        Parameters
        ----------
        text : str
            Input text to tokenize and count.

        Returns
        -------
        int
            Number of tokens in the text.

        Raises
        ------
        AiliaLLMError
            If the token counting operation fails.

        Examples
        --------
        >>> import ailia_llm
        >>> llm = ailia_llm.AiliaLLM()
        >>> llm.open("model.gguf")
        >>> llm.token_count("Hello, world!")
        4
        """
        count = ctypes.c_uint(0)
        text = self._string_to_c_char_p(text)
        self._check(dll.ailiaLLMGetTokenCount(self._instance, ctypes.byref(count), text))
        return count.value

    def prompt_token_count(self):
        """Get the token count of the most recently set prompt.

        Returns
        -------
        int
            Number of tokens in the prompt that was last passed to
            :meth:`generate` or :meth:`generate_multimodal`.

        Raises
        ------
        AiliaLLMError
            If the token count retrieval fails.

        Examples
        --------
        >>> for delta in llm.generate(messages):
        ...     print(delta, end="")
        >>> llm.prompt_token_count()
        42
        """
        count = ctypes.c_uint(0)
        self._check(dll.ailiaLLMGetPromptTokenCount(self._instance, ctypes.byref(count)))
        return count.value

    def generated_token_count(self):
        """Get the token count of the most recently generated output.

        Returns
        -------
        int
            Number of tokens generated in the last
            :meth:`generate` or :meth:`generate_multimodal` call.

        Raises
        ------
        AiliaLLMError
            If the token count retrieval fails.

        Examples
        --------
        >>> for delta in llm.generate(messages):
        ...     print(delta, end="")
        >>> llm.generated_token_count()
        128
        """
        count = ctypes.c_uint(0)
        self._check(dll.ailiaLLMGetGeneratedTokenCount(self._instance, ctypes.byref(count)))
        return count.value

    def __del__(self):
        if self._instance:
            dll.ailiaLLMDestroy(cast(self._instance, c_void_p))
            self._multimodal_projector_opened = False

