from __future__ import annotations

import json
import typing

import ollama
from .. import operator, entities, errors


@operator.operator_class(
    name='ollama',
    help='ollama平台操作',
    usage='!ollama\n!ollama show <模型名>\n!ollama pull <模型名>\n!ollama del <模型名>',
)
class OllamaOperator(operator.CommandOperator):
    async def execute(
        self, context: entities.ExecuteContext
    ) -> typing.AsyncGenerator[entities.CommandReturn, None]:
        try:
            content: str = '模型列表:\n'
            model_list: list = ollama.list().get('models', [])
            for model in model_list:
                content += f'名称: {model["name"]}\n'
                content += f'修改时间: {model["modified_at"]}\n'
                content += f'大小: {bytes_to_mb(model["size"])}MB\n\n'
            yield entities.CommandReturn(text=f'{content.strip()}')
        except ollama.ResponseError:
            yield entities.CommandReturn(
                error=errors.CommandError('无法获取模型列表，请确认 Ollama 服务正常')
            )


def bytes_to_mb(num_bytes):
    mb: float = num_bytes / 1024 / 1024
    return format(mb, '.2f')


@operator.operator_class(
    name='show', help='ollama模型详情', privilege=2, parent_class=OllamaOperator
)
class OllamaShowOperator(operator.CommandOperator):
    async def execute(
        self, context: entities.ExecuteContext
    ) -> typing.AsyncGenerator[entities.CommandReturn, None]:
        content: str = '模型详情:\n'
        try:
            show: dict = ollama.show(model=context.crt_params[0])
            model_info: dict = show.get('model_info', {})
            ignore_show: str = 'too long to show...'

            for key in ['license', 'modelfile']:
                show[key] = ignore_show

            for key in [
                'tokenizer.chat_template.rag',
                'tokenizer.chat_template.tool_use',
            ]:
                model_info[key] = ignore_show

            content += json.dumps(show, indent=4)
            yield entities.CommandReturn(text=content.strip())
        except ollama.ResponseError:
            yield entities.CommandReturn(
                error=errors.CommandError('无法获取模型详情，请确认 Ollama 服务正常')
            )


@operator.operator_class(
    name='pull', help='ollama模型拉取', privilege=2, parent_class=OllamaOperator
)
class OllamaPullOperator(operator.CommandOperator):
    async def execute(
        self, context: entities.ExecuteContext
    ) -> typing.AsyncGenerator[entities.CommandReturn, None]:
        try:
            model_list: list = ollama.list().get('models', [])
            if context.crt_params[0] in [model['name'] for model in model_list]:
                yield entities.CommandReturn(text='模型已存在')
                return
        except ollama.ResponseError:
            yield entities.CommandReturn(
                error=errors.CommandError('无法获取模型列表，请确认 Ollama 服务正常')
            )
            return

        on_progress: bool = False
        progress_count: int = 0
        try:
            for resp in ollama.pull(model=context.crt_params[0], stream=True):
                total: typing.Any = resp.get('total')
                if not on_progress:
                    if total is not None:
                        on_progress = True
                    yield entities.CommandReturn(text=resp.get('status'))
                else:
                    if total is None:
                        on_progress = False

                    completed: typing.Any = resp.get('completed')
                    if isinstance(completed, int) and isinstance(total, int):
                        percentage_completed = (completed / total) * 100
                        if percentage_completed > progress_count:
                            progress_count += 10
                            yield entities.CommandReturn(
                                text=f'下载进度: {completed}/{total} ({percentage_completed:.2f}%)'
                            )
        except ollama.ResponseError as e:
            yield entities.CommandReturn(text=f'拉取失败: {e.error}')


@operator.operator_class(
    name='del', help='ollama模型删除', privilege=2, parent_class=OllamaOperator
)
class OllamaDelOperator(operator.CommandOperator):
    async def execute(
        self, context: entities.ExecuteContext
    ) -> typing.AsyncGenerator[entities.CommandReturn, None]:
        try:
            ret: str = ollama.delete(model=context.crt_params[0])['status']
        except ollama.ResponseError as e:
            ret = f'{e.error}'
        yield entities.CommandReturn(text=ret)