"""
一个不依赖第三方 SDK 的简易 OpenAI LLM 接口实现
"""

import typing

import json
from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, TypedDict, Union

from typing_extensions import Literal, NotRequired

# import warnings
from zyplib.import_tools import import_lazy
from zyplib.utils.print import print_debug, print_error

if TYPE_CHECKING:
    import requests
    import requests as RQ
else:
    requests = import_lazy('requests')

# ========================= OpenAI LLM Completion ========================= #


class ImageContent(TypedDict, total=False):
    """Represents image content in a message."""

    type: Literal['image_url']
    image_url: Dict[str, str]  # {"url": "..."}


class TextContent(TypedDict, total=False):
    """Represents text content in a message."""

    type: Literal['text']
    text: str


MessageContent = Union[str, List[Union[TextContent, ImageContent]]]
MessageRole = Literal['user', 'assistant', 'system']


class IChatOption(TypedDict, total=False):
    """Options for the LLM completion request."""

    temperature: NotRequired[float]
    top_p: NotRequired[float]
    max_tokens: NotRequired[int]
    frequency_penalty: NotRequired[float]
    presence_penalty: NotRequired[float]
    stop: NotRequired[Union[str, List[str]]]
    stream: NotRequired[bool]


class ILLMModel(TypedDict, total=False):
    """Represents configuration for a specific model API endpoint."""

    modelId: NotRequired[str]  # Model name sent to the API
    modelName: str  # User-facing model identifier; 通常情况下, 可以和 modelId 相同, 不过特殊情况下, 可能需要配置; 比如 Ark 平台使用的是 ID 化的模型，不方便查看
    url: str
    apiKey: str


class MessageItem(TypedDict, total=False):
    role: MessageRole
    content: MessageContent
    reasoning_content: NotRequired[str]  # 推理内容 (用于 o1 等推理模型)
    tool_call_id: NotRequired[str]  # 工具调用 ID
    tool_calls: NotRequired[List[Dict[str, Any]]]  # 工具调用


class CompletionResponse(MessageItem, total=False):
    """Represents a single message response from the LLM API."""

    usage: NotRequired[Dict[str, Any]]  # 使用统计信息
    finished: NotRequired[bool]  # 是否为最后一个 chunk


class StreamChunk(TypedDict, total=False):
    """Represents a single chunk from streaming response."""

    content: NotRequired[str]
    reasoning_content: NotRequired[str]  # 推理内容 (用于 o1 等推理模型)
    cumulative_content: NotRequired[str]  # 累积的完整内容
    cumulative_reasoning: NotRequired[str]  # 累积的推理内容
    usage: NotRequired[Dict[str, Any]]
    finished: NotRequired[bool]  # 是否为最后一个 chunk


def _copy_message_field(result: MessageItem, key: str, value: object) -> None:
    if key == 'role' and value in ('user', 'assistant', 'system'):
        result['role'] = typing.cast(MessageRole, value)
    elif key == 'content' and (isinstance(value, str) or isinstance(value, list)):
        result['content'] = typing.cast(MessageContent, value)
    elif key == 'tool_call_id' and isinstance(value, str):
        result['tool_call_id'] = value
    elif key == 'tool_calls' and isinstance(value, list):
        result['tool_calls'] = typing.cast(List[Dict[str, Any]], value)


def _adapt_input_message(
    input_data: Union[str, typing.Sequence[MessageItem]],
) -> List[MessageItem]:
    """
    Adapt input message to standard format and filter unsupported content based on model capabilities.

    Parameters
    ----------
    - `input_data` : `Union[str, List[MessageItem]]`
        - Input data, either a string or list of message items

    Returns
    ----------
    - `List[MessageItem]`
        - Standardized list of message items
    """
    messages: List[MessageItem] = []

    # Convert string input to message format
    if isinstance(input_data, str):
        messages = [{'role': 'user', 'content': input_data}]
    else:
        # Filter allowed fields and create clean messages
        allowed_fields = [
            'role',
            'content',
            'tool_call_id',
            'tool_calls',
            # 'reasoning_content',
        ]
        messages = []

        for item in input_data:
            result: MessageItem = {}
            for key, value in item.items():
                if key in allowed_fields:
                    _copy_message_field(result, key, value)
            messages.append(result)

    return messages


