import pygame
import random
import math

# 基础初始化
pygame.init()
pygame.font.init()
WIDTH, HEIGHT = 800, 600
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("太空综合发射模拟器")
clock = pygame.time.Clock()
FPS = 60

# 颜色定义
BLACK = (0, 0, 0)
WHITE = (255, 255, 255)
RED = (255, 60, 0)
ORANGE = (255, 160, 0)
YELLOW = (255, 255, 0)
GREEN = (0, 255, 0)
BROWN = (80, 40, 20)
BLUE = (0, 150, 255)
GRAY = (120, 120, 120)
DARK_GRAY = (60, 60, 60)
CRASH_RED = (220, 30, 30)
LIGHT_GREEN = (80, 220, 80)
BASE_YELLOW = (220, 180, 0)
SKY_DAY = (10, 20, 60)
SKY_NIGHT = (0, 0, 10)

# 重力设置
EARTH_GRAVITY = 0.11
MOON_GRAVITY = 0.04
current_gravity = EARTH_GRAVITY
is_moon_mode = False

# 三款强化火箭配置（外观全面升级）
rocket_types = [
    {"name": "Basic", "body_color": WHITE, "boost": 1.0, "fuel_cap": 100, "width":34,"height":65},
    {"name": "Heavy", "body_color": GRAY, "boost": 0.85, "fuel_cap": 140, "width":40,"height":72},
    {"name": "Speed", "body_color": BLUE, "boost": 1.2, "fuel_cap": 75, "width":28,"height":58}
]
selected_rocket = 0

# 火箭物理变量
rocket_x = WIDTH // 2
rocket_y = HEIGHT - 120
speed_y = 0
speed_x = 0
thrust_up_base = -0.48
thrust_side_base = 0.22
air_resistance = 0.025
fuel = 100
max_fuel = rocket_types[selected_rocket]["fuel_cap"]
score = 0
ground_level = HEIGHT - 100
max_vertical_speed = 12
max_horizontal_speed = 6
stage2_separated = False
stage2_y_offset = 62

# 视角系统：永久锁定火箭平滑跟随
camera_offset_y = 0
camera_offset_x = 0
camera_smooth = 0.09
lock_camera = True  # 默认锁定火箭

# 粒子系统
explode_particles = []
launch_smoke = []
weather_particles = []

# 昼夜时间系统
day_time = 0
day_speed = 0.03
is_night = False

# 天气系统：晴天/小雨/暴雨
weather_state = 0  #0晴天 1小雨 2暴雨
weather_timer = 0

# 卫星系统
satellite_spawned = False
satellite_x = 0
satellite_y = 0
satellite_vx = 0
satellite_vy = 0
satellite_orbit = False

# 空间站固定轨道位置
station_x = WIDTH // 2 + 120
station_y = HEIGHT // 2 - 100

