from __future__ import annotations

import asyncio
import inspect
import logging
from collections.abc import Mapping, MutableMapping, Sequence
from contextlib import AbstractAsyncContextManager
from copy import deepcopy
from importlib import import_module
from pathlib import Path
from typing import Any, TypeVar, cast

from aiolimiter import AsyncLimiter
from openai.types.chat import ChatCompletionMessageParam
from pydantic import BaseModel
from tenacity import before_sleep_log, retry, stop_after_attempt, wait_chain, wait_fixed
from typing_extensions import override

from ragu.common.logger import logger
from ragu.models.caching import ResponseCachingMixin
from ragu.utils.ragu_utils import (
    FLOATS,
    LoguruAdapter,
    attach_async_contexts,
    get_disk_cache,
    save_args_on_exception,
)


T = TypeVar('T', BaseModel, str)


def _import_gigachat() -> type[Any]:
    try:
        module = import_module('gigachat')
    except ImportError as exc:
        raise ImportError(
            'The `gigachat` package is required for GigaChat support. '
            'Install it with `pip install .[gigachat]` or `pip install gigachat`.'
        ) from exc
    return cast(type[Any], getattr(module, 'GigaChat'))


def _maybe_import_chat_model() -> type[Any] | None:
    try:
        module = import_module('gigachat.models')
    except ImportError:
        return None
    return cast(type[Any] | None, getattr(module, 'Chat', None))


def _message_content_to_text(content: Any) -> str:
    if isinstance(content, str):
        return content

    if isinstance(content, list):
        parts: list[str] = []
        for item in content:
            if isinstance(item, Mapping):
                item_type = item.get('type')
                if item_type == 'text' and isinstance(item.get('text'), str):
                    parts.append(cast(str, item['text']))
        if parts:
            return ''.join(parts)

    raise TypeError(f'Unsupported message content for GigaChat: {content!r}')


def _message_to_dict(message: ChatCompletionMessageParam) -> dict[str, Any]:
    role = cast(str, message['role'])
    return {
        'role': role,
        'content': _message_content_to_text(message.get('content')),
    }


def _normalize_schema_for_gigachat(schema: Mapping[str, Any]) -> dict[str, Any]:
    """Inline local refs because GigaChat rejects array items expressed via $ref."""
    defs = cast(dict[str, Any], deepcopy(schema.get('$defs', {})))

    def _resolve(node: Any) -> Any:
        if isinstance(node, Mapping):
            ref = node.get('$ref')
            if isinstance(ref, str) and ref.startswith('#/$defs/'):
                ref_name = ref.removeprefix('#/$defs/')
                resolved = defs.get(ref_name)
                if resolved is None:
                    raise ValueError(f'Unsupported JSON schema reference: {ref}')
                merged = {
                    key: value
                    for key, value in node.items()
                    if key != '$ref'
                }
                resolved_node = _resolve(deepcopy(resolved))
                if merged:
                    assert isinstance(resolved_node, dict)
                    resolved_node.update({key: _resolve(value) for key, value in merged.items()})
                return resolved_node

            return {
                key: _resolve(value)
                for key, value in node.items()
                if key != '$defs'
            }

        if isinstance(node, list):
            return [_resolve(item) for item in node]

        return node

    normalized = _resolve(deepcopy(schema))
    assert isinstance(normalized, dict)
    return normalized


def _build_chat_request(
    conversation: list[ChatCompletionMessageParam],
    model_name: str,
    output_schema: type[T],
    temperature: Any | None,
    top_p: Any | None,
    max_tokens: int | None,
) -> Any:
    payload: dict[str, Any] = {
        'model': model_name,
        'messages': [_message_to_dict(message) for message in conversation],
    }
    if temperature is not None:
        payload['temperature'] = temperature
    if top_p is not None:
        payload['top_p'] = top_p
    if max_tokens is not None:
        payload['max_tokens'] = max_tokens

    if not issubclass(output_schema, str):
        function_name = output_schema.__name__
        payload['functions'] = [{
            'name': function_name,
            'description': f'Output data in the structure of {function_name}',
            'parameters': _normalize_schema_for_gigachat(
                output_schema.model_json_schema(),  # type: ignore[arg-type]
            ),
        }]
        payload['function_call'] = {'name': function_name}

    chat_model = _maybe_import_chat_model()
    return chat_model(**payload) if chat_model is not None else payload


