import pygame
import math
import random

pygame.init()
WIDTH, HEIGHT = 1200, 700
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("Dave The Diver Enhanced")
clock = pygame.time.Clock()
FPS = 60

# 色彩
COLOR_SKY = (120, 190, 255)
COLOR_SHALLOW = (25, 110, 180)
COLOR_DEEP = (5, 30, 70)
COLOR_DIVER_BODY = (10, 90, 210)
COLOR_DIVER_MASK = (230, 245, 255)
COLOR_OXY_TANK = (70, 70, 80)
COLOR_SHARK = (35, 38, 42)
COLOR_SHARK_DARK = (15, 18, 20)
COLOR_CORAL = (255, 90, 130)
COLOR_GLOW_PLANT = (80, 255, 180)
COLOR_OXY_GREEN = (30, 210, 70)
COLOR_OXY_RED = (220, 20, 20)
COLOR_UI_WHITE = (255, 255, 255)
COLOR_FISH = [(255, 170, 30), (40, 210, 150), (210, 50, 110), (170, 90, 240)]
COLOR_BUBBLE = (200, 240, 255)
COLOR_GOLD = (255, 210, 40)

# 世界参数
CAM_X = 0
CAM_Y = 0
WORLD_W = 3600
WORLD_H = 3000
SEA_SURFACE_Y = 200

# 潜水员物理
diver_x = WORLD_W // 2
diver_y = 200
vel_x = 0.0
vel_y = 0.0
GRAVITY = 0.2
BUOYANCY = -0.09
WATER_DRAG = 0.95
MOVE_SPEED = 3.5

oxygen = 100
max_oxygen = 100
base_oxy_consume = 0.013
max_carry = 15
carry_count = 0
gold_coins = 0
hurt_cd = 0
flip_right = True
flash_red_alpha = 0
low_oxy_blink = False

# 环境粒子
bubble_particles = []
wave_points = []

# 实体列表
coral_list = []
rock_list = []
glow_plants = []
fish_list = []
shark_list = []
oxy_box_list = []
treasure_list = []

# 鱼叉数据
spear_fly = False
spear_x = 0
spear_y = 0
spear_vx = 0
spear_vy = 0
spear_max_dist = 450
charge_time = 0

# 水流全局扰动
current_x = random.uniform(-0.12, 0.12)
current_y = random.uniform(-0.06, 0.06)

# 生成海底地图
def reset_sea_map():
    global coral_list, rock_list, fish_list, shark_list, oxy_box_list, treasure_list, glow_plants, wave_points
    coral_list.clear()
    rock_list.clear()
    fish_list.clear()
    shark_list.clear()
    oxy_box_list.clear()
    treasure_list.clear()
    glow_plants.clear()
    wave_points.clear()
    bubble_particles.clear()

    # 海浪点
    for _ in range(400):
        wx = random.randint(-3000, WORLD_W+3000)
        wy = random.randint(SEA_SURFACE_Y, WORLD_H)
        ws = random.randint(1,4)
        wp = random.uniform(0, math.pi*2)
        wave_points.append([wx, wy, ws, wp])
    # 发光海底植物
    for _ in range(120):
        gx = random.randint(0, WORLD_W)
        gy = random.randint(400, WORLD_H)
        gs = random.randint(10,30)
        glow_plants.append([gx, gy, gs, random.uniform(0, math.pi*2)])
    # 珊瑚
    for _ in range(260):
        cx = random.randint(0, WORLD_W)
        cy = random.randint(300, WORLD_H)
        cs = random.randint(10, 48)
        coral_list.append([cx, cy, cs])
    # 礁石
    for _ in range(180):
        rx = random.randint(0, WORLD_W)
        ry = random.randint(260, WORLD_H)
        rs = random.randint(20, 75)
        rock_list.append([rx, ry, rs])
    # 小鱼
    for _ in range(110):
        fx = random.randint(50, WORLD_W - 50)
        fy = random.randint(240, WORLD_H - 200)
        fs = random.randint(7, 24)
        fsp = random.uniform(0.35, 1.3)
        dirr = random.choice([-1, 1])
        fc = random.choice(COLOR_FISH)
        fish_list.append([fx, fy, fs, fsp, dirr, fc, True, random.uniform(0, math.pi*2)])
    # 鲨鱼
    for _ in range(14):
        sx = random.randint(350, WORLD_W - 350)
        sy = random.randint(550, WORLD_H)
        ssp = random.uniform(1.0, 2.2)
        shark_list.append([sx, sy, ssp])
    # 氧气罐
    for _ in range(32):
        ox = random.randint(0, WORLD_W)
        oy = random.randint(320, WORLD_H)
        oxy_box_list.append([ox, oy])
    # 宝箱
    for _ in range(20):
        tx = random.randint(0, WORLD_W)
        ty = random.randint(380, WORLD_H)
        treasure_list.append([tx, ty, random.randint(10,40)])

