import pygame
import random
import sys
import math

# ========== 初始化 ==========
pygame.init()
WIDTH, HEIGHT = 800, 500
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("🔥 魂斗罗风格 – Jungle Assault")
clock = pygame.time.Clock()

# ========== 色彩 ==========
SKY_TOP = (80, 140, 200)
SKY_BOTTOM = (180, 210, 240)
MOUNTAIN_COLOR = (100, 160, 100)
MOUNTAIN_DARK = (70, 130, 70)
GROUND = (90, 150, 60)
DIRT = (120, 80, 40)
PLAYER_BODY = (40, 120, 200)
PLAYER_SKIN = (255, 200, 150)
ENEMY_COLOR = (200, 60, 40)
BULLET_PLAYER = (255, 240, 100)
BULLET_ENEMY = (255, 100, 100)
EXPLOSION = (255, 200, 50)
FLAME = (255, 180, 40)
WHITE = (255, 255, 255)
BLACK = (20, 20, 20)
RED = (220, 50, 50)

# ========== 字体 ==========
font_score = pygame.font.Font(None, 40)
font_big = pygame.font.Font(None, 60)

# ========== 常量 ==========
GRAVITY = 0.8
PLAYER_SPEED = 5
JUMP_POWER = -14
SHOOT_COOLDOWN = 12
ENEMY_SHOOT_COOLDOWN = 40
SCROLL_SPEED = 2

# ========== 背景元素 ==========
class Background:
    def __init__(self):
        # 树木：列表的列表 [x, y, height]
        self.trees = []
        for _ in range(6):
            x = random.randint(0, WIDTH)
            y = random.randint(280, 380)
            h = random.randint(70, 110)
            self.trees.append([x, y, h])
        # 云朵：列表的列表 [x, y, size]
        self.clouds = [[random.randint(0, WIDTH), random.randint(30, 100), random.randint(40, 60)] for _ in range(5)]

    def update(self):
        # 树木移动
        for tree in self.trees:
            tree[0] -= SCROLL_SPEED * 0.4
            if tree[0] < -60:
                tree[0] = WIDTH + random.randint(20, 80)
                tree[1] = random.randint(280, 380)
                tree[2] = random.randint(70, 110)
        # 云移动
        for cloud in self.clouds:
            cloud[0] -= SCROLL_SPEED * 0.15
            if cloud[0] < -80:
                cloud[0] = WIDTH + random.randint(20, 100)
                cloud[1] = random.randint(30, 100)
                cloud[2] = random.randint(40, 60)

    def draw(self, surface):
        # 天空渐变
        for y in range(300):
            ratio = y / 300
            r = int(SKY_TOP[0] + (SKY_BOTTOM[0]-SKY_TOP[0])*ratio)
            g = int(SKY_TOP[1] + (SKY_BOTTOM[1]-SKY_TOP[1])*ratio)
            b = int(SKY_TOP[2] + (SKY_BOTTOM[2]-SKY_TOP[2])*ratio)
            pygame.draw.line(surface, (r,g,b), (0,y), (WIDTH,y))
        # 云朵
        for x, y, s in self.clouds:
            pygame.draw.ellipse(surface, WHITE, (x, y, s*1.5, s*0.6))
            pygame.draw.ellipse(surface, WHITE, (x+10, y-8, s, s*0.5))
        # 远山
        for i in range(4):
            offset = (SCROLL_SPEED * 0.2 * i) % 200
            px = 100 + i*200 - offset
            pygame.draw.polygon(surface, MOUNTAIN_COLOR if i%2==0 else MOUNTAIN_DARK,
                                [(px-80,300),(px+20,200),(px+120,300)])
        # 草地
        pygame.draw.rect(surface, GROUND, (0, 400, WIDTH, 100))
        # 地面纹理
        for x in range(0, WIDTH, 30):
            h = random.randint(3, 8)
            pygame.draw.line(surface, (60,120,40), (x, 400), (x+10, 400-h), 2)
        # 树木
        for x, y, h in self.trees:
            trunk = pygame.Rect(x-8, y-h, 16, h)
            pygame.draw.rect(surface, (100,70,40), trunk)
            pygame.draw.ellipse(surface, (40,110,40), (x-30, y-h-30, 60, 50))
            pygame.draw.ellipse(surface, (30,90,30), (x-20, y-h-25, 40, 40))

# ========== 粒子 ==========
class Particle:
    def __init__(self, x, y, color, size=4):
        self.x, self.y = x, y
        self.vx = random.uniform(-3, 3)
        self.vy = random.uniform(-7, -2)
        self.color = color
        self.life = 20
        self.size = size
    def update(self):
        self.x += self.vx
        self.y += self.vy
        self.vy += 0.3
        self.life -= 1
    def draw(self, surf):
        if self.life > 0:
            alpha = int(255 * self.life / 20)
            pygame.draw.circle(surf, (*self.color, alpha), (int(self.x), int(self.y)), self.size)

