import asyncio
import base64
import io
import json
import re
from pathlib import Path

from PIL.Image import Image
from injected_utils.injected_cache_utils import async_cached, sqlite_dict
from loguru import logger
from openai import AsyncOpenAI, RateLimitError
from pinjected import injected, Injected, instances


def to_content(img: Image):
    # convert Image into jpeg bytes
    jpg_bytes = io.BytesIO()
    img.convert('RGB').save(jpg_bytes, format='jpeg', quality=95)
    b64_image = base64.b64encode(jpg_bytes.getvalue()).decode('utf-8')
    mb_of_b64 = len(b64_image) / 1024 / 1024
    logger.info(f"image size: {mb_of_b64:.2f} MB in base64.")
    return {
        "type": 'image_url',
        "image_url": f"data:image/jpeg;base64,{b64_image}"
    }


@injected
async def a_repeat_for_rate_limit(logger, /, task):
    while True:
        try:
            return await task()
        except RateLimitError as e:
            logger.error(f"rate limit error: {e}")
            pat = "Please retry after (\d+) seconds."
            match = re.search(pat, e.message)
            if match:
                seconds = int(match.group(1))
                logger.info(f"sleeping for {seconds} seconds")
                await asyncio.sleep(seconds)
            else:
                logger.warning(f"failed to parse rate limit error message: {e.message}")
                await asyncio.sleep(10)


@injected
async def a_vision_llm__gpt4(
        async_openai_client: AsyncOpenAI,
        a_repeat_for_rate_limit,
        /,
        text: str,
        images: list[Image]) -> str:
    assert isinstance(async_openai_client, AsyncOpenAI)

    for img in images:
        assert isinstance(img, Image), f"image is not Image, but {type(img)}"

    async def task():
        chat_completion = await async_openai_client.chat.completions.create(
            messages=[
                {
                    "role": "user",
                    "content": [
                        {
                            "type": 'text',
                            "text": text
                        },
                        *[to_content(img) for img in images]
                    ]
                }
            ],
            model="gpt-4-vision-preview",
            max_tokens=2048
        )
        return chat_completion

    chat_completion = await a_repeat_for_rate_limit(task)
    res = chat_completion.choices[0].message.content
    assert isinstance(res, str)
    logger.info(f"vision_llm__gpt4 result:\n{res}")
    return res


@injected
async def a_llm__openai(
        async_openai_client: AsyncOpenAI,
        a_repeat_for_rate_limit,
        /,
        text: str,
        model_name: str,
        max_completion_tokens=4096,
) -> str:
    assert isinstance(async_openai_client, AsyncOpenAI)

    async def task():
        # import tiktoken
        # enc = tiktoken.get_encoding("cl100k_base")
        # n_token = len(enc.encode(text))
        chat_completion = await async_openai_client.chat.completions.create(
            messages=[
                {
                    "role": "user",
                    "content": [
                        {
                            "type": 'text',
                            "text": text
                        },
                    ]
                }
            ],
            model=model_name,
            max_tokens=max_completion_tokens
        )
        return chat_completion

    chat_completion = await a_repeat_for_rate_limit(task)
    res = chat_completion.choices[0].message.content
    assert isinstance(res, str)
    logger.info(f"result:\n{res}")
    return res


@injected
async def a_llm__gpt4_turbo(
        a_llm__openai,
        /,
        text: str,
        max_completion_tokens=4096
) -> str:
    return await a_llm__openai(text, max_completion_tokens=max_completion_tokens, model_name="gpt-4-turbo-preview")

a_llm__gpt4_turbo_cached = async_cached(
    sqlite_dict(str(Path("~/.cache/a_llm__gpt4_turbo.sqlite").expanduser()))
)(a_llm__gpt4_turbo)

@injected
async def a_llm__gpt35_turbo(
        a_llm__openai,
        /,
        text: str,
        max_completion_tokens=4096
) -> str:
    return await a_llm__openai(text, max_completion_tokens=max_completion_tokens, model_name="gpt-3.5-turbo")


@injected
async def a_json_llm__openai(
        logger,
        async_openai_client: AsyncOpenAI,
        a_repeat_for_rate_limit,
        /,
        text: str,
        max_completion_tokens=4096,
        model="gpt-4-0125-preview"
) -> str:
    assert isinstance(async_openai_client, AsyncOpenAI)

    async def task():
        # import tiktoken
        # enc = tiktoken.get_encoding("cl100k_base")
        # n_token = len(enc.encode(text))
        chat_completion = await async_openai_client.chat.completions.create(
            messages=[
                {
                    "role": "user",
                    "content": [
                        {
                            "type": 'text',
                            "text": text
                        },
                    ]
                }
            ],
            model=model,
            max_tokens=max_completion_tokens,
            response_format={"type": "json_object"}
        )
        return chat_completion

    chat_completion = await a_repeat_for_rate_limit(task)
    res = chat_completion.choices[0].message.content
    assert isinstance(res, str)
    logger.info(f"\n{res}")
    return json.loads(res)


@injected
async def a_json_llm__gpt4_turbo(
        a_json_llm__openai,
        /,
        text: str,
        max_completion_tokens=4096
) -> str:
    return await a_json_llm__openai(
        text=text,
        max_completion_tokens=max_completion_tokens,
        model="gpt-4-turbo-preview"
    )


test_vision_llm__gpt4 = a_vision_llm__gpt4(
    text="What are inside this image?",
    images=Injected.list(
    ),
)
"""
('The image appears to be an advertisement or an informational graphic about '
 'infant and newborn nutrition. It features a baby with light-colored hair who '
 'is lying down and holding onto a baby bottle, seemingly feeding themselves. '
 'The baby is looking directly towards the camera. The image uses a soft pink '
 'color palette, which is common for baby-related products or information. '
 'There are texts that read "Infant & Newborn Nutrition" and "Absolutely New," '
 'along with the word "PINGUIN" at the top, which could be a brand name or '
 "logo. The layout and design of this image suggest it's likely used for "
 'marketing purposes or as part of educational material regarding baby '
 'nutrition.')
"""

test_llm__gpt4_turbo = a_llm__gpt4_turbo(
    "Hello world"
)

test_json_llm__gpt4_turbo = a_json_llm__gpt4_turbo(
    "Hello world, respond to me in json"
)

__meta_design__ = instances()