def _print_stream(chunk: 'StreamChunk'):
    print(chunk.get('content', ''), end='')


def openai_complete(
    prompt: Union[str, MessageItem, typing.Sequence[MessageItem]],
    model: Optional[ILLMModel] = None,
    options: Optional[IChatOption] = None,
    systemPrompt: Optional[str] = None,
    history: typing.Sequence[MessageItem] | None = None,
    proxy: Literal['system', 'none'] | str | None = None,
    return_type: Literal[
        'text', 'message_item', 'message_history', 'completion_response'
    ] = 'text',
    stream: bool = False,
    on_chunk_cb: Optional[Callable[[StreamChunk], None]] = _print_stream,
    chunk_cb_interval: int = 1,
) -> Union[str, MessageItem, List[MessageItem], CompletionResponse]:
    """
    Complete text using OpenAI-compatible API with optional streaming support.

    Supports reasoning models (e.g., o1-series) by automatically handling `reasoning_content`
    fields in the response.

    Parameters
    ----------
    - `prompt` : `Union[str, List[MessageItem]]`
        - The input prompt text or list of messages
    - `model` : `ILLMModel`, optional
        - Model configuration including URL and API key
    - `options` : `IChatOption`, optional
        - Chat completion options like temperature, max_tokens, etc.
    - `systemPrompt` : `str`, optional
        - System prompt to prepend to the conversation
    - `history` : `List[MessageItem]`, optional
        - Previous conversation history
    - `proxy` : `Literal['system', 'none'] | str`, optional
        - Proxy configuration, by default None
    - `return_type` : `Literal['text', 'message_item', 'message_history', 'completion_response']`, optional
        - Format of the return value, by default 'text'
    - `stream` : `bool`, optional
        - Enable streaming response, by default False
    - `on_chunk_cb` : `Callable[[StreamChunk], None]`, optional
        - Callback function for streaming chunks. Receives StreamChunk with content,
          reasoning_content, cumulative_content, and cumulative_reasoning fields, by default None

    Returns
    ----------
    - `Union[str, MessageItem, List[MessageItem], CompletionResponse]`
        - Response in the specified format. For reasoning models, includes formatted reasoning content.
    """
    options = options or {}

    # 1. Resolve the model identifier to get API details
    if not model:
        from zyplib._config import config

        api_url = config.OPENAI_BASE_URL
        api_key = config.OPENAI_KEY
        model_name_to_use = config.OPENAI_MODEL
    else:
        api_url = model.get('url')
        api_key = model.get('apiKey')
        model_name_to_use = model.get('modelId', model.get('modelName'))

    if not api_url or not model_name_to_use or not api_key:
        raise ValueError(
            'Model configuration is incomplete. Please provide a valid model with URL and API key.'
        )

    if not api_url or not model_name_to_use:
        raise RuntimeError(
            f"Resolved model configuration is incomplete for '{model}'. Missing URL or model name."
        )

    # api url trim last `/` if endsd with v1 / v2 v3 ... add /chat/completions
    if api_url.endswith('/'):
        api_url = api_url[:-1]  # Remove trailing slash
    if api_url.endswith('/chat/completions'):
        ...
    else:
        api_url = api_url + '/chat/completions'

    messages: List[MessageItem] = [] if history is None else list(history)
    if isinstance(prompt, str):
        messages.append({'role': 'user', 'content': prompt})
    elif isinstance(prompt, dict):
        # If prompt is a single MessageItem, add it directly
        messages.append(prompt)
    else:
        # If prompt is a list/sequence of MessageItem
        messages.extend(list(prompt))

    full_messages = _adapt_input_message(messages)

    if systemPrompt:
        # 独立处理系统提示; 防止 complete 返回List[MessageItem]中包含系统提示导致重复插入
        full_messages.insert(0, {'role': 'system', 'content': systemPrompt})

    # 2. Construct the request payload (OpenAI chat completions format)
    payload: Dict[str, Any] = {
        'model': model_name_to_use,
        'messages': full_messages,
        **options,
    }

    # 设置 stream 模式
    if stream:
        payload['stream'] = True
        # 启用 usage 数据返回
        payload['stream_options'] = {'include_usage': True}
    else:
        payload.pop('stream', None)

    # 3. Construct headers
    headers = {'Content-Type': 'application/json'}
    if api_key:  # Only add Authorization if an API key is provided
        headers['Authorization'] = f'Bearer {api_key}'

    print_debug(
        f'[LLM Request] POST {api_url}; {len(payload["messages"])} -> {model_name_to_use} (stream={stream})'
    )

    # 4. Make the API call using requests
    try:
        # Prepare session with proxy configuration
        session = requests.Session()

        if proxy is None or proxy == 'none':
            session.trust_env = False  # Disable system proxy
        elif proxy == 'system':
            session.trust_env = True  # Use system proxy
        elif isinstance(proxy, str):
            session.proxies = {'http': proxy, 'https': proxy}

        response = session.post(api_url, json=payload, headers=headers, stream=stream)

        # 5. Handle the response
        response_status = response.status_code

        if 200 <= response_status < 300:
            print_debug(f'[LLM Response] Status: {response_status}')

            if stream:
                return _handle_stream_response(
                    response, messages, return_type, on_chunk_cb, chunk_cb_interval
                )
            else:
                return _handle_normal_response(response, messages, return_type)
        else:
            response_text = response.text
            print_error(
                f'LLM API request failed with status {response_status}. Response: {response_text}'
            )
            raise RuntimeError(
                f'API request failed. Status: {response_status}, Response: {response_text[:500]}'
            )

    except requests.RequestException as e:
        print_error(f'Network error during API call: {e}')
        raise RuntimeError(f'Network error during API call: {e}') from e
    except Exception as e:
        print_error(f'An unexpected error occurred in openai_complete: {e}')
        raise RuntimeError(f'An unexpected error occurred: {e}') from e


