import os
from langchain.memory import ConversationBufferMemory
from langchain.utilities import GoogleSearchAPIWrapper
from langchain.agents import AgentType, initialize_agent, Tool
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']


search = GoogleSearchAPIWrapper()
tools = [
    Tool(
    name ="Search" ,
    func=search.run,
    description="useful when you need to answer questions about current events"
    ),
]

llm = G4F(model=model)
agent_chain = initialize_agent(tools, llm, agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION, verbose=True)


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

@app.post('/linlada')
async def hello_post(request: Request):
    llm = G4F(model=model)
    data = await request.json()
    prompt = data['prompt']
    chat = llm(prompt)
    return chat

@app.post('/search')
async def searches(request: Request):
    data = await request.json()
    prompt = data['prompt']
    response = agent_chain.run(input=prompt)
    return response

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}'}