import os
from googleapiclient.discovery import build
from lang import G4F
from fastapi import FastAPI, Request, Path
from pydantic import BaseModel
from fastapi.middleware.cors import CORSMiddleware
from ImageCreator import generate_image_prodia

app = FastAPI()

app.add_middleware(  # add the middleware
    CORSMiddleware,
    allow_credentials=True,  # allow credentials
    allow_origins=["*"],  # allow all origins
    allow_methods=["*"],  # allow all methods
    allow_headers=["*"],  # allow all headers
)


google_api_key = os.environ["GOOGLE_API_KEY"]
cse_id = os.environ["GOOGLE_CSE_ID"]
model = os.environ['default_model']

def search_google(query):
    service = build("customsearch", "v1", developerKey=google_api_key)
    result = service.cse().list(
        q=query,
        cx=cse_id
    ).execute()
    return result['items']


@app.get("/")
def hello():
    return "Hello! My name is Linlada."



def gpt_with_google_search(prompt):
    search_results = search_google(prompt)
    text = ""
    ref = ""
    for item in search_results:
        text += item['title'] + "\n" + item['snippet'] + "\n\n"
        ref += "- {} ({})\n".format(item['title'], item['link'])
    results = generate_text(f'Summarize: {text}')
    res = "{} \n\n {}".format(results, ref)
    return res

def gpt_with_google_search(prompt):
    search_results = search_google(prompt)
    text = ""
    ref = ""
    ref_count = 1
    for item in search_results:
        text += item['title'] + "\n" + item['snippet'] + "\n\n"
        ref += " [{}] {} ({})\n".format(ref_count, item['title'], item['link'])
        ref_count += 1
    results = generate_text(f'Summarize: {text}')
    res = "{} \n\n{}".format(results, ref)
    return res

class Linlada(BaseModel):
    prompt: str
    web_access: str = False
    sampler: str = None
    seed: int = None
    neg: str = None


@app.post('/linlada')
def linlada(request: Linlada):
    llm = G4F(model=model)
    prompt = request.prompt
    web_access = request.web_access
    if web_access == True:
       chat = gpt_with_google_search(prompt)
    else:
        chat = llm(prompt)
    return chat


class User(BaseModel):
    prompt: str
    model: str = None
    sampler: str = None
    seed: int = None
    neg: str = None

@app.post("/imagen")
def generate_image(request: User):
    prompt = request.prompt
    model = request.model
    sampler = request.sampler
    seed = request.seed
    neg = request.neg

    response = generate_image_prodia(prompt, model, sampler, seed, neg)
    return {"image": response}


details = {
    1: {'Absolute Reality V1.6': 'absolutereality_V16.safetensors [37db0fc3]',
        'Analog V1': 'analog-diffusion-1.0.ckpt [9ca13f02]',
        'Anything V3': 'anythingv3_0-pruned.ckpt [2700c435]',
        'Anything V4.5': 'anything-v4.5-pruned.ckpt [65745d25]',
        'Anything V5': 'anythingV5_PrtRE.safetensors [893e49b9]',
        'AbyssOrangeMix V3': 'AOM3A3_orangemixs.safetensors [9600da17]',
        'Deliberate V2': 'deliberate_v2.safetensors [10ec4b29]',
        'Dreamlike Diffusion V1': 'dreamlike-diffusion-1.0.safetensors [5c9fd6e0]',
        'Dreamlike Diffusion V2': 'dreamlike-diffusion-2.0.safetensors [fdcf65e7]',
        'Dreamshaper 6 baked vae': 'dreamshaper_6BakedVae.safetensors [114c8abb]',
        'Dreamshaper 7': 'dreamshaper_7.safetensors [5cf5ae06]',
        'Dreamshaper 8': 'dreamshaper_8.safetensors [9d40847d]',
        'Eimis Anime Diffusion V1.0': 'EimisAnimeDiffusion_V1.ckpt [4f828a15]',
        "Elldreth's Vivid": 'elldreths-vivid-mix.safetensors [342d9d26]',
        'Lyriel V1.6': 'lyriel_v16.safetensors [68fceea2]',
        'MechaMix V1.0': 'mechamix_v10.safetensors [ee685731]',
        'MeinaMix Meina V9': 'meinamix_meinaV9.safetensors [2ec66ab0]',
        'MeinaMix Meina V11': 'meinamix_meinaV11.safetensors [b56ce717]',
        'Openjourney V4': 'openjourney_V4.ckpt [ca2f377f]',
        'Portrait+ V1': 'portraitplus_V1.0.safetensors [1400e684]',
        'Realistic Vision V1.4': 'Realistic_Vision_V1.4-pruned-fp16.safetensors [8d21810b]',
        'Realistic Vision V4.0': 'Realistic_Vision_V4.0.safetensors [29a7afaa]',
        'Realistic Vision V5.0': 'Realistic_Vision_V5.0.safetensors [614d1063]',
        'Redshift Diffusion V1.0': 'redshift_diffusion-V10.safetensors [1400e684]',
        'ReV Animated V1.2.2': 'revAnimated_v122.safetensors [3f4fefd9]',
        'SD V1.4': 'sdv1_4.ckpt [7460a6fa]',
        'SD V1.5': 'v1-5-pruned-emaonly.ckpt [81761151]',
        "Shonin's Beautiful People V1.0": 'shoninsBeautiful_v10.safetensors [25d8c546]',
        "TheAlly's Mix II": 'theallys-mix-ii-churned.safetensors [5d9225a4]',
        'Timeless V1': 'timeless-1.0.ckpt [7c4971d4]'
    },
    2: {
        'Euler': 'Euler',
        'Euler a': 'Euler a',
        'Heun': 'Heun',
        'DPM++ 2M Karras': 'DPM++ 2M Karras',
        'DPM++ SDE Karras': 'DPM++ SDE Karras',
        'DDIM': 'DDIM'
    }
}

@app.get("/imagen-details/{detail_id}")
def image_detail(detail_id: int = Path(None, description="The ID of 1.model id and 2.sampler id", gt=0, lt=3)):
    return details[detail_id]


@app.post("/test")
def test(request: User):
        return {'data': f'Prompt is {request.prompt} Model is {request.model}'}