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
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
    sampler: str
    seed: int
    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}


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