def _handle_normal_response(
    response: 'RQ.Response', messages: typing.Sequence[MessageItem], return_type: str
) -> Union[str, MessageItem, List[MessageItem], CompletionResponse]:
    """Handle non-streaming response."""
    try:
        response_data = response.json()

        # Extract content and reasoning - adjust path based on actual API response structure
        message = response_data.get('choices', [{}])[0].get('message', {})
        content = message.get('content')
        reasoning = message.get('reasoning_content', '') or message.get('reasoning', '')

        if content is None:
            print_error(
                f'Failed to extract content from LLM response. Structure: {response_data}'
            )
            raise RuntimeError('LLM response format unexpected: Could not find content.')

        # Format content with reasoning
        # final_content = _format_content_with_reasoning(str(content), reasoning)

        print_debug(
            f'[LLM Response] Success, Content: {len(content)}, Reasoning: {len(reasoning)}'
        )

        # Create response message
        response_message: CompletionResponse = {
            'role': 'assistant',
            'content': content,
            'reasoning_content': reasoning,
            'finished': True,
        }
        usage = response_data.get('usage')
        if isinstance(usage, dict):
            response_message['usage'] = usage

        if return_type == 'text':
            return content
        elif return_type == 'message_item':
            return {
                'role': 'assistant',
                'content': content,
                'reasoning_content': reasoning,
            }
        elif return_type == 'message_history':
            return list(messages) + [
                {'role': 'assistant', 'content': content, 'reasoning_content': reasoning}
            ]
        elif return_type == 'completion_response':
            return response_message

        raise ValueError(f'Unsupported return_type: {return_type}')

    except json.JSONDecodeError as e:
        print_error(f'Failed to parse JSON response: {e}')
        raise RuntimeError(f'Failed to parse JSON response: {e}') from e


