<?php

namespace OtomaticAi\Controllers;

use Exception;
use OtomaticAi\Api\OpenAi\Client as OpenAiClient;
use OtomaticAi\Api\OtomaticAi\Client;
use OtomaticAi\Api\StabilityAi\Client as StabilityAiClient;
use OtomaticAi\Api\Ideogram\Client as IdeogramClient;
use OtomaticAi\Api\XAi\Client as XAiClient;
use OtomaticAi\Api\Flux\Client as FluxClient;
use OtomaticAi\Api\Gemini\Client as GeminiClient;
use OtomaticAi\Models\Presets\Preset;
use OtomaticAi\Utils\Auth;
use OtomaticAi\Utils\Image;
use OtomaticAi\Utils\Language;
use OtomaticAi\Vendors\Illuminate\Support\Arr;
use OtomaticAi\Vendors\Illuminate\Support\Str;
use OtomaticAi\Vendors\Illuminate\Validation\Rule;

class ImageGeneratorController extends Controller
{
    public function generate()
    {
        $this->verifyNonce();

        $this->validate([
            "image_model" => ["required"],
            "prompt" => ["required", "string"],
        ]);

        // refresh auth
        // Auth::refreshDomain();

        if (!Auth::isPremium()) {
            $this->response(["message" => "An error occurred", "error" => "This feature requires a premium subscription."], 503);
        }

        // save start event
        $this->collectEvent(
            "generate-image-start",
            [
                "image_model" => $this->input("image_model")
            ]
        );

        if (Arr::get(Image::imageModelToProviderAndModel($this->input("image_model")), "provider") === "dall_e") {
            $prompt = $this->input("prompt");
        } else {
            $prompt = $this->translatePrompt($this->input("prompt"));
        }

        try {
            $image = null;
            $providerAndModel = Image::imageModelToProviderAndModel($this->input("image_model"));

            switch ($providerAndModel["provider"]) {
                case "dall_e":
                    $image = $this->generateDallEImage($prompt, $providerAndModel["model"]);
                    break;
                case "stable_diffusion":
                    $image = $this->generateStableDiffusionImage($prompt, $providerAndModel["model"]);
                    break;
                case "ideogram":
                    $image = $this->generateIdeogramImage($prompt, $providerAndModel["model"]);
                    break;
                case "flux":
                    $image = $this->generateFluxImage($prompt, $providerAndModel["model"]);
                    break;
                case "grok":
                    $image = $this->generateGrokImage($prompt, $providerAndModel["model"]);
                    break;
                case "gemini":
                    $image = $this->generateGeminiImage($prompt, $providerAndModel["model"]);
                    break;
            }

            if (!empty($image)) {

                // refresh auth
                // Auth::refreshDomain();

                // save end event
                $this->collectEvent("generate-image-end",  [
                    "image_model" => $this->input("image_model")
                ]);

                $this->response($image);
            }
        } catch (Exception $e) {

            $this->collectEvent("generate-image-failed",  [
                "image_model" => $this->input("image_model"),
                "message" => $e->getMessage(),
            ]);

            $this->response(["message" => "An error occurred", "error" => $e->getMessage()], 503);
        }

        $this->response(["message" => "An error occurred", "error" => "No image generated."], 503);
    }

    public function save()
    {
        $this->verifyNonce();

        $this->validate([
            "name" => ["nullable"],
            "description" => ["nullable"],
            "legend" => ["nullable"],
            "post_id" => ["nullable"],
            "path" => ["required", "string"],
        ]);

        $image = Image::fromPath($this->input("path"));
        if (($attachmentId = $image->save($this->input("name"), $this->input("description"), $this->input("legend"))) !== null) {
            if (!empty($this->input("post_id"))) {
                if (set_post_thumbnail($this->input("post_id"), $attachmentId) === false) {
                }
            }
        }

        $this->response([]);
    }

    /**
     * Generate an image description from title
     *
     */
    public function generateDescription()
    {
        $this->verifyNonce();

        $this->validate([
            "title" => ["required", "string"],
        ]);

        try {
            $preset = Preset::findFromAPI("generate_alt_image");

            $response = $preset->process([
                "language" => Language::findFromLocale()->value,
                "request" => $this->input("title"),
            ]);

            // get the response content
            $response = json_decode(Arr::get($response, 'choices.0.message.content'), true, 512, JSON_THROW_ON_ERROR);
            $content = Arr::get($response, "value");
            $content = Str::clean($content);

            $this->response([
                "description" => Str::lower($content),
            ]);
        } catch (Exception $e) {
        }

        return $this->input("title");
    }

