import logging
import base64
import asyncio

from typing import Optional

from fastapi import FastAPI, Query

from webscout import WEBS

import requests
from requests.exceptions import RequestException
from random import randint

class Image:
    """
    This class provides methods for generating images based on prompts.
    """

    def create(self, prompt):
        """
        Create a new image generation based on the given prompt.

        Args:
            prompt (str): The prompt for generating the image.

        Returns:
            resp: The generated image content
        """
        headers = {
            "user-agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/104.0.5112.79 Safari/537.36",
        }
        try:
            resp = requests.get(
                "https://api.prodia.com/generate",
                params={
                    "new": "true",
                    "prompt": prompt,
                    "model": "dreamshaper_6BakedVae.safetensors [114c8abb]",
                    "negative_prompt": "(nsfw:1.5),verybadimagenegative_v1.3, ng_deepnegative_v1_75t, (ugly face:0.5),cross-eyed,sketches, (worst quality:2), (low quality:2.1), (normal quality:2), lowres, normal quality, ((monochrome)), ((grayscale)), skin spots, acnes, skin blemishes, bad anatomy, DeepNegative, facing away, tilted head, {Multiple people}, lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worstquality, low quality, normal quality, jpegartifacts, signature, watermark, username, blurry, bad feet, cropped, poorly drawn hands, poorly drawn face, mutation, deformed, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, extra fingers, fewer digits, extra limbs, extra arms,extra legs, malformed limbs, fused fingers, too many fingers, long neck, cross-eyed,mutated hands, polar lowres, bad body, bad proportions, gross proportions, text, error, missing fingers, missing arms, missing legs, extra digit, extra arms, extra leg, extra foot, repeating hair",
                    "steps": "50",
                    "cfg": "9.5",
                    "seed": randint(1, 10000),
                    "sampler": "Euler",
                    "aspect_ratio": "square",
                },
                headers=headers,
                timeout=30,
            )
            data = resp.json()
            while True:
                resp = requests.get(f"https://api.prodia.com/job/{data['job']}", headers=headers)
                json = resp.json()
                if json["status"] == "succeeded":
                    return requests.get(
                        f"https://images.prodia.xyz/{data['job']}.png?download=1",
                        headers=headers,
                    ).content
        except RequestException as exc:
            raise RequestException("Unable to fetch the response.") from exc

imag = Image()
TIMEOUT = 10
PROXY = None

class SearchAPIApp:
    def __init__(self):
        self.app = FastAPI(
            docs_url="/",
            title="Webscout API",
            swagger_ui_parameters={"defaultModelsExpandDepth": -1},
            version="1.3",
        )
        self.setup_routes()

    def setup_routes(self):
        @self.app.get("/api/search")
        async def search_text(
            q: str = Query(..., description="The search query string"),
            max_results: int = Query(10, ge=1, le=100, description="The maximum number of results to return"),
            timelimit: Optional[str] = Query(None, description="The time limit for the search (e.g., 'd' for day, 'w' for week, 'm' for month, 'y' for year)"),
            safesearch: str = Query("moderate", description="The safe search level ('on', 'moderate', or 'off')"),
            region: str = Query("wt-wt", description="The region for the search (e.g., 'us-en', 'uk-en', 'ru-ru')"),
        ):
            """
            Perform a text search using the Webscout API.
            """
            WEBS_instance = WEBS()
            results = []
            # Assuming WEBS_instance.text returns a list directly
            results = WEBS_instance.text(q, max_results=max_results, timelimit=timelimit, safesearch=safesearch, region=region)

            return {"results": results}

        @self.app.get("/api/images")
        async def search_images(
            q: str = Query(..., description="The search query string"),
            max_results: int = Query(10, ge=1, le=100, description="The maximum number of results to return"),
            safesearch: str = Query("moderate", description="The safe search level ('on', 'moderate', or 'off')"),
            region: str = Query("wt-wt", description="The region for the search (e.g., 'us-en', 'uk-en', 'ru-ru')"),
        ):
            """
            Perform an image search using the Webscout API.
            """
            WEBS_instance = WEBS()
            results = WEBS_instance.images(q, max_results=max_results, safesearch=safesearch, region=region)
            return {"results": results}

        @self.app.get("/api/videos")
        async def search_videos(
            q: str = Query(..., description="The search query string"),
            max_results: int = Query(10, ge=1, le=100, description="The maximum number of results to return"),
            safesearch: str = Query("moderate", description="The safe search level ('on', 'moderate', or 'off')"),
            region: str = Query("wt-wt", description="The region for the search (e.g., 'us-en', 'uk-en', 'ru-ru')"),
            timelimit: Optional[str] = Query(None, description="The time limit for the search (e.g., 'd' for day, 'w' for week, 'm' for month)"),
            resolution: Optional[str] = Query(None, description="The resolution of the videos ('small', 'medium', 'large')"),
            duration: Optional[str] = Query(None, description="The duration of the videos ('short', 'medium', 'long')"),
        ):
            """
            Perform a video search using the Webscout API.
            """
            WEBS_instance = WEBS()
            results = WEBS_instance.videos(q, max_results=max_results, safesearch=safesearch, region=region, timelimit=timelimit, resolution=resolution, duration=duration)
            return {"results": results}

        @self.app.get("/api/news")
        async def search_news(
            q: str = Query(..., description="The search query string"),
            max_results: int = Query(10, ge=1, le=100, description="The maximum number of results to return"),
            safesearch: str = Query("moderate", description="The safe search level ('on', 'moderate', or 'off')"),
            region: str = Query("wt-wt", description="The region for the search (e.g., 'us-en', 'uk-en', 'ru-ru')"),
            timelimit: Optional[str] = Query(None, description="The time limit for the search (e.g., 'd' for day, 'w' for week, 'm' for month)"),
        ):
            """
            Perform a news search using the Webscout API.
            """
            WEBS_instance = WEBS()
            results = WEBS_instance.news(q, max_results=max_results, safesearch=safesearch, region=region, timelimit=timelimit)
            return {"results": results}


        @self.app.get("/api/image/gen")
        async def generate_image(
            prompt: str = Query(..., description="The prompt to generate an image from")
        ):
            """
            Generate an image using the Image class.
            """
            try:
                image_content = imag.create(prompt)
                image_base64 = base64.b64encode(image_content).decode('utf-8')
                return {"image": image_base64}
            except Exception as e:
                return {"error": str(e)}, 500


app = SearchAPIApp().app