def _parse_chat_response(response: Any, output_schema: type[T]) -> T:
    choices = getattr(response, 'choices', None)
    if not choices:
        raise ValueError('GigaChat returned no choices.')

    message = choices[0].message
    if issubclass(output_schema, str):
        content = getattr(message, 'content', '')
        return cast(T, content if content is not None else '')

    if (parsed := getattr(message, 'parsed', None)) is not None:
        return cast(T, parsed)

    function_call = getattr(message, 'function_call', None)
    if function_call is not None and getattr(function_call, 'arguments', None):
        arguments = function_call.arguments
        if isinstance(arguments, str):
            return cast(T, output_schema.model_validate_json(arguments))
        return cast(T, output_schema.model_validate(arguments))

    tool_calls = getattr(message, 'tool_calls', None)
    if tool_calls:
        arguments = getattr(tool_calls[0].function, 'arguments', None)
        if arguments:
            if isinstance(arguments, str):
                return cast(T, output_schema.model_validate_json(arguments))
            return cast(T, output_schema.model_validate(arguments))

    content = getattr(message, 'content', None)
    if isinstance(content, str) and content.strip():
        return cast(T, output_schema.model_validate_json(content))

    raise ValueError('GigaChat did not return structured data in a supported format.')


async def _call_maybe_async(method: Any, *args: Any, **kwargs: Any) -> Any:
    if inspect.iscoroutinefunction(method):
        return await method(*args, **kwargs)
    return await asyncio.to_thread(method, *args, **kwargs)


def _method_accepts_kwarg(method: Any, name: str) -> bool:
    try:
        signature = inspect.signature(method)
    except (TypeError, ValueError):
        return True

    for parameter in signature.parameters.values():
        if parameter.kind is inspect.Parameter.VAR_KEYWORD:
            return True
    return name in signature.parameters