    /**
     * Translate the prompt in english
     *
     * @param string $prompt
     * @return string
     */
    private function translatePrompt(string $prompt): string
    {
        try {
            // get the openai preset
            $preset = Preset::findFromAPI("translate_content");

            // make payload
            $payload = [
                "language" => Language::find("en")->value,
                "content" => $prompt,
            ];


            // run the preset
            $response = $preset->process($payload);

            // get the response content
            $content = Arr::get($response, 'choices.0.message.content');
            $content = Str::clean($content);

            // return translated content
            if (!empty($content)) {
                return $content;
            }
        } catch (Exception $e) {
        }

        // return prompt if translation was incorrect
        return $prompt;
    }

    /**
     * save the base64 image to temp folder
     *
     * @param string $base64
     * @return string|null
     */
    private function imageUrlToBase64(string $url): ?string
    {
        $mediaFile = wp_remote_get($url);

        if (is_wp_error($mediaFile)) {
            throw new Exception("Can not download image url `" . $url . "`. Error: " . $mediaFile->get_error_message());
        }

        return base64_encode($mediaFile["body"]);
    }

    /**
     * save the base64 image to temp folder
     *
     * @param string $base64
     * @return string|null
     */
    private function saveImageInTemp(string $base64): ?string
    {
        $extension = "png";
        $name = uniqid();
        $path = get_temp_dir() . $name . '.' . $extension;

        $base64 = base64_decode($base64);

        if (file_put_contents($path, $base64) !== false) {
            return $path;
        }

        return null;
    }

    private function collectEvent(string $key, array $payload = [])
    {
        $api = new Client;
        $api->collectEvent($key, $payload);
    }

    /**
     * Generate a stable diffusion image
     *
     * @param string $prompt
     * @param string $model
     * @return array|null
     */
    private function generateStableDiffusionImage(string $prompt, string $model): ?array
    {
        $api = new StabilityAiClient();

        switch ($model) {
            case "sdxl":
                // sdxl model
                $payload = [
                    "text_prompts" => [
                        [
                            "text" => $prompt,
                            "weight" => 1
                        ]
                    ]
                ];

                $result = $api->textToImage($payload);

                if (!empty(Arr::get($result, "artifacts.0.base64"))) {
                    $path = $this->saveImageInTemp(Arr::get($result, "artifacts.0.base64"));

                    return [
                        "path" => $path,
                        "provider" => "stable_diffusion",
                        "image" => [
                            "base64" => Arr::get($result, "artifacts.0.base64"),
                            "height" => 768,
                            "width" => 1344,
                        ]
                    ];
                }
                break;
            default:
                // others models
                $payload = [
                    "prompt" => $prompt,
                    "model" => $model
                ];

                $result = $api->stableImageGenerate($payload);

                if (!empty(Arr::get($result, "image"))) {
                    $path = $this->saveImageInTemp(Arr::get($result, "image"));

                    if (!empty($path)) {
                        return [
                            "path" => $path,
                            "provider" => "stable_diffusion",
                            "image" => [
                                "base64" => Arr::get($result, "image"),
                                "height" => 768,
                                "width" => 1344,
                            ]
                        ];
                    }
                }
        }

        return null;
    }

    /**
     * Generate a ideogram image
     *
     * @param string $prompt
     * @param string $model
     * @return array|null
     */
    private function generateIdeogramImage(string $prompt, string $model)
    {
        $api = new IdeogramClient();

        $payload = [
            "prompt" => $prompt,
            "rendering_speed" => "TURBO"
        ];

        $m = Image::imageModelToProviderAndModel($this->input("image_model"));
        $payload["rendering_speed"] = Arr::get($m, "params.rendering_speed", "TURBO");

        $result = $api->generate($payload);

        if (!empty(Arr::get($result, "data.0.url"))) {
            $base64 = $this->imageUrlToBase64(Arr::get($result, "data.0.url"));
            $path = $this->saveImageInTemp($base64);

            if (!empty($path)) {
                return [
                    "path" => $path,
                    "provider" => "ideogram",
                    "image" => [
                        "base64" => $base64,
                        "width" => explode("x", Arr::get($result, "data.0.resolution", "1280x768"))[0],
                        "height" => explode("x", Arr::get($result, "data.0.resolution", "1280x768"))[1],
                    ]
                ];
            }
        }

        return null;
    }

