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

@app.route("/imagen", methods=["POST"])
async def generate_image(request: Request):
    data = await request.json()
    prompt = data['prompt']
    model = data.get["model"]
    sampler = data.get["sampler"]
    seed = int(data.get["seed"])
    neg = data.get["neg"]

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

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

@app.route("/test", methods=["POST"])
async def test(request: User):
    data = {
        'prompt': request.prompt,
        'model': request.model,
        'sampler': request.sampler,
        'seed': request.seed,
        'neg': request.neg
    }
    return data