class CachedAsyncGigaChat(ResponseCachingMixin):
    """
    Native GigaChat client wrapper with caching, retrying and rate limiting.

    The public interface mirrors :class:`CachedAsyncOpenAI` to keep the logic
    layer unchanged. Conversations are accepted in OpenAI-style message format
    and translated to GigaChat requests internally.
    """

    def __init__(
        self,
        credentials: str | None = None,
        access_token: str | None = None,
        client: Any | None = None,
        rate_min_delay: float | None = None,
        rate_max_per_minute: int | None = None,
        rate_max_simultaneous: int | None = None,
        retry_times_sec: Sequence[float] | None = None,
        max_completion_tokens: int | None = None,
        cache: MutableMapping[str, Any] | str | Path | None = None,
        cache_prefix: str = 'gigachat',
        debug_errors_storage: MutableMapping[str, Any] | str | Path | None = None,
        verify_ssl_certs: bool=False,
        **client_kwargs: Any,
    ):
        self.max_completion_tokens = max_completion_tokens
        if client is None:
            GigaChat = _import_gigachat()
            init_kwargs = dict(client_kwargs)
            if credentials is not None:
                init_kwargs.setdefault('credentials', credentials)
            if access_token is not None:
                init_kwargs.setdefault('access_token', access_token)
            self.client = GigaChat(verify_ssl_certs=verify_ssl_certs, **init_kwargs)
        else:
            self.client = client

        ResponseCachingMixin.__init__(self, cache=cache, cache_prefix=cache_prefix)

        self._uncached_raw_chat_completion = self._uncached_chat_completion
        self._uncached_raw_embed_text = self._uncached_embed_text

        self.debug_errors_storage: MutableMapping[str, Any] | None
        match debug_errors_storage:
            case None:
                self.debug_errors_storage = None
            case str() | Path():
                self.debug_errors_storage = get_disk_cache(debug_errors_storage)
            case _:
                self.debug_errors_storage = debug_errors_storage
        if self.debug_errors_storage is not None:
            self._uncached_chat_completion = save_args_on_exception(
                self._uncached_chat_completion, self.debug_errors_storage)
            self._uncached_embed_text = save_args_on_exception(
                self._uncached_embed_text, self.debug_errors_storage)

        contexts: list[AbstractAsyncContextManager[Any]] = []
        if rate_max_per_minute:
            contexts.append(AsyncLimiter(rate_max_per_minute, time_period=60))
        if rate_max_simultaneous:
            contexts.append(asyncio.Semaphore(rate_max_simultaneous))
        if rate_min_delay:
            contexts.append(AsyncLimiter(1, time_period=rate_min_delay))
        if contexts:
            self._uncached_chat_completion = attach_async_contexts(
                self._uncached_chat_completion, *contexts)
            self._uncached_embed_text = attach_async_contexts(
                self._uncached_embed_text, *contexts)

        if retry_times_sec:
            retrying_decorator = retry(
                stop=stop_after_attempt(len(retry_times_sec) + 1),
                wait=wait_chain(*[wait_fixed(t) for t in retry_times_sec]),
                before_sleep=before_sleep_log(LoguruAdapter('logger'), logging.DEBUG),
                reraise=True,
            )
            self._uncached_chat_completion = retrying_decorator(self._uncached_chat_completion)
            self._uncached_embed_text = retrying_decorator(self._uncached_embed_text)

    async def chat_completion(
        self,
        model_name: str,
        conversation: list[ChatCompletionMessageParam],
        output_schema: type[T] = str,
        **kwargs: Any,
    ) -> T:
        return await self._cached_chat_completion(
            model_name=model_name,
            conversation=conversation,
            output_schema=output_schema,
            **kwargs,
        )

    async def embed_text(
        self,
        model_name: str,
        text: str,
        **kwargs: Any,
    ) -> list[float] | FLOATS:
        return await self._cached_embed_text(
            model_name=model_name,
            text=text,
            **kwargs,
        )

    async def score(
        self,
        model_name: str,
        text_1: str,
        text_2: list[str],
        **kwargs: Any,
    ) -> list[tuple[int, float]]:
        raise NotImplementedError('GigaChat scoring is not implemented in phase 1.')

    @override
    async def _uncached_chat_completion(
        self,
        model_name: str,
        conversation: list[ChatCompletionMessageParam],
        output_schema: type[T] = str,
        as_tool: bool = False,
        **kwargs: Any,
    ) -> T:
        logger.debug(f'Sending GigaChat chat_completion request with schema {output_schema.__name__}...')

        temperature = kwargs.pop('temperature', None)
        top_p = kwargs.pop('top_p', None)
        kwargs.pop('as_tool', as_tool)
        assert not kwargs, f'Guard triggered: add this to supported kwargs: {kwargs}'

        request = _build_chat_request(
            conversation=conversation,
            model_name=model_name,
            output_schema=output_schema,
            temperature=temperature,
            top_p=top_p,
            max_tokens=self.max_completion_tokens,
        )

        if hasattr(self.client, 'achat'):
            response = await _call_maybe_async(self.client.achat, request)
        elif hasattr(self.client, 'chat'):
            response = await _call_maybe_async(self.client.chat, request)
        else:
            raise TypeError('Configured GigaChat client has neither `achat` nor `chat`.')

        return _parse_chat_response(response, output_schema)

    @override
    async def _uncached_embed_text(
        self,
        model_name: str,
        text: str,
        **kwargs: Any,
    ) -> list[float] | FLOATS:
        debug_text = text[:20].replace('\n', '\\n')
        logger.debug(f'Sending GigaChat embed_text request with text {debug_text}...')
        assert not kwargs, f'Guard triggered: add this to supported kwargs: {kwargs}'

        if hasattr(self.client, 'aembeddings'):
            method = self.client.aembeddings
        elif hasattr(self.client, 'embeddings'):
            method = self.client.embeddings
        else:
            raise TypeError('Configured GigaChat client has neither `aembeddings` nor `embeddings`.')

        call_kwargs = {'texts': [text]} if _method_accepts_kwarg(method, 'texts') else {}
        if _method_accepts_kwarg(method, 'model'):
            call_kwargs['model'] = model_name
        if not call_kwargs:
            first_parameter = next(iter(inspect.signature(method).parameters.values()), None)
            if first_parameter is not None:
                call_kwargs = {first_parameter.name: [text]}
        if not call_kwargs:
            response = await _call_maybe_async(method, [text])
        else:
            response = await _call_maybe_async(method, **call_kwargs)

        if hasattr(response, 'data'):
            return cast(list[float] | FLOATS, response.data[0].embedding)
        if isinstance(response, list) and response and hasattr(response[0], 'embedding'):
            return cast(list[float] | FLOATS, response[0].embedding)

        raise ValueError('GigaChat embeddings response is not in a supported format.')
