from fastapi import FastAPI, HTTPException
from fastapi.responses import JSONResponse
from pydantic import BaseModel
import os
import random
import regex
import requests

app = FastAPI()

BING_URL = os.getenv("BING_URL", "https://www.bing.com")
FORWARDED_IP = (
    f"13.{random.randint(104, 107)}.{random.randint(0, 255)}.{random.randint(0, 255)}"
)
HEADERS = {
    "accept": "text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,image/png,*/*;q=0.8,application/signed-exchange;v=b3;q=0.7",
    "accept-language": "en-US,en;q=0.9",
    "cache-control": "max-age=0",
    "content-type": "application/x-www-form-urlencoded",
    "referrer": "https://www.bing.com/images/create/",
    "origin": "https://www.bing.com",
    "user-agent": "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/110.0.0.0 Safari/537.36 Edg/110.0.1587.63",
    "x-forwarded-for": FORWARDED_IP,
}

error_blocked_prompt = "Your prompt has been blocked by Bing. Try to change any bad words and try again."
error_being_reviewed_prompt = "Your prompt is being reviewed by Bing. Try to change any sensitive words and try again."
error_unsupported_lang = "\nthis language is currently not supported by bing"
error_noresults = "Could not get results"
error_redirect = "Redirect failed"
error_timeout = "Your request has timed out."
error_no_images = "No images"
error_bad_images = "Bad images"

auth_cookie = "1lstVnWXKWfcOAckdDg-ZR7mRlJrbxSz1GiD0o1k6gZB7rlucEuLUL9aKMdCOpkfWP-FksnuAjhckH7QrBrbkk7th5f4jDnfkWIgYlrWQKwClbxrbHVFd_0VKBguzLrOVKRv6wWuT3E2x6s9V3xz5ApBgb6aFmMOlUy1qstrLpX1SYfOlo-JN4xlGuJbasjTeSV96yqeVaNDMhy8nmVsb5avt6DgGisw6jQmWDu6Bzwo"
auth_cookie_SRCHHPGUSR = "SRCHLANG=en&IG=F3E10F2F1A0645129FD73C4ED349FFF6&BRW=XW&BRH=M&CW=1872&CH=958&SCW=1857&SCH=2262&DPR=1.0&UTC=420&DM=0&PV=15.0.0&WTS=63846716463&HV=1712888613&PRVCW=1872&PRVCH=958&CIBV=1.1686.0&EXLTT=2"

class Item(BaseModel):
    prompt: str

@app.post("/api/v1/smartchat/image_generate")
def get_images(a: Item):
    session = requests.Session()
    session.headers = HEADERS
    session.cookies.set("_U", auth_cookie)
    session.cookies.set("SRCHHPGUSR", auth_cookie_SRCHHPGUSR)

    url_encoded_prompt = requests.utils.quote(a.prompt)
    payload = f"q={url_encoded_prompt}&qs=ds"
    url = f"{BING_URL}/images/create?q={url_encoded_prompt}&rt=4&FORM=GENCRE"

    response = session.post(
        url,
        allow_redirects=False,
        data=payload,
        timeout=200,
    )

    if "this prompt is being reviewed" in response.text.lower():
        raise HTTPException(status_code=400, detail=error_being_reviewed_prompt)
    if "this prompt has been blocked" in response.text.lower():
        raise HTTPException(status_code=400, detail=error_blocked_prompt)
    if "we're working hard to offer image creator in more languages" in response.text.lower():
        raise HTTPException(status_code=400, detail=error_unsupported_lang)
    if response.status_code != 302:
        url = f"{BING_URL}/images/create?q={url_encoded_prompt}&rt=3&FORM=GENCRE"
        response = session.post(url, allow_redirects=False, timeout=200)
        if response.status_code != 302:
            raise HTTPException(status_code=400, detail=error_redirect)

    redirect_url = response.headers["Location"].replace("&nfy=1", "")
    request_id = redirect_url.split("id=")[-1]
    session.get(f"{BING_URL}{redirect_url}")

    polling_url = f"{BING_URL}/images/create/async/results/{request_id}?q={url_encoded_prompt}"
    while True:
        response = session.get(polling_url)
        if response.status_code != 200:
            raise HTTPException(status_code=400, detail=error_noresults)
        if not response.text or response.text.find("errorMessage") != -1:
            continue
        else:
            break

    image_links = regex.findall(r'src="([^"]+)"', response.text)
    normal_image_links = [link.split("?w=")[0] for link in image_links]
    normal_image_links = list(set(normal_image_links))

    bad_images = [
        "https://r.bing.com/rp/in-2zU3AJUdkgFe7ZKv19yPBHVs.png",
        "https://r.bing.com/rp/TX9QuO3WzcCJz1uaaSwQAz39Kb0.jpg",
    ]
    for img in normal_image_links:
        if img in bad_images:
            raise HTTPException(status_code=400, detail=error_bad_images)

    if not normal_image_links:
        raise HTTPException(status_code=400, detail=error_no_images)
 
    return JSONResponse(content={"data": normal_image_links})