# ========== 玩家 ==========
class Player:
    def __init__(self):
        self.x, self.y = 150, 360
        self.vx, self.vy = 0, 0
        self.on_ground = False
        self.shoot_timer = 0
        self.lives = 3
        self.invincible = 0
        self.dir = 1
        self.muzzle_flash = 0

    def update(self, keys):
        self.vx = 0
        if keys[pygame.K_LEFT] or keys[pygame.K_a]:
            self.vx = -PLAYER_SPEED
            self.dir = -1
        if keys[pygame.K_RIGHT] or keys[pygame.K_d]:
            self.vx = PLAYER_SPEED
            self.dir = 1
        if (keys[pygame.K_UP] or keys[pygame.K_w]) and self.on_ground:
            self.vy = JUMP_POWER
            self.on_ground = False
        self.vy += GRAVITY
        self.x += self.vx
        self.y += self.vy
        self.x = max(30, min(WIDTH-30, self.x))
        if self.y >= 400:
            self.y = 400
            self.vy = 0
            self.on_ground = True
        if self.shoot_timer > 0:
            self.shoot_timer -= 1
        if self.invincible > 0:
            self.invincible -= 1
        if self.muzzle_flash > 0:
            self.muzzle_flash -= 1

    def shoot(self):
        if self.shoot_timer <= 0:
            self.shoot_timer = SHOOT_COOLDOWN
            self.muzzle_flash = 4
            return True
        return False

    def draw(self, surface):
        if self.invincible > 0 and self.invincible % 4 < 2:
            return
        pygame.draw.ellipse(surface, (0,0,0,80), (self.x-12, self.y+5, 24, 6))
        body_rect = pygame.Rect(self.x-12, self.y-24, 24, 30)
        pygame.draw.rect(surface, PLAYER_BODY, body_rect)
        pygame.draw.rect(surface, BLACK, body_rect, 2)
        pygame.draw.rect(surface, (80,40,20), (self.x-12, self.y-10, 24, 6))
        pygame.draw.circle(surface, PLAYER_SKIN, (self.x, self.y-32), 10)
        pygame.draw.circle(surface, BLACK, (self.x, self.y-32), 10, 2)
        pygame.draw.rect(surface, RED, (self.x-12, self.y-42, 24, 8), border_radius=3)
        eye_x = self.x + 4*self.dir
        pygame.draw.circle(surface, BLACK, (eye_x, self.y-34), 2)
        gun_x = self.x + 16*self.dir
        pygame.draw.rect(surface, (60,60,60), (gun_x-2, self.y-20, 14, 6))
        if self.muzzle_flash > 0:
            flame_x = gun_x + 12*self.dir
            pygame.draw.circle(surface, FLAME, (flame_x, self.y-17), 6)
            pygame.draw.circle(surface, (255,255,100), (flame_x, self.y-17), 3)
        leg_offset = 3 if self.on_ground else 0
        pygame.draw.line(surface, BLACK, (self.x-6, self.y-6), (self.x-8, self.y+8+leg_offset), 4)
        pygame.draw.line(surface, BLACK, (self.x+6, self.y-6), (self.x+8, self.y+8-leg_offset), 4)

    def get_rect(self):
        return pygame.Rect(self.x-12, self.y-24, 24, 30)

# ========== 敌人 ==========
class Enemy:
    def __init__(self, x, y):
        self.x, self.y = x, y
        self.vx = -2
        self.shoot_timer = random.randint(20, ENEMY_SHOOT_COOLDOWN)
        self.alive = True

    def update(self):
        self.x += self.vx
        if self.shoot_timer > 0:
            self.shoot_timer -= 1
        if self.x < -30:
            self.alive = False

    def can_shoot(self):
        return self.shoot_timer <= 0

    def draw(self, surface):
        body_rect = pygame.Rect(self.x-10, self.y-20, 20, 26)
        pygame.draw.rect(surface, ENEMY_COLOR, body_rect)
        pygame.draw.rect(surface, BLACK, body_rect, 2)
        pygame.draw.rect(surface, (100,30,20), (self.x-10, self.y-20, 20, 6))
        pygame.draw.circle(surface, (255,180,150), (self.x, self.y-28), 9)
        pygame.draw.circle(surface, BLACK, (self.x, self.y-28), 9, 2)
        pygame.draw.rect(surface, BLACK, (self.x-9, self.y-32, 18, 8))
        pygame.draw.rect(surface, (60,60,60), (self.x-14, self.y-18, 10, 5))

    def get_rect(self):
        return pygame.Rect(self.x-10, self.y-20, 20, 26)

