import json
from fastapi import FastAPI, Query
from fastapi.responses import JSONResponse
from playwright.async_api import async_playwright

app = FastAPI()

NANO_URL = "https://nanobanana.ai/generator"
CREATE_URL = "https://nanobanana.ai/api/generate-image"
STATUS_URL = "https://nanobanana.ai/api/generate-image/status?taskId="

#Manba @thewwiw & @wwiw_repository
#Manba @thewwiw & @wwiw_repository
#Manba @thewwiw & @wwiw_repository

async def fetch_with_playwright(prompt: str):
    async with async_playwright() as pw:
        browser = await pw.chromium.launch(headless=True)
        ctx = await browser.new_context()

        page = await ctx.new_page()
        await page.goto(NANO_URL, wait_until="networkidle")

        await page.wait_for_timeout(1500)

#Manba @thewwiw & @wwiw_repository
#Manba @thewwiw & @wwiw_repository
#Manba @thewwiw & @wwiw_repository

        cookies = await ctx.cookies()
        cookie_header = "; ".join([f"{c['name']}={c['value']}" for c in cookies])

        csrf_token = None
        for c in cookies:
            if c["name"] == "__Host-next-auth.csrf-token":
                csrf_token = c["value"]
                break

        if csrf_token is None:
            await browser.close()
            return {"error": "CSRF token topilmadi", "cookies": cookie_header}

        payload = {
            "prompt": prompt,
            "styleId": "realistic",
            "mode": "text",
            "imageUrl": None,
            "imageSize": "auto",
            "quality": "standard",
            "numImages": 1,
            "outputFormat": "png",
            "model": "nano-banana",
            "resolution": "1024*1024",
            "aspectRatio": "1:1"
        }

        js = f"""
            async () => {{
                const res = await fetch("{CREATE_URL}", {{
                    method: "POST",
                    headers: {{
                        "Content-Type": "application/json",
                        "Cookie": "{cookie_header}"
                    }},
                    body: JSON.stringify({json.dumps(payload)})
                }});
                return res.json();
            }}
        """

        create_res = await page.evaluate(js)

        if "taskId" not in create_res:
            await browser.close()
            return {
                "error": "Failed to create task",
                "response": create_res,
                "cookies": cookie_header
            }

        task_id = create_res["taskId"]

        while True:
            status_js = f"""
                async () => {{
                    const res = await fetch("{STATUS_URL}{task_id}", {{
                        method: "GET",
                        headers: {{
                            "Cookie": "{cookie_header}"
                        }}
                    }});
                    return res.json();
                }}
            """

            status = await page.evaluate(status_js)

            if status.get("status") == "completed":
                await browser.close()
                return {
                    "taskId": task_id,
                    "imageUrl": status.get("imageUrl"),
                    "status": "completed"
                }

            if status.get("status") == "failed":
                await browser.close()
                return {"error": "Generation failed", "status": status}

#Manba @thewwiw & @wwiw_repository
#Manba @thewwiw & @wwiw_repository
#Manba @thewwiw & @wwiw_repository

@app.get("/generate")
async def generate(prompt: str = Query(...)):
    try:
        result = await fetch_with_playwright(prompt)
        return JSONResponse(result)
    except Exception as e:
        return JSONResponse({"error": str(e)})