reset_sea_map()

# 字体兼容兜底
try:
    font_main = pygame.font.SysFont("msyh", 24)
    font_small = pygame.font.SysFont("msyh", 16)
except:
    font_main = pygame.font.Font(None, 24)
    font_small = pygame.font.Font(None, 16)

# 生成气泡
def spawn_bubbles(x,y):
    for _ in range(2):
        bubble_particles.append({
            "x":x, "y":y,
            "vx":random.uniform(-0.3,0.3),
            "vy":random.uniform(-1.2,-0.4),
            "size":random.randint(2,6),
            "life":120
        })

# 绘制潜水员（修复多边形）
def draw_diver(screen_x, screen_y, flip):
    surf = pygame.Surface((48, 56), pygame.SRCALPHA)
    pygame.draw.ellipse(surf, COLOR_DIVER_BODY, (8, 14, 32, 36))
    pygame.draw.circle(surf, COLOR_DIVER_MASK, (24, 11), 10)
    pygame.draw.rect(surf, COLOR_OXY_TANK, (3, 16, 8, 28))
    # 脚蹼 标准成对坐标
    pygame.draw.polygon(surf, COLOR_DIVER_BODY, [(12, 48), (6, 56), (20, 50)])
    pygame.draw.polygon(surf, COLOR_DIVER_BODY, [(36, 48), (42, 56), (28, 50)])
    if not flip:
        surf = pygame.transform.flip(surf, True, False)
    screen.blit(surf, (screen_x - 24, screen_y - 28))

# 圆形碰撞检测
def circle_collide(x1, y1, r1, x2, y2, r2):
    return math.hypot(x1 - x2, y1 - y2) < r1 + r2

# 小地图绘制
def draw_minimap():
    map_w, map_h = 180, 120
    map_surf = pygame.Surface((map_w, map_h), pygame.SRCALPHA)
    map_surf.fill((0,20,40,180))
    scale_x = map_w / WORLD_W
    scale_y = map_h / WORLD_H
    # 礁石
    for rx,ry,rs in rock_list:
        mx = rx * scale_x
        my = ry * scale_y
        pygame.draw.circle(map_surf, (80,80,80), (int(mx),int(my)),2)
    # 鲨鱼红点
    for sx,sy,sp in shark_list:
        mx = sx * scale_x
        my = sy * scale_y
        pygame.draw.circle(map_surf, (255,30,30), (int(mx),int(my)),3)
    # 玩家
    px = diver_x * scale_x
    py = diver_y * scale_y
    pygame.draw.circle(map_surf, (40,220,80), (int(px),int(py)),4)
    screen.blit(map_surf, (WIDTH-map_w-10, HEIGHT-map_h-10))

