import streamlit as st
import os
import pygame
from PIL import Image
import requests
import io
from streamlit_autorefresh import st_autorefresh
import base64
st.set_page_config(layout="wide")
os.environ["SDL_VIDEODRIVER"] = "dummy"
pygame.init()
pygame.display.set_mode((1, 1))

# Initialize session state for game variables
if 'page' not in st.session_state:
    st.session_state.page = 'menu'




# Helper to load and convert image to base64 data URL
def img_to_data_url(path):
    img = Image.open(path)
    buf = io.BytesIO()
    img.save(buf, format="PNG")
    b64 = base64.b64encode(buf.getvalue()).decode()
    return f"data:image/png;base64,{b64}"

# Helper to load and convert audio to base64 data URL
def audio_to_data_url(path, mime_type):
    with open(path, "rb") as f:
        b64 = base64.b64encode(f.read()).decode()
    return f"data:{mime_type};base64,{b64}"


# Player image
plane_data_url = img_to_data_url('images/player.png')

# Audio base64 URLs
shoot_audio_url = audio_to_data_url('game_sounds/shooting/shoot.mp3', 'audio/mpeg')
explosion_audio_url = audio_to_data_url('game_sounds/explosions/explosion1.wav', 'audio/wav')
gameover_audio_url = audio_to_data_url('game_sounds/gameover.mp3', 'audio/mpeg')
bgm_audio_url = audio_to_data_url('game_sounds/background_music.mp3', 'audio/mpeg')

# Enemy images
enemy_imgs = [
    img_to_data_url('images/enemy/enemy1_1.png'),
    img_to_data_url('images/enemy/enemy1_2.png'),
    img_to_data_url('images/enemy/enemy1_3.png'),
    img_to_data_url('images/enemy/enemy2_1.png'),
    img_to_data_url('images/enemy/enemy2_2.png'),
]

# Boss images
boss_imgs = [
    img_to_data_url('images/boss/boss1.png'),
    img_to_data_url('images/boss/boss2.png'),
    img_to_data_url('images/boss/boss2_1.png'),
    img_to_data_url('images/boss/boss3.png'),
]

# Bullet images
bullet_imgs = [
    img_to_data_url('images/bullets/bullet1.png'),
    img_to_data_url('images/bullets/bullet2.png'),
    img_to_data_url('images/bullets/bullet3.png'),
    img_to_data_url('images/bullets/bullet4.png'),
]
# Explosion images (first 5 frames for animation)
explosion_imgs = [
    img_to_data_url('images/explosion/explosion0.png'),
    img_to_data_url('images/explosion/explosion1.png'),
    img_to_data_url('images/explosion/explosion2.png'),
    img_to_data_url('images/explosion/explosion3.png'),
    img_to_data_url('images/explosion/explosion4.png'),
]

# Meteor images
meteor_imgs = [
    img_to_data_url('images/meteors/meteor_1.png'),
    img_to_data_url('images/meteors/meteor_2.png'),
    img_to_data_url('images/meteors/meteor_3.png'),
    img_to_data_url('images/meteors/meteor_4.png'),
    img_to_data_url('images/meteors/meteor2_1.png'),
    img_to_data_url('images/meteors/meteor2_2.png'),
    img_to_data_url('images/meteors/meteor2_3.png'),
    img_to_data_url('images/meteors/meteor2_4.png'),
]
# Refill images
refill_imgs = [
    img_to_data_url('images/refill/bullet_refill.png'),
    img_to_data_url('images/refill/double_refill.png'),
    img_to_data_url('images/refill/health_refill.png'),
]
# Hole images
hole_imgs = [
    img_to_data_url('images/hole/black_hole.png'),
    img_to_data_url('images/hole/black_hole2.png'),
]
# Score images
score_imgs = [
    img_to_data_url('images/score/score_coin.png'),
]

# Menu page
if st.session_state.page == 'menu':
    st.title('Cosmic Heat')
    menu_img = Image.open('images/mainmenu.jpg')
    st.image(menu_img, use_container_width=True)
    if st.button('Play'):
        st.session_state.page = 'game'
        st.rerun()
    if st.button('Exit'):
        st.stop()