# ========== 子弹 ==========
class Bullet:
    def __init__(self, x, y, vx, vy, color):
        self.x, self.y = x, y
        self.vx, self.vy = vx, vy
        self.color = color
        self.radius = 3

    def update(self):
        self.x += self.vx
        self.y += self.vy
        return self.x < 0 or self.x > WIDTH or self.y < 0 or self.y > HEIGHT

    def draw(self, surface):
        pygame.draw.circle(surface, self.color, (int(self.x), int(self.y)), self.radius)
        pygame.draw.circle(surface, WHITE, (int(self.x), int(self.y)), self.radius-1)

# ========== 游戏主类 ==========
class ContraGame:
    def __init__(self):
        self.bg = Background()
        self.player = Player()
        self.bullets = []
        self.enemies = []
        self.particles = []
        self.score = 0
        self.game_over = False
        self.spawn_timer = 0

    def spawn_enemy(self):
        y = random.choice([370, 340])
        self.enemies.append(Enemy(WIDTH+20, y))

    def handle_player_shoot(self):
        if self.player.shoot():
            bx = self.player.x + 16*self.player.dir
            by = self.player.y - 18
            self.bullets.append(Bullet(bx, by, 8*self.player.dir, 0, BULLET_PLAYER))

    def update(self):
        if self.game_over:
            return
        keys = pygame.key.get_pressed()
        self.player.update(keys)
        self.bg.update()

        if keys[pygame.K_SPACE] or keys[pygame.K_j]:
            self.handle_player_shoot()

        self.spawn_timer += 1
        if self.spawn_timer > 80 and len(self.enemies) < 5:
            self.spawn_timer = 0
            self.spawn_enemy()

        for b in self.bullets[:]:
            if b.update():
                self.bullets.remove(b)

        for enemy in self.enemies[:]:
            enemy.update()
            if not enemy.alive:
                self.enemies.remove(enemy)
                continue
            if enemy.can_shoot():
                enemy.shoot_timer = ENEMY_SHOOT_COOLDOWN + random.randint(0, 30)
                dx = self.player.x - enemy.x
                dy = self.player.y - enemy.y
                dist = math.hypot(dx, dy)
                if dist > 0:
                    vx = dx/dist * 4
                    vy = dy/dist * 4
                    self.bullets.append(Bullet(enemy.x, enemy.y-10, vx, vy, BULLET_ENEMY))
            if self.player.invincible <= 0 and self.player.get_rect().colliderect(enemy.get_rect()):
                self.player_hit()

        for b in self.bullets[:]:
            if b.color == BULLET_PLAYER:
                for enemy in self.enemies[:]:
                    if enemy.get_rect().collidepoint(b.x, b.y):
                        self.bullets.remove(b)
                        self.enemies.remove(enemy)
                        self.score += 100
                        for _ in range(12):
                            self.particles.append(Particle(enemy.x, enemy.y, EXPLOSION, 5))
                        break

        for b in self.bullets[:]:
            if b.color == BULLET_ENEMY and self.player.invincible <= 0:
                if self.player.get_rect().collidepoint(b.x, b.y):
                    self.bullets.remove(b)
                    self.player_hit()

        for p in self.particles[:]:
            p.update()
            if p.life <= 0:
                self.particles.remove(p)

    def player_hit(self):
        self.player.lives -= 1
        if self.player.lives <= 0:
            self.game_over = True
            for _ in range(20):
                self.particles.append(Particle(self.player.x, self.player.y-20, EXPLOSION, 6))
        else:
            self.player.invincible = 60
            self.player.vy = -8

    def draw(self):
        # 清屏防止拖尾
        screen.fill(SKY_TOP)

        self.bg.draw(screen)
        for b in self.bullets:
            b.draw(screen)
        for enemy in self.enemies:
            enemy.draw(screen)
        self.player.draw(screen)
        for p in self.particles:
            p.draw(screen)

        score_txt = font_score.render(f"Score: {self.score}", True, WHITE)
        screen.blit(score_txt, (20, 20))
        lives_txt = font_score.render(f"Lives: {self.player.lives}", True, RED)
        screen.blit(lives_txt, (WIDTH-150, 20))
        if self.game_over:
            over = font_big.render("GAME OVER", True, RED)
            screen.blit(over, (WIDTH//2-120, HEIGHT//2-30))
            restart = font_score.render("Press R to Restart", True, WHITE)
            screen.blit(restart, (WIDTH//2-120, HEIGHT//2+30))

    def reset(self):
        self.__init__()

# ========== 主循环 ==========
def main():
    game = ContraGame()
    running = True
    while running:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False
            if event.type == pygame.KEYDOWN:
                if event.key == pygame.K_r and game.game_over:
                    game.reset()
        game.update()
        game.draw()
        pygame.display.flip()
        clock.tick(60)
    pygame.quit()
    sys.exit()

if __name__ == "__main__":
    main()