running = True
time_global = 0
while running:
    delta = clock.tick(FPS) / 1000
    keys = pygame.key.get_pressed()
    time_global += delta
    t = time_global

    # 事件
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False
        # 鼠标左键蓄力发射鱼叉
        if event.type == pygame.MOUSEBUTTONDOWN and event.button == 1 and not spear_fly:
            charge_time = 0.3
        if event.type == pygame.MOUSEBUTTONUP and event.button == 1 and charge_time>0:
            mx, my = pygame.mouse.get_pos()
            world_mx = mx + CAM_X
            world_my = my + CAM_Y
            angle = math.atan2(world_my - diver_y, world_mx - diver_x)
            power_mult = min(1.8, 1 + charge_time*1.2)
            spear_fly = True
            spear_x = diver_x
            spear_y = diver_y
            spear_vx = math.cos(angle) * 15 * power_mult
            spear_vy = math.sin(angle) * 15 * power_mult
            charge_time = 0
    # 蓄力计时
    if charge_time > 0:
        charge_time -= delta

    # 深度压力：越深氧气消耗越快
    depth = max(0, diver_y - SEA_SURFACE_Y)
    depth_oxy_mod = 1 + depth / 1200
    oxygen -= base_oxy_consume * depth_oxy_mod
    # 低氧闪烁标记
    low_oxy_blink = oxygen < 30

    # 氧气归零重置
    if oxygen <= 0:
        oxygen = max_oxygen
        diver_x = WORLD_W // 2
        diver_y = SEA_SURFACE_Y + 50
        carry_count = 0
        reset_sea_map()

    # 上浮上岸回城（海面上方恢复全部氧气，清空鱼货换金币）
    if diver_y < SEA_SURFACE_Y + 20:
        gold_coins += carry_count * 12
        carry_count = 0
        oxygen = min(max_oxygen, oxygen + 0.3)

    # 玩家移动输入
    move_x = 0
    move_y = 0
    if keys[pygame.K_a] or keys[pygame.K_LEFT]:
        move_x = -1
        flip_right = False
    if keys[pygame.K_d] or keys[pygame.K_RIGHT]:
        move_x = 1
        flip_right = True
    if keys[pygame.K_w] or keys[pygame.K_UP]:
        move_y = -1
    if keys[pygame.K_s] or keys[pygame.K_DOWN]:
        move_y = 1

    # 水下物理 + 全局水流扰动
    vel_x += move_x * MOVE_SPEED * delta * 60
    vel_y += GRAVITY + move_y * MOVE_SPEED * delta * 60
    vel_y += BUOYANCY
    vel_x += current_x
    vel_y += current_y
    vel_x *= WATER_DRAG
    vel_y *= WATER_DRAG

    # 更新潜水员位置+边界限制
    diver_x += vel_x
    diver_y += vel_y
    diver_x = max(40, min(WORLD_W - 40, diver_x))
    diver_y = max(SEA_SURFACE_Y, min(WORLD_H - 40, diver_y))

    # 持续生成气泡
    if random.random() < 0.08:
        spawn_bubbles(diver_x, diver_y+20)

    # 镜头跟随
    CAM_X = diver_x - WIDTH // 2
    CAM_Y = diver_y - HEIGHT // 2

    # 受伤冷却 & 受伤泛红衰减
    if hurt_cd > 0:
        hurt_cd -= delta
        flash_red_alpha = min(120, flash_red_alpha + 8)
    else:
        flash_red_alpha = max(0, flash_red_alpha - 6)

    # 鲨鱼AI追击 + 危险预警红光
    shark_near = False
    for shark in shark_list:
        sx, sy, sp = shark
        dist_shark = math.hypot(diver_x-sx, diver_y-sy)
        if dist_shark < 350:
            shark_near = True
        ang = math.atan2(diver_y - sy, diver_x - sx)
        shark[0] += math.cos(ang) * sp
        shark[1] += math.sin(ang) * sp
        # 撞击扣氧气+击退
        if circle_collide(diver_x, diver_y, 20, sx, sy, 35) and hurt_cd <= 0:
            oxygen -= 25
            hurt_cd = 1.3
            vel_x *= -1.4
            vel_y *= -1.4

    # 小鱼摆动动画、鱼叉捕捉
    for fish in fish_list:
        fx, fy, fs, fsp, fdir, fcol, alive, wave_t = fish
        if not alive:
            continue
        fish[7] += delta * 2.5
        fish[0] += fdir * fsp
        # 左右边界反弹
        if fx < 80 or fx > WORLD_W - 80:
            fish[4] *= -1
        # 鱼叉命中捕获
        if spear_fly and circle_collide(spear_x, spear_y, 7, fx, fy, fs):
            if carry_count < max_carry:
                fish[6] = False
                carry_count += 1
            spear_fly = False

    # 鱼叉飞行逻辑
    if spear_fly:
        spear_x += spear_vx
        spear_y += spear_vy
        dist = math.hypot(spear_x - diver_x, spear_y - diver_y)
        if dist > spear_max_dist:
            spear_fly = False
        # 撞礁石消失
        for (rx, ry, rs) in rock_list:
            if circle_collide(spear_x, spear_y, 7, rx, ry, rs):
                spear_fly = False
                break

    # 拾取氧气箱
    new_oxy = []
    for ox, oy in oxy_box_list:
        if circle_collide(diver_x, diver_y, 20, ox, oy, 18):
            oxygen = min(max_oxygen, oxygen + 38)
        else:
            new_oxy.append([ox, oy])
    oxy_box_list = new_oxy

    # 拾取宝箱金币
    new_treasure = []
    for tx, ty, gold in treasure_list:
        if circle_collide(diver_x, diver_y, 20, tx, ty, 22):
            gold_coins += gold
        else:
            new_treasure.append([tx, ty, gold])
    treasure_list = new_treasure

    # 更新气泡粒子
    new_bub = []
    for b in bubble_particles:
        b["x"] += b["vx"]
        b["y"] += b["vy"]
        b["life"] -= 1
        if b["life"] > 0 and b["y"] > SEA_SURFACE_Y:
            new_bub.append(b)
    bubble_particles = new_bub

    # ============ 分层绘制 ============
    # 海面天空
    screen.fill(COLOR_SKY)
    sea_line = SEA_SURFACE_Y - CAM_Y
    # 浅水半透明层
    shallow_surface = pygame.Surface((WIDTH, HEIGHT - sea_line), pygame.SRCALPHA)
    shallow_surface.fill((*COLOR_SHALLOW, 180))
    screen.blit(shallow_surface, (0, sea_line))
    # 深海雾化，越深越暗
    deep_fog = pygame.Surface((WIDTH, HEIGHT), pygame.SRCALPHA)
    fog_alpha = min(170, diver_y / WORLD_H * 150)
    deep_fog.fill((*COLOR_DEEP, int(fog_alpha)))
    screen.blit(deep_fog, (0, 0))

    # 流动海浪波纹
    for wx, wy, ws, wp in wave_points:
        sx = wx - CAM_X
        sy = wy - CAM_Y + math.sin(t*1.2 + wp)*3
        pygame.draw.circle(screen, (160,220,255,100), (int(sx), int(sy)), ws)

    # 发光海底植物
    for gx, gy, gs, gp in glow_plants:
        sx = gx - CAM_X
        sy = gy - CAM_Y
        pulse = math.sin(t*2.2 + gp) * 0.4 + 1
        size_pulse = int(gs * pulse)
        pygame.draw.circle(screen, COLOR_GLOW_PLANT, (int(sx), int(sy)), size_pulse, 2)

    # 珊瑚、礁石、氧气箱、宝箱
    for cx, cy, cs in coral_list:
        sx = cx - CAM_X
        sy = cy - CAM_Y
        pygame.draw.circle(screen, COLOR_CORAL, (int(sx), int(sy)), cs)
    for rx, ry, rs in rock_list:
        sx = rx - CAM_X
        sy = ry - CAM_Y
        pygame.draw.circle(screen, (45, 48, 52), (int(sx), int(sy)), rs)
        pygame.draw.circle(screen, (20, 22, 25), (int(sx), int(sy)), rs - 9)
    for ox, oy in oxy_box_list:
        sx = ox - CAM_X
        sy = oy - CAM_Y
        pygame.draw.rect(screen, COLOR_OXY_GREEN, [sx - 13, sy - 13, 26, 26])
    for tx, ty, g in treasure_list:
        sx = tx - CAM_X
        sy = ty - CAM_Y
        pygame.draw.rect(screen, COLOR_GOLD, [sx - 15, sy - 11, 30, 22])

    # 小鱼（带摆动）
    for fish in fish_list:
        fx, fy, fs, fsp, fdir, fcol, alive, wave_t = fish
        if not alive:
            continue
        sx = fx - CAM_X
        sy = fy - CAM_Y + math.sin(wave_t)*2
        pygame.draw.ellipse(screen, fcol, [sx - fs, sy - fs//2, fs*2, fs])

    # 鲨鱼 + 预警红光
    if shark_near:
        warn_surf = pygame.Surface((WIDTH, HEIGHT), pygame.SRCALPHA)
        warn_surf.fill((255,0,0,40))
        screen.blit(warn_surf, (0,0))
    for sx, sy, sp in shark_list:
        draw_x = sx - CAM_X
        draw_y = sy - CAM_Y
        pygame.draw.ellipse(screen, COLOR_SHARK, [draw_x - 34, draw_y - 16, 68, 32])
        tail_points = [(draw_x+34, draw_y), (draw_x+50, draw_y-14), (draw_x+50, draw_y+14)]
        pygame.draw.polygon(screen, COLOR_SHARK_DARK, tail_points)

    # 气泡粒子
    for b in bubble_particles:
        bx = b["x"] - CAM_X
        by = b["y"] - CAM_Y
        alpha = int(255 * (b["life"] / 120))
        bub_surf = pygame.Surface((b["size"]*2, b["size"]*2), pygame.SRCALPHA)
        pygame.draw.circle(bub_surf, (*COLOR_BUBBLE, alpha), (b["size"], b["size"]), b["size"])
        screen.blit(bub_surf, (bx, by))

    # 飞行鱼叉 + 蓄力指示器
    if spear_fly:
        sx = spear_x - CAM_X
        sy = spear_y - CAM_Y
        pygame.draw.line(screen, (240,240,240), (sx,sy), (sx-spear_vx*2.8, sy-spear_vy*2.8), 5)
    if charge_time > 0:
        charge_len = int(charge_time / 0.3 * 80)
        pygame.draw.rect(screen, (255,200,0), (WIDTH//2-40, HEIGHT//2+40, charge_len, 6))

    # 玩家潜水员
    draw_diver(WIDTH//2, HEIGHT//2, flip_right)

    # 受伤全屏泛红
    if flash_red_alpha > 0:
        hurt_surf = pygame.Surface((WIDTH, HEIGHT), pygame.SRCALPHA)
        hurt_surf.fill((255,0,0,int(flash_red_alpha)))
        screen.blit(hurt_surf, (0,0))

    # ============ UI界面 ============
    # 氧气条，低氧闪烁
    bar_w = 230
    bar_h = 24
    fill_w = bar_w * (oxygen / max_oxygen)
    oxy_back = COLOR_OXY_RED if (low_oxy_blink and int(t*4)%2==0) else COLOR_OXY_RED
    pygame.draw.rect(screen, oxy_back, (12, 12, bar_w, bar_h))
    pygame.draw.rect(screen, COLOR_OXY_GREEN, (12, 12, fill_w, bar_h))
    pygame.draw.rect(screen, COLOR_UI_WHITE, (12, 12, bar_w, bar_h), 2)
    oxy_text = font_main.render(f"Oxygen: {oxygen:.0f}%", True, COLOR_UI_WHITE)
    screen.blit(oxy_text, (250, 10))

    # 捕获、金币、深度
    carry_text = font_main.render(f"Catch: {carry_count}/{max_carry}", True, COLOR_UI_WHITE)
    gold_text = font_main.render(f"Gold: {gold_coins}", True, COLOR_GOLD)
    depth_text = font_small.render(f"Depth: {int(depth)}m", True, COLOR_UI_WHITE)
    screen.blit(carry_text, (12, 42))
    screen.blit(gold_text, (12, 68))
    screen.blit(depth_text, (12, 94))

    # 操作提示
    hint = font_small.render("WASD Move | Hold LMB Charge Spear | Surface to Sell Fish", True, COLOR_UI_WHITE)
    screen.blit(hint, (12, HEIGHT - 30))

    # 右下角小地图
    draw_minimap()

    pygame.display.update()

pygame.quit()