# Game page
if st.session_state.page == 'game':
    st.title('Cosmic Heat - Game')
    st_autorefresh(interval=100, key="game_refresh")

    # Pass image arrays as JSON to JS
    import json
    enemy_imgs_js = json.dumps(enemy_imgs)
    boss_imgs_js = json.dumps(boss_imgs)
    bullet_imgs_js = json.dumps(bullet_imgs)
    explosion_imgs_js = json.dumps(explosion_imgs)
    meteor_imgs_js = json.dumps(meteor_imgs)
    refill_imgs_js = json.dumps(refill_imgs)
    hole_imgs_js = json.dumps(hole_imgs)
    score_imgs_js = json.dumps(score_imgs)

    game_data = st.components.v1.html(f"""
<div id=\"game-canvas\" style=\"position: relative; width: 100%; height: 600px; overflow: hidden; background: url('images/bg/background.jpg') no-repeat; background-size: cover; position: relative; z-index: 1;\">
    <audio id=\"shoot-audio\" src=\"{shoot_audio_url}\" preload=\"auto\"></audio>
    <audio id=\"explosion-audio\" src=\"{explosion_audio_url}\" preload=\"auto\"></audio>
    <audio id=\"gameover-audio\" src=\"{gameover_audio_url}\" preload=\"auto\"></audio>
    <audio id=\"bgm-audio\" src=\"{bgm_audio_url}\" preload=\"auto\" loop></audio>
    <div id=\"player\" style=\"position: absolute; width: 50px; height: 50px; background: url('{plane_data_url}') no-repeat; background-size: 100% 100%; border: 1px solid red; z-index: 2;\"></div>
    <div id=\"enemies\"></div>
    <div id=\"bosses\"></div>
    <div id=\"meteors\"></div>
    <div id=\"refills\"></div>
    <div id=\"holes\"></div>
    <div id=\"scores\"></div>
    <div id=\"explosions\"></div>
    <canvas id=\"bullet-canvas\" style=\"position: absolute; top: 0; left: 0; width: 100%; height: 100%; pointer-events: none; z-index: 3;\"></canvas>
    <input type=\"hidden\" id=\"game-frame\" name=\"game-frame-output\">

    <!-- Game Over Message -->
    <div id="game-over-message" style="display:none; position:absolute; top:50%; left:50%; transform:translate(-50%,-50%); color:red; font-size:24px; z-index:10;">
        Game Over
    </div>
</div>
<script>
    var bgImgs = [
        'images/bg/background.jpg',
        'images/bg/background2.png',
        'images/bg/background3.png',
        'images/bg/background4.png',
    ];
    const enemyImgs = {enemy_imgs_js};
    const bossImgs = {boss_imgs_js};
    const bulletImgs = {bullet_imgs_js};
    const explosionImgs = {explosion_imgs_js};
    const meteorImgs = {meteor_imgs_js};
    const refillImgs = {refill_imgs_js};
    const holeImgs = {hole_imgs_js};
    const scoreImgs = {score_imgs_js};

    let level = 0;
    let score = 0;
    let explosions = [];
    const explosionsDiv = document.getElementById('explosions');

    const canvas = document.getElementById('game-canvas');
    const player = document.getElementById('player');
    const bulletCanvas = document.getElementById('bullet-canvas');
    const ctx = bulletCanvas.getContext('2d');
    const gameFrame = document.getElementById('game-frame');
    const enemiesDiv = document.getElementById('enemies');
    const bossesDiv = document.getElementById('bosses');
    const meteorsDiv = document.getElementById('meteors');
    const refillsDiv = document.getElementById('refills');
    const holesDiv = document.getElementById('holes');
    const scoresDiv = document.getElementById('scores');
    let playerX = canvas.offsetWidth / 2 - 25;
    let playerY = canvas.offsetHeight / 2 - 25;
    let bullets = [];
    let enemyObjs = [];
    let bossObjs = [];
    let meteorObjs = [];
    let refillObjs = [];
    let holeObjs = [];
    let scoreObjs = [];
    let bulletCounter = 2000;
    let lastShotTime = 0;
    const bulletSpeed = 5;
    const bulletSize = 10;
    let firing = false;
    let gameOver = false;

    function setBackground() {{
    canvas.style.background = `url('${{bgImgs[level % bgImgs.length]}}') no-repeat`;
    canvas.style.backgroundSize = 'cover';
    }}

    function spawnEnemies() {{
        enemyObjs = [];
        for (let i = 0; i < 3 + level; i++) {{
            enemyObjs.push({{
                x: Math.random() * (canvas.offsetWidth - 50),
                y: -60, // Start above the visible area
                img: enemyImgs[i % enemyImgs.length],
                w: 50,
                h: 50,
                dx: (Math.random() - 0.5) * 2,
                dy: 1 + Math.random(),
                entering: true // Mark as entering
            }});
        }}
    }}

    function spawnBoss() {{
        bossObjs = [];
        bossObjs.push({{
            x: canvas.offsetWidth / 2 - 75,
            y: 20,
            img: bossImgs[level % bossImgs.length],
            w: 150,
            h: 100,
            dx: 1,
            dy: 0.5,
        }});
    }}

    function spawnMeteors() {{
        meteorObjs = [];
        for (let i = 0; i < 2 + level; i++) {{
            meteorObjs.push({{
                x: Math.random() * (canvas.offsetWidth - 40),
                y: -60,
                img: meteorImgs[i % meteorImgs.length],
                w: 40,
                h: 40,
                dy: 2 + Math.random() * 2
            }});
        }}
    }}
    function spawnRefills() {{
        refillObjs = [];
        for (let i = 0; i < 1; i++) {{
            refillObjs.push({{
                x: Math.random() * (canvas.offsetWidth - 30),
                y: -60,
                img: refillImgs[i % refillImgs.length],
                w: 30,
                h: 30,
                dy: 2
            }});
        }}
    }}
    function spawnHoles() {{
        holeObjs = [];
        for (let i = 0; i < 1; i++) {{
            holeObjs.push({{
                x: Math.random() * (canvas.offsetWidth - 40),
                y: -60,
                img: holeImgs[i % holeImgs.length],
                w: 40,
                h: 40,
                dy: 2.5
            }});
        }}
    }}
    function spawnScores() {{
        scoreObjs = [];
        for (let i = 0; i < 1; i++) {{
            scoreObjs.push({{
                x: Math.random() * (canvas.offsetWidth - 30),
                y: -60,
                img: scoreImgs[i % scoreImgs.length],
                w: 30,
                h: 30,
                dy: 2
            }});
        }}
    }}

    setBackground();
    spawnEnemies();
    spawnBoss();
    spawnMeteors();
    spawnRefills();
    spawnHoles();
    spawnScores();

    // Handle mouse movement
    canvas.addEventListener('mousemove', function(e) {{
        const rect = canvas.getBoundingClientRect();
        playerX = Math.max(0, Math.min(rect.width - 50, e.clientX - rect.left - 25));
        playerY = Math.max(0, Math.min(rect.height - 50, e.clientY - rect.top - 25));
        updatePlayer();
    }});

    canvas.addEventListener('mousedown', function(e) {{
        firing = true;
    }});
    canvas.addEventListener('mouseup', function(e) {{
        firing = false;
    }});
    canvas.addEventListener('mouseleave', function(e) {{
        firing = false;
    }});

    // Collision detection and playerRect are JS only, not Python
    function rectsOverlap(a, b) {{
        return a.x < b.x + b.w && a.x + a.w > b.x && a.y < b.y + b.h && a.y + a.h > b.y;
    }}
    function playerRect() {{
        return {{ x: playerX, y: playerY, w: 50, h: 50 }};
    }}

    // Collision detection
    function rectsOverlap(a, b) {{
        return a.x < b.x + b.w && a.x + a.w > b.x && a.y < b.y + b.h && a.y + a.h > b.y;
    }}

    function playerRect() {{
    return {{ x: playerX, y: playerY, w: 50, h: 50 }};
    }}

    // Update game state
    function updateGame() {{
        if (gameOver) return;
        // Update bullets
        bullets = bullets.filter(function(bullet) {{
            bullet.y += bullet.dy;
            return bullet.y > 0;
        }});

        // Move enemies
        enemyObjs.forEach(function(enemy) {{
            enemy.x += enemy.dx;
            enemy.y += enemy.dy;
            if (enemy.entering && enemy.y >= Math.random() * 200) {{
                enemy.entering = false; // Stop marking as entering after it reaches random y
            }}
            if (enemy.x < 0 || enemy.x > canvas.offsetWidth - enemy.w) enemy.dx *= -1;
            if (enemy.y > canvas.offsetHeight) {{ enemy.y = -60; enemy.x = Math.random() * (canvas.offsetWidth - enemy.w); enemy.entering = true; }}
        }});
        // Move bosses
        bossObjs.forEach(function(boss) {{
            boss.x += boss.dx;
            boss.y += boss.dy;
            if (boss.x < 0 || boss.x > canvas.offsetWidth - boss.w) boss.dx *= -1;
        }});
        // Move meteors
        meteorObjs.forEach(function(meteor) {{
            meteor.y += meteor.dy;
            if (meteor.y > canvas.offsetHeight) {{ meteor.y = -60; meteor.x = Math.random() * (canvas.offsetWidth - meteor.w); }}
        }});
        // Move refills
        refillObjs.forEach(function(refill) {{
            refill.y += refill.dy;
            if (refill.y > canvas.offsetHeight) {{ refill.y = -60; refill.x = Math.random() * (canvas.offsetWidth - refill.w); }}
        }});
        // Move holes
        holeObjs.forEach(function(hole) {{
            hole.y += hole.dy;
            if (hole.y > canvas.offsetHeight) {{ hole.y = -60; hole.x = Math.random() * (canvas.offsetWidth - hole.w); }}
        }});
        // Move scores
        scoreObjs.forEach(function(score) {{
            score.y += score.dy;
            if (score.y > canvas.offsetHeight) {{ score.y = -60; score.x = Math.random() * (canvas.offsetWidth - score.w); }}
        }});


        // Bullet-enemy collision
        let newBullets = [];
        bullets.forEach(function(bullet) {{
            let hit = false;
            for (let i = 0; i < enemyObjs.length; i++) {{
                let enemy = enemyObjs[i];
                if (rectsOverlap({{x: bullet.x-5, y: bullet.y-5, w: bulletSize, h: bulletSize}}, enemy)) {{
                    // Add explosion animation at enemy position
                    explosions.push({{
                        x: enemy.x,
                        y: enemy.y,
                        frame: 0
                    }});
                    playExplosion();
                    enemyObjs.splice(i, 1);
                    score += 100;
                    hit = true;
                    break;
                }}
            }}
            if (!hit) newBullets.push(bullet);
        }});
        bullets = newBullets;

        // Check collision with meteors
        meteorObjs.forEach(function(meteor) {{
            if (rectsOverlap(playerRect(), meteor)) {{
                explosions.push({{ x: playerX, y: playerY, frame: 0 }});
                playExplosion();
                gameOver = true;
                playGameOver();
            }}
        }});
        // Check collision with refills
        refillObjs.forEach(function(refill) {{
            if (rectsOverlap(playerRect(), refill)) {{
                explosions.push({{ x: playerX, y: playerY, frame: 0 }});
                playExplosion();
                gameOver = true;
                playGameOver();
            }}
        }});
        // Check collision with holes
        holeObjs.forEach(function(hole) {{
            if (rectsOverlap(playerRect(), hole)) {{
                explosions.push({{ x: playerX, y: playerY, frame: 0 }});
                playExplosion();
                gameOver = true;
                playGameOver();
            }}
        }});
        // Check collision with score coins
        scoreObjs.forEach(function(score) {{
            if (rectsOverlap(playerRect(), score)) {{
                explosions.push({{ x: playerX, y: playerY, frame: 0 }});
                playExplosion();
                gameOver = true;
                playGameOver();
            }}
        }});

        // Check collision with enemies
        enemyObjs.forEach(function(enemy) {{
            if (rectsOverlap(playerRect(), enemy)) {{
                explosions.push({{ x: playerX, y: playerY, frame: 0 }});
                playExplosion();
                gameOver = true;
                playGameOver();
            }}
        }});

        // Check collision with bosses
        bossObjs.forEach(function(boss) {{
            if (rectsOverlap(playerRect(), boss)) {{
                explosions.push({{ x: playerX, y: playerY, frame: 0 }});
                playExplosion();
                gameOver = true;
                playGameOver();
            }}
        }});

        // Render explosions
        explosionsDiv.innerHTML = '';
        let newExplosions = [];
        explosions.forEach(function(explosion) {{
            if (explosion.frame < explosionImgs.length) {{
                let el = document.createElement('div');
                el.style.position = 'absolute';
                el.style.left = explosion.x + 'px';
                el.style.top = explosion.y + 'px';
                el.style.width = '50px';
                el.style.height = '50px';
                el.style.background = `url('${{explosionImgs[explosion.frame]}}') no-repeat`;
                el.style.backgroundSize = '100% 100%';
                el.style.zIndex = 5;
                explosionsDiv.appendChild(el);
                explosion.frame++;
                newExplosions.push(explosion);
            }}
        }});
        explosions = newExplosions;

        // Level up if all enemies are gone
        if (enemyObjs.length === 0) {{
            level++;
            setBackground();
            spawnEnemies();
            // Optionally: spawnBoss();
        }}

        // Draw bullets
        ctx.clearRect(0, 0, bulletCanvas.width, bulletCanvas.height);
        ctx.fillStyle = 'yellow';
        bullets.forEach(function(bullet) {{
            ctx.beginPath();
            ctx.arc(bullet.x, bullet.y, bulletSize / 2, 0, Math.PI * 2);
            ctx.fill();
        }});

        // Render enemies
        enemiesDiv.innerHTML = '';
        enemyObjs.forEach(function(enemy) {{
            let el = document.createElement('div');
            el.style.position = 'absolute';
            el.style.left = enemy.x + 'px';
            el.style.top = enemy.y + 'px';
            el.style.width = enemy.w + 'px';
            el.style.height = enemy.h + 'px';
            el.style.background = `url('${{enemy.img}}') no-repeat`;
            el.style.backgroundSize = '100% 100%';
            el.style.zIndex = 2;
            enemiesDiv.appendChild(el);
        }});
        // Render bosses
        bossesDiv.innerHTML = '';
        bossObjs.forEach(function(boss) {{
            let el = document.createElement('div');
            el.style.position = 'absolute';
            el.style.left = boss.x + 'px';
            el.style.top = boss.y + 'px';
            el.style.width = boss.w + 'px';
            el.style.height = boss.h + 'px';
            el.style.background = `url('${{boss.img}}') no-repeat`;
            el.style.backgroundSize = '100% 100%';
            el.style.zIndex = 2;
            bossesDiv.appendChild(el);
        }});
        // Render meteors
        meteorsDiv.innerHTML = '';
        meteorObjs.forEach(function(meteor) {{
            let el = document.createElement('div');
            el.style.position = 'absolute';
            el.style.left = meteor.x + 'px';
            el.style.top = meteor.y + 'px';
            el.style.width = meteor.w + 'px';
            el.style.height = meteor.h + 'px';
            el.style.background = `url('${{meteor.img}}') no-repeat`;
            el.style.backgroundSize = '100% 100%';
            el.style.zIndex = 2;
            meteorsDiv.appendChild(el);
        }});
        // Render refills
        refillsDiv.innerHTML = '';
        refillObjs.forEach(function(refill) {{
            let el = document.createElement('div');
            el.style.position = 'absolute';
            el.style.left = refill.x + 'px';
            el.style.top = refill.y + 'px';
            el.style.width = refill.w + 'px';
            el.style.height = refill.h + 'px';
            el.style.background = `url('${{refill.img}}') no-repeat`;
            el.style.backgroundSize = '100% 100%';
            el.style.zIndex = 2;
            refillsDiv.appendChild(el);
        }});
        // Render holes
        holesDiv.innerHTML = '';
        holeObjs.forEach(function(hole) {{
            let el = document.createElement('div');
            el.style.position = 'absolute';
            el.style.left = hole.x + 'px';
            el.style.top = hole.y + 'px';
            el.style.width = hole.w + 'px';
            el.style.height = hole.h + 'px';
            el.style.background = `url('${{hole.img}}') no-repeat`;
            el.style.backgroundSize = '100% 100%';
            el.style.zIndex = 2;
            holesDiv.appendChild(el);
        }});
        // Render scores
        scoresDiv.innerHTML = '';
        scoreObjs.forEach(function(score) {{
            let el = document.createElement('div');
            el.style.position = 'absolute';
            el.style.left = score.x + 'px';
            el.style.top = score.y + 'px';
            el.style.width = score.w + 'px';
            el.style.height = score.h + 'px';
            el.style.background = `url('${{score.img}}') no-repeat`;
            el.style.backgroundSize = '100% 100%';
            el.style.zIndex = 2;
            scoresDiv.appendChild(el);
        }});

        // Create data URL of the canvas for Streamlit
        const frameData = bulletCanvas.toDataURL();
        gameFrame.value = `${{playerX}},${{playerY}},${{frameData}},${{bulletCounter}}`;
    }}

    // Update player position
    function updatePlayer() {{
        player.style.left = `${{playerX}}px`;
        player.style.top = `${{playerY}}px`;
    }}

    // Game loop
    function gameLoop() {{
        if (!gameOver) {{
            if (firing) {{
                const now = Date.now();
                if (now - lastShotTime > 50 && bulletCounter > 0) {{
                    bullets.push({{ x: playerX + 20, y: playerY, dy: -bulletSpeed }});
                    bulletCounter--;
                    lastShotTime = now;
                    playShoot();
                }}
            }}
            updateGame();
            requestAnimationFrame(gameLoop);
        }} else {{
            // Show Game Over message
            let overDiv = document.getElementById('game-over');
            if (!overDiv) {{
                overDiv = document.createElement('div');
                overDiv.id = 'game-over';
                overDiv.style.position = 'absolute';
                overDiv.style.top = '40%';
                overDiv.style.left = '0';
                overDiv.style.width = '100%';
                overDiv.style.textAlign = 'center';
                overDiv.style.fontSize = '48px';
                overDiv.style.color = 'red';
                overDiv.style.fontWeight = 'bold';
                overDiv.style.zIndex = 100;
                overDiv.innerText = 'GAME OVER';
                canvas.appendChild(overDiv);
            }}
        }}
    }}

    // Resize canvas on window resize
    window.addEventListener('resize', function() {{
        bulletCanvas.width = canvas.offsetWidth;
        bulletCanvas.height = canvas.offsetHeight;
    }});

    // Initialize

    bulletCanvas.width = canvas.offsetWidth;
    bulletCanvas.height = canvas.offsetHeight;
    updatePlayer();
    gameLoop();

    // Only allow sound after first user interaction (browser autoplay policy)
    let soundInitialized = false;
    function initSound() {{
        if (!soundInitialized) {{
            document.getElementById('bgm-audio').volume = 0.25;
            document.getElementById('bgm-audio').play();
            soundInitialized = true;
        }}
    }}
    // Listen for first mousedown on canvas to enable sound
    canvas.addEventListener('mousedown', initSound, {{ once: true }});

    // Play shoot sound
    function playShoot() {{
        let audio = document.getElementById('shoot-audio');
        audio.currentTime = 0;
        audio.play();
    }}

    // Play explosion sound
    function playExplosion() {{
        let audio = document.getElementById('explosion-audio');
        audio.currentTime = 0;
        audio.play();
    }}

    // Play game over sound
    function playGameOver() {{
        let audio = document.getElementById('gameover-audio');
        audio.currentTime = 0;
        audio.play();
        document.getElementById('bgm-audio').pause();
    }}
</script>
""", height=600).value

    if game_data and isinstance(game_data, str):
        try:
            player_x, player_y, frame_data, bullet_count = game_data.split(',')
            player_x, player_y = map(int, (player_x, player_y))
            bullet_count = int(bullet_count)
            # Render the frame using the data URL
            response = requests.get(frame_data)
            frame = Image.open(io.BytesIO(response.content))
            st.image(frame, caption='Cosmic Heat Frame', use_container_width=True)
            st.write(f"Bullets: {bullet_count}")
        except ValueError as e:
            st.write("Error processing frame:", e)
            st.image(Image.open('images/bg/background.jpg'), caption='Cosmic Heat Frame', use_container_width=True)
            st.write("Bullets: 200")

    if st.button('Back to Menu'):
        st.session_state.page = 'menu'
        st.rerun()