import gradio as gr
import asyncio, time
from pathlib import Path
from typing import AsyncGenerator, Tuple
from playwright.async_api import async_playwright

MOVE_ORDER_DEFAULT = ["ArrowUp", "ArrowLeft", "ArrowRight", "ArrowDown"]

async def get_score(page):
    el = await page.query_selector(".score-container")
    if not el:
        return 0
    text = (await el.inner_text()).strip()
    try:
        return int(text.split()[0].replace(",", ""))
    except Exception:
        return 0

async def is_game_over(page):
    msg = await page.query_selector(".game-message.game-over, .game-over")
    return msg is not None

async def board_signature(page):
    tiles = await page.query_selector_all(".tile")
    parts = []
    for t in tiles:
        classes = (await t.get_attribute("class")) or ""
        parts.append(classes)
    return "|".join(sorted(parts))

async def autoplay_stream(url:str, moves:int, fps:int, strategy:str) -> AsyncGenerator[Tuple[str,str], None]:
    """Yields (frame_path, logs_text) repeatedly so the first output streams frames."""
    logs = []
    img_path = Path("live_frame.png")
    final_path = Path("final_frame.png")
    move_order = MOVE_ORDER_DEFAULT.copy()
    if strategy == "Left/Down Bias":
        move_order = ["ArrowLeft", "ArrowDown", "ArrowRight", "ArrowUp"]
    elif strategy == "Clockwise":
        move_order = ["ArrowUp", "ArrowRight", "ArrowDown", "ArrowLeft"]

    frame_delay_ms = max(50, int(1000 / max(1, fps)))

    async with async_playwright() as p:
        browser = await p.chromium.launch(
            headless=True,
            args=["--no-sandbox", "--disable-setuid-sandbox"]
        )
        page = await browser.new_page(viewport={"width": 900, "height": 1100})
        await page.goto(url, wait_until="load", timeout=60000)
        logs.append(f"Loaded: {url}")
        yield str(img_path), "\n".join(logs)

        # Make sure the game surface has focus
        await page.mouse.click(50, 50)
        last_sig = await board_signature(page)
        invalid_count = 0
        played = 0
        last_frame_time = 0

        for i in range(moves):
            await page.keyboard.press(move_order[i % 4])
            await page.wait_for_timeout(60)

            sig = await board_signature(page)
            if sig == last_sig:
                for alt in move_order[1:]:
                    await page.keyboard.press(alt)
                    await page.wait_for_timeout(50)
                    new_sig = await board_signature(page)
                    if new_sig != sig:
                        sig = new_sig
                        break
                else:
                    invalid_count += 1
            else:
                invalid_count = 0

            last_sig = sig
            played += 1
            sc = await get_score(page)
            logs.append(f"Move {played:03d} | Score {sc}")
            # throttle logs to last 60 lines
            if len(logs) > 60:
                logs = logs[-60:]

            # Stream a frame at desired fps
            now = time.time()*1000
            if now - last_frame_time >= frame_delay_ms:
                await page.screenshot(path=str(img_path), full_page=True)
                last_frame_time = now
                yield str(img_path), "\n".join(logs)

            if await is_game_over(page):
                logs.append("Game Over detected. Stopping.")
                break

            if invalid_count >= 3:
                first = move_order.pop(0)
                move_order.append(first)
                invalid_count = 0
                logs.append(f"Rotated move order: {move_order}")

        await page.screenshot(path=str(final_path), full_page=True)
        yield str(final_path), "\n".join(logs)  # final frame
        await browser.close()

with gr.Blocks(title="Online Game Auto-Player (Live Stream)") as demo:
    gr.Markdown("# Online Game Auto-Player (Live Stream)")
    gr.Markdown("Plays an **ONLINE** game in a headless browser and streams frames while it plays. "
                "Default URL is play2048.co. You can replace with any accessible game URL that reacts to arrow keys or clicks.")
    with gr.Row():
        url_in = gr.Textbox(label="Game URL", value="https://play2048.co")
    with gr.Row():
        moves = gr.Slider(50, 1200, value=300, step=10, label="Max moves")
        fps = gr.Slider(1, 15, value=5, step=1, label="Stream FPS")
        strat = gr.Dropdown(choices=["Up/Left Bias","Left/Down Bias","Clockwise"], value="Up/Left Bias", label="Move Strategy")
    start = gr.Button("Start")
    with gr.Row():
        stream_out = gr.Image(label="Live View (frames)", streaming=True)
        logs_out = gr.Textbox(label="Logs (tail)", lines=18)
    gr.Markdown("**Note:** Internet access must be enabled in Space settings for online URLs. "
                "On free (CPU Basic) Spaces this option is generally unavailable.")

    start.click(fn=autoplay_stream, inputs=[url_in, moves, fps, strat], outputs=[stream_out, logs_out])

if __name__ == "__main__":
    demo.launch()