    /**
     * Generate a grok image
     *
     * @param string $prompt
     * @param string $model
     * @return array|null
     */
    private function generateGrokImage(string $prompt, string $model)
    {
        $api = new XAiClient();

        $payload = [
            "prompt" => $prompt,
            "model" => $model
        ];

        $result = $api->image($payload);

        if (!empty(Arr::get($result, "data.0.b64_json"))) {
            $base64 = Arr::get($result, "data.0.b64_json");
            $path = $this->saveImageInTemp($base64);

            if (!empty($path)) {
                return [
                    "path" => $path,
                    "provider" => "grok",
                    "image" => [
                        "base64" => $base64,
                        "width" => 1024,
                        "height" => 768,
                    ]
                ];
            }
        }

        return null;
    }

    /**
     * Generate a google image
     *
     * @param string $prompt
     * @param string $model
     * @return array|null
     */
    private function generateGeminiImage(string $prompt, string $model)
    {
        $api = new GeminiClient();

        $payload = [
            "model" => $model,
            "contents" => [
                ["parts" => [["text" => $prompt]]]
            ],
            "generationConfig" => [
                "imageConfig" => [
                    "aspectRatio" => "16:9"
                ]
            ]
        ];

        $result = $api->image($payload);

        $base64 = GeminiClient::getImageInlineData($result);
        if (!empty($base64)) {
            $path = $this->saveImageInTemp($base64);

            if (!empty($path)) {
                return [
                    "path" => $path,
                    "provider" => "google",
                    "image" => [
                        "base64" => $base64,
                        "width" => 1024,
                        "height" => 768,
                    ]
                ];
            }
        }

        return null;
    }

    /**
     * Generate a Flux image
     *
     * @param string $prompt
     * @param string $model
     * @return array|null
     */
    private function generateFluxImage(string $prompt, string $model)
    {
        $api = new FluxClient();

        $payload = [
            "prompt" => $prompt,
            "model" => $model
        ];

        $result = $api->createTask($payload);

        $pollingUrl = Arr::get($result, "polling_url");
        if (!empty($pollingUrl)) {

            while (true) {

                // wait 3 seconds for the task to be ready
                sleep(3);

                $result = $api->getResult($pollingUrl);

                switch (Arr::get($result, "status")) {
                    case "Pending":
                        break;
                    case "Ready":
                        $base64 = $this->imageUrlToBase64(Arr::get($result, "result.sample"));
                        $path = $this->saveImageInTemp($base64);

                        if (!empty($path)) {
                            return [
                                "path" => $path,
                                "provider" => "flux",
                                "image" => [
                                    "base64" => $base64,
                                    "width" => 1344,
                                    "height" => 1024,
                                ]
                            ];
                        }
                        break;
                    default:
                        throw new Exception("Flux image generation failed. Status: " . Arr::get($result, "status"));
                }
            }
        }

        return null;
    }

    /**
     * Generate a dall-e image
     *
     * @param string $prompt
     * @param string $model
     * @return array|null
     */
    private function generateDallEImage(string $prompt, string $model)
    {
        $api = new OpenAiClient();

        $m = Image::imageModelToProviderAndModel($this->input("image_model"));
        $payload = [
            "model" => $model,
            "quality" => Arr::get($m, "params.quality", "auto"),
            "size" => "auto",
            "moderation" => "low",
        ];

        $result = $api->image($prompt, $payload);

        if (!empty(Arr::get($result, "data.0.b64_json"))) {
            $path = $this->saveImageInTemp(Arr::get($result, "data.0.b64_json"));

            if (!empty($path)) {
                return [
                    "path" => $path,
                    "provider" => "dall_e",
                    "image" => [
                        "base64" => Arr::get($result, "data.0.b64_json"),
                        "width" => $model === "dall-e-3" ? 1792 : 1024,
                        "height" => 1024,
                    ]
                ];
            }
        }

        return null;
    }
}