# 星空背景
stars = []
for _ in range(220):
    x = random.randint(0, WIDTH)
    y = random.randint(0, HEIGHT // 2)
    size = random.randint(1, 3)
    twinkle = random.randint(0, 100)
    stars.append([x, y, size, twinkle])

font = pygame.font.Font(None, 24)

def reset_rocket():
    global rocket_x, rocket_y, speed_x, speed_y, fuel, stage2_separated, camera_offset_y, camera_offset_x
    global launch_smoke, satellite_spawned, satellite_orbit
    rocket_x = WIDTH // 2
    rocket_y = HEIGHT - 120
    speed_x = 0
    speed_y = 0
    max_fuel = rocket_types[selected_rocket]["fuel_cap"]
    fuel = max_fuel
    stage2_separated = False
    camera_offset_y = 0
    camera_offset_x = 0
    launch_smoke.clear()
    satellite_spawned = False
    satellite_orbit = False

def spawn_explosion(x, y):
    for _ in range(35):
        angle = random.uniform(0, 2 * math.pi)
        spd = random.uniform(2, 8)
        vx = spd * math.cos(angle)
        vy = spd * math.sin(angle)
        life = random.randint(25, 55)
        color = random.choice([RED, ORANGE, YELLOW])
        explode_particles.append([x, y, vx, vy, life, color])

def spawn_launch_smoke(x, y):
    for _ in range(7):
        sx = x + random.randint(-22, 22)
        sy = y
        svx = random.uniform(-1.2, 1.2)
        svy = random.uniform(-3.5, -1.2)
        slife = random.randint(30, 50)
        launch_smoke.append([sx, sy, svx, svy, slife])

def spawn_rain():
    # 生成雨滴粒子
    for _ in range(15):
        rx = random.randint(0, WIDTH)
        ry = random.randint(-50, HEIGHT)
        rvx = random.uniform(-0.5, 0.5)
        rvy = random.uniform(6, 12)
        weather_particles.append([rx, ry, rvx, rvy])

def draw_space_station(offset_y, offset_x):
    """绘制空间站"""
    sx = station_x + offset_x
    sy = station_y + offset_y
    # 主体舱段
    pygame.draw.rect(screen, GRAY, (sx-35, sy-12, 70, 24))
    # 太阳能帆板
    pygame.draw.rect(screen, BLUE, (sx-70, sy-6, 30, 12))
    pygame.draw.rect(screen, BLUE, (sx+40, sy-6, 30, 12))
    # 对接端口
    pygame.draw.circle(screen, WHITE, (sx-35, sy), 6)
    pygame.draw.circle(screen, WHITE, (sx+35, sy), 6)

def draw_satellite(offset_y, offset_x):
    """绘制释放后的卫星"""
    sx = satellite_x + offset_x
    sy = satellite_y + offset_y
    pygame.draw.circle(screen, WHITE, (sx, sy), 8)
    pygame.draw.rect(screen, BLUE, (sx-18, sy-3, 36, 6))

def draw_launch_base(base_screen_y):
    """顶配强化发射基地"""
    center_x = WIDTH // 2
    # 主平台底座
    pygame.draw.rect(screen, DARK_GRAY, (center_x - 100, base_screen_y, 200, 48))
    pygame.draw.rect(screen, GRAY, (center_x - 90, base_screen_y + 6, 180, 36))
    # 着陆靶心
    pygame.draw.circle(screen, LIGHT_GREEN, (center_x, base_screen_y + 24), 32, 4)
    pygame.draw.circle(screen, BASE_YELLOW, (center_x, base_screen_y + 24), 16, 2)
    pygame.draw.circle(screen, RED, (center_x, base_screen_y + 24), 7)
    # 火焰导流槽
    pygame.draw.rect(screen, BLACK, (center_x - 38, base_screen_y + 40, 76, 9))
    # 双侧发射塔架
    tl_x = center_x - 70
    tr_x = center_x + 58
    pygame.draw.rect(screen, DARK_GRAY, (tl_x, base_screen_y - 170, 14, 170))
    pygame.draw.rect(screen, DARK_GRAY, (tr_x, base_screen_y - 170, 14, 170))
    # 横向桁架
    for h in range(0,140,32):
        pygame.draw.rect(screen, GRAY, (tl_x, base_screen_y - h -22, 52,7))
        pygame.draw.rect(screen, GRAY, (tr_x -40, base_screen_y - h -22, 52,7))
    # 塔顶信号灯
    pygame.draw.circle(screen, RED, (tl_x+7, base_screen_y-172),5)
    pygame.draw.circle(screen, GREEN, (tr_x+7, base_screen_y-172),5)
    # 两侧扶梯
    pygame.draw.polygon(screen, BROWN, [(center_x-90,base_screen_y),(center_x-115,base_screen_y+38),(center_x-80,base_screen_y+38)])
    pygame.draw.polygon(screen, BROWN, [(center_x+90,base_screen_y),(center_x+115,base_screen_y+38),(center_x+80,base_screen_y+38)])
    # 护栏
    pygame.draw.rect(screen, BROWN, (center_x-105, base_screen_y-9,210,8))

def draw_rocket(x, y, fire_main, fire_left, fire_right):
    cfg = rocket_types[selected_rocket]
    bw, bh = cfg["width"], cfg["height"]
    boost_mod = cfg["boost"]
    # 加粗箭体
    pygame.draw.rect(screen, cfg["body_color"], (x - bw//2, y, bw, bh))
    # 流线弹头
    pygame.draw.polygon(screen, RED, [(x - bw//2, y), (x + bw//2, y), (x, y - 30)])
    # 加固尾翼
    pygame.draw.polygon(screen, BROWN, [(x - bw//2, y+bh), (x - bw//2-18, y+bh+16), (x - bw//2, y+bh-15)])
    pygame.draw.polygon(screen, BROWN, [(x + bw//2, y+bh), (x + bw//2+18, y+bh+16), (x + bw//2, y+bh-15)])
    # 二级助推
    if not stage2_separated:
        pygame.draw.rect(screen, GRAY, (x - bw//2-8, y+stage2_y_offset,9,38))
        pygame.draw.rect(screen, GRAY, (x + bw//2-1, y+stage2_y_offset,9,38))
    # 主火焰
    if fire_main:
        fh = random.randint(28,60)*boost_mod
        pygame.draw.polygon(screen, ORANGE, [(x-bw//2+2,y+bh),(x+bw//2-2,y+bh),(x,y+bh+fh)])
        pygame.draw.polygon(screen, YELLOW, [(x-bw//4,y+bh),(x+bw//4,y+bh),(x,y+bh+fh*0.7)])
        spawn_launch_smoke(x,y+bh)
    # 姿态喷口
    if fire_left:
        fl = random.randint(14,26)
        pygame.draw.polygon(screen, BLUE, [(x-bw//2,y+bh-12),(x-bw//2-12,y+bh-12),(x-bw//2-6,y+bh-12+fl)])
    if fire_right:
        fr = random.randint(14,26)
        pygame.draw.polygon(screen, BLUE, [(x+bw//2,y+bh-12),(x+bw//2+12,y+bh-12),(x+bw//2+6,y+bh-12+fr)])

running = True
while running:
    # 昼夜背景渐变
    day_time += day_speed
    if day_time > 1:
        day_time = 0
    is_night = day_time > 0.5
    bg_r = int(SKY_DAY[0]*(1-day_time) + SKY_NIGHT[0]*day_time)
    bg_g = int(SKY_DAY[1]*(1-day_time) + SKY_NIGHT[1]*day_time)
    bg_b = int(SKY_DAY[2]*(1-day_time) + SKY_NIGHT[2]*day_time)
    screen.fill((bg_r,bg_g,bg_b))

    clock.tick(FPS)
    keys = pygame.key.get_pressed()
    fire_main = fire_left = fire_right = False

    # 星空闪烁
    for star in stars:
        sx,sy,sz,tw = star
        tw +=1
        if tw>120:tw=0
        alpha = 255 if tw<90 else 60
        if is_night:
            alpha = min(255, int(alpha*1.8))
        pygame.draw.circle(screen,(alpha,alpha,alpha),(sx,sy),sz)
        star[3]=tw

    # 天气刷新
    weather_timer +=1
    if weather_state>0 and weather_timer%3==0:
        spawn_rain()
    # 雨滴渲染
    for ra in weather_particles[:]:
        ra[0] += ra[2]
        ra[1] += ra[3]
        if ra[1]>HEIGHT:
            weather_particles.remove(ra)
        else:
            rain_alpha = 120 if weather_state==1 else 180
            pygame.draw.line(screen,(rain_alpha,rain_alpha,255),(ra[0],ra[1]),(ra[0]+ra[2]*3,ra[1]+ra[3]*3),1)

    # 事件按键
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False
        if event.type == pygame.KEYDOWN:
            # 星球重力
            if event.key == pygame.K_m:
                is_moon_mode = not is_moon_mode
                current_gravity = MOON_GRAVITY if is_moon_mode else EARTH_GRAVITY
            # 切换火箭
            if event.key == pygame.K_1: selected_rocket=0;reset_rocket()
            if event.key == pygame.K_2: selected_rocket=1;reset_rocket()
            if event.key == pygame.K_3: selected_rocket=2;reset_rocket()
            # 二级分离
            if event.key == pygame.K_x and not stage2_separated:
                stage2_separated=True;score+=50
            # 重置
            if event.key == pygame.K_r:
                reset_rocket();explode_particles.clear();weather_particles.clear()
            # 释放卫星（高空才能释放）
            if event.key == pygame.K_s and not satellite_spawned:
                h = max(0, ground_level - rocket_y)
                if h>150:
                    satellite_spawned=True
                    satellite_x = rocket_x
                    satellite_y = rocket_y
                    satellite_vx = speed_x + 0.8
                    satellite_vy = speed_y
                    satellite_orbit = True
                    score +=150
            # 切换天气 Q
            if event.key == pygame.K_q:
                weather_state = (weather_state+1)%3
                weather_particles.clear()
            # 锁定/解除视角锁定 C
            if event.key == pygame.K_c:
                lock_camera = not lock_camera

    cfg = rocket_types[selected_rocket]
    boost_mod = cfg["boost"]

    # 推进控制
    if keys[pygame.K_SPACE] and fuel>0:
        speed_y += thrust_up_base*boost_mod
        fuel -=0.38
        fire_main=True
    if keys[pygame.K_a] and fuel>0:
        speed_x -= thrust_side_base;fuel-=0.22;fire_left=True
    if keys[pygame.K_d] and fuel>0:
        speed_x += thrust_side_base;fuel-=0.22;fire_right=True

    # 物理运算
    speed_y += current_gravity
    speed_x *= (1-air_resistance)
    speed_y = min(speed_y, max_vertical_speed)
    speed_x = max(-max_horizontal_speed, min(speed_x, max_horizontal_speed))
    rocket_x += speed_x
    rocket_y += speed_y
    rocket_x = max(20, min(WIDTH-20, rocket_x))

    # 火箭锁定视角跟随
    if lock_camera:
        target_cam_y = min(320, max(0, ground_level - rocket_y - 160))
        target_cam_x = (WIDTH//2) - rocket_x
        camera_offset_y += (target_cam_y - camera_offset_y)*camera_smooth
        camera_offset_x += (target_cam_x - camera_offset_x)*camera_smooth

    # 落地判定+精准靶心加分
    if rocket_y >= ground_level:
        rocket_y = ground_level
        center_dist = abs(rocket_x - WIDTH//2)
        if abs(speed_y)<2.2:
            bonus = 100
            if center_dist<30:bonus=200
            score += bonus
            fuel = min(max_fuel, fuel+12)
        else:
            spawn_explosion(rocket_x, rocket_y+60)
            score = max(0,score-80)
            fuel=0
        speed_y=0
        speed_x *=0.5

    height = max(0, ground_level - rocket_y)
    score += height*0.003

    # 卫星轨道运动
    if satellite_orbit:
        satellite_x += satellite_vx
        satellite_y += satellite_vy
        satellite_vy += current_gravity*0.3

    # 爆炸粒子
    for p in explode_particles[:]:
        p[0]+=p[2];p[1]+=p[3]+camera_offset_y;p[4]-=1
        if p[4]<=0:explode_particles.remove(p)
        else:pygame.draw.circle(screen,p[5],(int(p[0]+camera_offset_x),int(p[1])),4)

    # 发射烟雾
    for sm in launch_smoke[:]:
        sm[0]+=sm[2];sm[1]+=sm[3]+camera_offset_y;sm[4]-=1
        if sm[4]<=0:launch_smoke.remove(sm)
        else:
            a = int(sm[4]/50*190)
            pygame.draw.circle(screen,(a,a,a),(int(sm[0]+camera_offset_x),int(sm[1])),7)

    # 绘制基地、空间站、卫星、火箭
    plat_y = HEIGHT - 40 + camera_offset_y
    draw_launch_base(plat_y)
    draw_space_station(camera_offset_y, camera_offset_x)
    if satellite_spawned:
        draw_satellite(camera_offset_y, camera_offset_x)
    draw_rocket(rocket_x+camera_offset_x, rocket_y+camera_offset_y, fire_main, fire_left, fire_right)

    # 燃料颜色
    fuel_color = GREEN if fuel>30 else ((255,200,0) if fuel>15 else RED)
    # UI文本
    ui_texts = [
        font.render(f"Fuel:{int(fuel)}%",True,fuel_color),
        font.render(f"Alt:{int(height)}m Score:{int(score)}",True,WHITE),
        font.render(f"Rocket:{cfg['name']} Grav:{'Moon' if is_moon_mode else 'Earth'}",True,YELLOW),
        font.render(f"Time:{'Night' if is_night else 'Day'} Weather:{['Clear','Rain','Storm'][weather_state]}",True,WHITE),
        font.render(f"CamLock:{'ON' if lock_camera else 'OFF'} Sat:{'Launched' if satellite_spawned else 'Ready(S高空释放)'}",True,GREEN)
    ]
    for i,t in enumerate(ui_texts):
        screen.blit(t,(12,12+i*26))

    # 操作提示
    hint = font.render("空格上升 A/D平移 X分离 M重力 Q天气 C视角锁 R重置 1/2/3换火箭",True,WHITE)
    screen.blit(hint,(WIDTH-hint.get_width()-10,HEIGHT-30))

    # 警告提示
    if fuel<=0:
        screen.blit(font.render("FUEL EMPTY GLIDE ONLY",True,CRASH_RED),(WIDTH//2-130,100))
    if stage2_separated:
        screen.blit(font.render("Booster Separated +50",True,GREEN),(WIDTH//2-110,130))

    pygame.display.update()
pygame.quit()