from fastapi import FastAPI, HTTPException
from fastapi.responses import StreamingResponse
import requests
from urllib.parse import quote
import asyncio
from io import BytesIO
import os

urls = os.environ.get("url")

app = FastAPI()

def create_job(prompt, model, sampler, seed, neg):
    if model is None:
        model = 'Realistic_Vision_V5.0.safetensors [614d1063]'
    if sampler is None:
        sampler = 'DPM++ 2M Karras'
    if seed is None:
        seed = '-1'
    if neg is None:
        neg = "(nsfw:1.5),verybadimagenegative_v1.3, ng_deepnegative_v1_75t, (ugly face:0.8),cross-eyed,sketches, (worst quality:2), (low quality:2), (normal quality:2), lowres, normal quality, ((monochrome)), ((grayscale)), skin spots, acnes, skin blemishes, bad anatomy, DeepNegative, facing away, tilted head, {Multiple people}, lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worstquality, low quality, normal quality, jpegartifacts, signature, watermark, username, blurry, bad feet, cropped, poorly drawn hands, poorly drawn face, mutation, deformed, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, extra fingers,fewer digits ,extra limbs ,extra arms ,extra legs ,malformed limbs ,fused fingers ,too many fingers ,long neck ,cross-eyed ,mutated hands ,polar lowres ,bad body ,bad proportions ,gross proportions ,text ,error ,missing fingers ,missing arms ,missing legs ,extra digit ,extra arms ,extra leg ,extra foot ,repeating hair ,nsfw ,[bad-artist-anime],[sketch by bad-artist] ,[mutation],[lowres],[bad hands],[text],[signature],[watermark],[username],[blurry],[monochrome],[grayscale],[realistic],[simple background],[limited palette],close-up,(swimsuit),(cleavage),(armpits),(ass),(navel),(cleavage cutout),(forehead jewel:1.2),(forehead mark:1.5),(bad and mutated hands:1.3),(worst quality:2.0),(low quality:2.0),(blurry:2.0),multiple limbs,bad anatomy,(interlocked fingers:1.2),(interlocked leg:1.2),Ugly Fingers,(extra digit and hands and fingers and legs and arms:1.4),crown braid,(deformed fingers:1.2),(long fingers:1.2)"
    url = f'https://api.{urls}.com/generate'
    params = {
        'new': 'true',
        'prompt': quote(prompt),
        'model': model,
        'negative_prompt': neg,
        'steps': '100',
        'cfg': '9.5',
        'seed': seed,
        'sampler': sampler,
        'upscale': 'True',
        'aspect_ratio': 'square'
    }
    response = requests.get(url, params=params)
    response.raise_for_status()
    return response.json()['job']

@app.get('/')
def hello_world():
    return """
    This is an API for AI-generated images.
    Params:
        - prompt: The prompt for image generation
        - model: The model to use for image generation (default: Realistic_Vision_V5.0.safetensors [614d1063])
        - negative_prompt: Negative prompts for image generation
        - seed: The seed for image generation (default: -1)
        - sampler: The sampler to use for image generation (default: DPM++ 2M Karras)
    Sample: /generate_image?prompt=your_prompt&model=model_name&sampler=sampler_name&seed=seed_number&neg=negative_prompts
    """

@app.get("/generate_image")
async def generate_image(prompt: str, model: str = None, sampler: str = None, seed: str = None, neg: str = None):
    job_id = create_job(prompt, model, sampler, seed, neg)
    url = f'https://api.{urls}.com/job/{job_id}'
    headers = {'accept': '*/*'}
    
    while True:
        response = requests.get(url=url, headers=headers)
        response.raise_for_status()
        job_response = response.json()
        if job_response['status'] == 'succeeded':
            image_url = f'https://images.{urls}.xyz/{job_id}.png'
            image_response = requests.get(image_url)
            image_response.raise_for_status()
            image = BytesIO(image_response.content)
            return StreamingResponse(image, media_type='image/png')
        await asyncio.sleep(2)  # Add a delay to prevent excessive requests

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000, debug=True)