def _handle_stream_response(
    response: 'RQ.Response',
    messages: typing.Sequence[MessageItem],
    return_type: str,
    on_chunk_cb: Optional[Callable[[StreamChunk], None]],
    chunk_cb_interval: int = 1,
) -> Union[str, MessageItem, List[MessageItem], CompletionResponse]:
    """Handle streaming response."""
    cumulative_content = ''
    cumulative_reasoning = ''
    usage_info = None

    chunk_idx = 0  # 用于跟踪 chunk 调用次数

    try:
        for raw_chunk in response.iter_lines(decode_unicode=False):
            line = raw_chunk.decode('utf-8')  # 手动处理 UTF-8 解码
            if not line or not line.strip():
                continue

            if line.strip() == 'data: [DONE]':
                # 发送最终 chunk
                if on_chunk_cb:
                    final_chunk: StreamChunk = {
                        'content': '',
                        'cumulative_content': cumulative_content,
                        'cumulative_reasoning': cumulative_reasoning,
                        'finished': True,
                    }
                    if isinstance(usage_info, dict):
                        final_chunk['usage'] = usage_info
                    on_chunk_cb(final_chunk)
                break

            if line.startswith('data: '):
                try:
                    chunk_data = json.loads(line[6:])  # Remove 'data: ' prefix

                    # Handle usage information
                    if 'usage' in chunk_data:
                        usage_info = chunk_data['usage']

                    # Extract content from delta
                    choices = chunk_data.get('choices', [])
                    if choices:
                        delta = choices[0].get('delta', {})
                        chunk_content = delta.get('content', '')
                        chunk_reasoning = delta.get('reasoning_content', '') or delta.get(
                            'reasoning', ''
                        )

                        # 更新累积内容
                        if chunk_content:
                            cumulative_content += chunk_content
                        if chunk_reasoning:
                            cumulative_reasoning += chunk_reasoning

                        # 调用回调函数 (只有在有新内容时才调用)
                        if (
                            (chunk_content or chunk_reasoning)
                            and on_chunk_cb
                            and (chunk_idx % chunk_cb_interval == 0)
                        ):
                            stream_chunk: StreamChunk = {
                                'content': chunk_content,
                                'reasoning_content': chunk_reasoning,
                                'cumulative_content': cumulative_content,
                                'cumulative_reasoning': cumulative_reasoning,
                                'finished': False,
                            }
                            on_chunk_cb(stream_chunk)

                        chunk_idx += 1  # 增加 chunk 调用次数

                except json.JSONDecodeError as e:
                    print_error(f'Failed to parse stream chunk: {line} - {e}')
                    continue
                except Exception as e:
                    print_error(f'Error processing stream chunk: {e}')
                    continue

        print_debug(
            f'[LLM Stream] Complete, Content: {len(cumulative_content)}, Reasoning: {len(cumulative_reasoning)}'
        )

        # Create response message
        response_message: CompletionResponse = {
            'role': 'assistant',
            'content': cumulative_content,
            'reasoning_content': cumulative_reasoning,
            'finished': True,
        }
        if isinstance(usage_info, dict):
            response_message['usage'] = usage_info

        if return_type == 'text':
            return cumulative_content
        elif return_type == 'message_item':
            return {
                'role': 'assistant',
                'content': cumulative_content,
                'reasoning_content': cumulative_reasoning,
            }
        elif return_type == 'message_history':
            return list(messages) + [
                {
                    'role': 'assistant',
                    'content': cumulative_content,
                    'reasoning_content': cumulative_reasoning,
                }
            ]
        elif return_type == 'completion_response':
            return response_message

        raise ValueError(f'Unsupported return_type: {return_type}')

    except Exception as e:
        print_error(f'Error in stream response handling: {e}')
        raise RuntimeError(f'Error in stream response handling: {e}') from e
