from openai import Client
from flask import Flask, request
import os
import g4f

app = Flask(__name__)

def get_client():
    return Client(
        api_key=os.environ.get("API_KEY"),
        base_url="https://zukijourney.xyzbot.net/v1"
    )

def get_g4f_client():
    return g4f.client.Client()

@app.route('/gen_text', methods=['GET'])
def gen_text():
    prompt = request.args.get('prompt')

    if not prompt:
        return "Error: No prompt provided", 400
    
    print(prompt)
    
    try:
        text = _gen_text(prompt, use_zukijourney=True)
    except Exception as e:
        print(f"Error generating text: {e}")
        return f"Error generating text: {str(e)}", 500
    
    return { 'text': str(text) }


    
@app.route('/gen_prompt', methods=['GET'])
def gen_prompt():
    prompt = request.args.get('prompt')
    style = request.args.get('style')

    if not prompt:
        return "Error: No prompt provided", 400
    
    print(prompt)
    
    try:
        gen_prompt = _gen_text(f"""
Your job is to generate stable diffusion prompts. Here's how they're supposed to look:

```
score_9, score_8_up, score_7_up, score_6_up, score_5_up, score_4_up, waterelements,beautiful expressive cute face, underwater, extremely detailed, beautiful underwater cute dragon creature, no humans, extremely detailed scales, seaweed mane, fishtail, beautiful eyes, moonlight, realistic, high quality, detailed, fantasy dark cave background, sparkling water, dark shadows, glowing in the dark, bright colors, high contrast, dark background, vivid lighting,
```

```
score_9, score_8_up, score_7_up, score_6_up, score_5_up, score_4_up, best quality, highres, detailed background, office BREAK <lora:Depresso:1> depresso, blue body, purple eyes, looking aside, eye bags, tired, bored, dress shirt, white shirt, glasses, table, sitting, coffee mug
```

```
anime art, surreal (biomorphic landscape:1.1) on a cloud, art made by complex AI art prompt machine and studio ghibli, crazy and beautiful, hyperdetailed, colorful, waterfall, strange style, meadow elements, whimsical (biomechanical kingdom shire:1.2), digital art and alcohol ink manga painting mixed into one wonderful uncanny masterwork
```

Here's the request: {prompt} {f"\nKeep to the following style: {style}" if style else ""}

Output the prompt in a txt code block.""", 
    use_zukijourney=True
)
        
        print(gen_prompt)
        
        try:
            prompt = gen_prompt.split("```")[1].strip().lstrip('```\n').rstrip('\n```')
        except IndexError:
            return "Error: No prompt could be generated", 500

        
    except Exception as e:
        print(f"Error generating image prompt: {e}")
        return f"Error generating image prompt: {str(e)}", 500
    
    return { 'raw_generation': gen_prompt, 'text': prompt }


    


def _gen_text(prompt, use_zukijourney=False):
    if use_zukijourney:
        client = get_client()
        response = client.chat.completions.create(
                model="llama-3-70b-instruct",
                messages=[{"role": "user", "content": prompt}],
                max_tokens=800,
                temperature=0.7,
            )
    else:
        client = get_g4f_client()
        response = client.chat.completions.create(
                provider=g4f.Provider.PerplexityLabs,
                model="llama-3-70b-instruct",
                messages=[{"role": "user", "content": prompt}],
                max_tokens=800,
                temperature=0.7,
            )
    print(response)
    text = response.choices[0].message.content
    return text


@app.route('/gen_image', methods=['GET'])
def gen_image():
    print(request.args)
    prompt = request.args.get('prompt')
    model = request.args.get('model')
    model = model if model else "playground-v2.5"

    if not prompt:
        return "Error: No prompt provided", 400
    
    
    print(prompt)
    
    try:
        client = get_client()
        response = client.images.generate(
            model=model,
            prompt=prompt,
        )
        image_url = response.data[0].url
    except Exception as e:
        print(f"Error generating image: {e}")
        return f"Error generating image: {str(e)}", 500
    
    return { 'url': str(image_url) }

if __name__ == '__main__':
    app.run()