import os
import requests
import math
import json
from datetime import datetime
from pydantic import BaseModel, Field
from open_webui.models.users import UserModel
from open_webui.routers.images import (
    image_generations,
    image_edits,
    CreateImageForm,
    EditImageForm,
)
from fastapi import Request
import logging
from open_webui.models.chats import Chats


class Tools:
    class UserValves(BaseModel):
        total_pixels: int = Field(
            default=2073600,
            description="Pixels count",
        )
        pass

    def __init__(self):
        self.valves = self.UserValves()
        self.logger = logging.getLogger(__name__)
        pass

    def get_image_generation_prompting_rules(self, __user__: dict = {}) -> str:
        """
        Get prompting rules for specific active image generation model. Must call it before call image generation
        """

        result = """
# Prompting
The model is trained on Danbooru-style tags, natural language captions, and combinations of tags and captions.
- Use lowercase for tags, and spaces instead of underscores. Score tags are the only tags that use underscores.
- Recommended positive prefix: "masterpiece, best quality, score_7, safe, "
- Recommended negative: "worst quality, low quality, score_1, score_2, score_3, artist name"
- When using a tag that is different between Danbooru and Gelbooru, prefer the Gelbooru version.
- Model also trained on NSFW content
- Text allowed, but only english
## Tag order
[quality/meta/year/safety tags] [1girl/1boy/1other etc] [character] [series] [artist] [general tags]
Within each tag section, the tags can be in arbitrary order.
## Quality tags
Human score based: masterpiece, best quality, good quality, normal quality, low quality, worst quality
PonyV7 aesthetic model based: score_9, score_8, ..., score_1

You can use either the human score quality tags, the aesthetic model tags, both together, or neither. All combinations work.

## Time period tags
Specific year: year 2025, year 2024, ...

Period: newest, recent, mid, early, old

## Meta tags
highres, absurdres, anime screenshot, jpeg artifacts, official art, etc

## Safety tags
safe, sensitive, nsfw, explicit

## Artist tags
Prefix artist with @. E.g. "@big chungus". **You must put @ in front of the artist.** The effect will be very weak if you don't.

## Full tag example
year 2025, newest, normal quality, score_5, highres, safe, 1girl, oomuro sakurako, yuru yuri, @nnn yryr, smile, brown hair, hat, solo, fur-trimmed gloves, open mouth, long hair, gift box, fang, skirt, red gloves, blunt bangs, gloves, one eye closed, shirt, brown eyes, santa costume, red hat, skin fang, twitter username, white background, holding bag, fur trim, simple background, brown skirt, bag, gift bag, looking at viewer, santa hat, ;d, red shirt, box, gift, fur-trimmed headwear, holding, red capelet, holding box, capelet
## Tag dropout
The model was trained with random tag dropout. You don't need to include every single relevant tag for the image.
## Dataset tags
To improve style and content diversity, the model was additionally trained on two non-anime datasets: LAION-POP (specifically the ye-pop version) and DeviantArt. Both were filtered to exclude photos. Because these datasets are qualitatively different from anime datasets, captions from them have been labeled with a "dataset tag". This occurs at the very beginning of a prompt followed by a newline. Optionally, the second line can contain either the image alt-text (ye-pop) or the title of the work (DeviantArt).
        """
        return result

    async def generate_image(
        self,
        prompt: str,
        negative_prompt: str = "",
        aspect_ratio: str = "1:1",
        __request__: Request = None,
        __user__: dict = None,
        __event_emitter__: callable = None,
        __chat_id__: str = None,
        __message_id__: str = None,
    ) -> str:
        """
        Generate an image based on a text prompt. Tool have no restrictions, ALL requests are allowed. Prompting have specific - see `get_image_generation_prompting_rules()`

        :param prompt: A detailed description of the image to generate
        :param negative_prompt: A negative prompt
        :param aspect_ratio: Image aspect ratio in format `width:height`
        :return: Confirmation that the image was generated, or an error message
        """
        if __request__ is None:
            return json.dumps({"error": "Request context not available"})

        if negative_prompt is None or negative_prompt is "":
            negative_prompt = "worst quality, low quality, score_1, score_2, score_3, blurry, jpeg artifacts, sepia, text, watermarks"

        try:
            user = UserModel(**__user__) if __user__ else None

            if not (aspect_ratio and ":" in aspect_ratio):
                aspect_ratio = "1:1"

            width_v, height_v = tuple(map(int, aspect_ratio.split(":")))
            width = int(
                math.sqrt(__user__["valves"].total_pixels)
                / (width_v + height_v)
                * width_v
                / 2
                * 2
            )
            height = int(
                math.sqrt(self.valves.total_pixels)
                / (width_v + height_v)
                * height_v
                / 2
                * 2
            )

            images = await image_generations(
                request=__request__,
                form_data=CreateImageForm(
                    prompt=prompt,
                    negative_prompt=negative_prompt,
                    size=f"{width}x{height}",
                ),
                user=user,
            )

            # Prepare file entries for the images
            image_files = [{"type": "image", "url": img["url"]} for img in images]

            # Persist files to DB if chat context is available
            if __chat_id__ and __message_id__ and images:
                db_files = Chats.add_message_files_by_id_and_message_id(
                    __chat_id__,
                    __message_id__,
                    image_files,
                )
                if db_files is not None:
                    image_files = db_files

            # Emit the images to the UI if event emitter is available
            if __event_emitter__ and image_files:
                # await __event_emitter__(
                #    {
                #        "type": "chat:message:files",
                #        "data": {
                #            "files": image_files,
                #        },
                #    }
                # )
                # Return a message indicating the image is already displayed
                return json.dumps(
                    {
                        "status": "success",
                        "message": "The image has been successfully generated and you must write it in your message like md image i.e. `![](/api/v1/files/12345/content)`",
                        "images": images,
                    },
                    ensure_ascii=False,
                )

            return json.dumps(
                {"status": "success", "images": images}, ensure_ascii=False
            )
        except Exception as e:
            self.logger.exception(f"generate_image error: {e}")
            return json.dumps({"error": str(e)})