import pygame
import random
import sys
import math

# ========== 初始化 ==========
pygame.init()
WIDTH, HEIGHT = 800, 600
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("🐠 Aquarium Tycoon – Beautiful Edition")
clock = pygame.time.Clock()

# ========== 配色方案（水下氛围） ==========
WATER_TOP = (15, 55, 95)
WATER_BOTTOM = (5, 20, 40)
SAND = (210, 190, 150)
SAND_DARK = (180, 160, 130)
ROCK = (140, 130, 120)
PLANT_GREEN = (100, 200, 100)
PLANT_DARK = (60, 140, 60)
CORAL_PINK = (255, 140, 140)
CORAL_ORANGE = (255, 180, 100)
GOLD = (255, 210, 80)
WHITE = (255, 255, 255)
BLACK = (30, 30, 40)
UI_BG = (15, 35, 55)

# ========== 字体 ==========
font_tiny = pygame.font.Font(None, 18)
font_small = pygame.font.Font(None, 22)
font_med = pygame.font.Font(None, 30)
font_large = pygame.font.Font(None, 44)

# ========== 鱼类数据 ==========
fish_species = [
    {"name": "Guppy",        "icon": "🐟", "color": (255, 150, 50),  "buy": 10,  "sell": 25,  "grow_time": 400, "food": "flakes"},
    {"name": "Neon Tetra",   "icon": "🐠", "color": (50, 200, 255),  "buy": 20,  "sell": 50,  "grow_time": 500, "food": "flakes"},
    {"name": "Angelfish",    "icon": "🐡", "color": (220, 220, 220), "buy": 40,  "sell": 100, "grow_time": 600, "food": "pellets"},
    {"name": "Betta",        "icon": "🦈", "color": (200, 50, 50),   "buy": 60,  "sell": 150, "grow_time": 700, "food": "pellets"},
    {"name": "Discus",       "icon": "🐙", "color": (255, 100, 200), "buy": 100, "sell": 250, "grow_time": 900, "food": "worms"},
]

# ========== 装饰品 ==========
decorations = [
    {"name": "Seaweed",    "cost": 20,  "type": "plant"},
    {"name": "Castle",     "cost": 50,  "type": "structure"},
    {"name": "Treasure",   "cost": 40,  "type": "structure"},
    {"name": "Coral",      "cost": 30,  "type": "plant"},
    {"name": "Bubbler",    "cost": 35,  "type": "device"},
]

# ========== 水箱升级 ==========
tank_upgrades = [
    {"name": "Small Tank",  "capacity": 4, "cost": 0,   "owned": True},
    {"name": "Medium Tank", "capacity": 6, "cost": 100, "owned": False},
    {"name": "Large Tank",  "capacity": 10,"cost": 250, "owned": False},
]

# ========== 特效类 ==========
class Bubble:
    def __init__(self, x, y):
        self.x = x
        self.y = y
        self.radius = random.randint(2, 5)
        self.speed = random.uniform(0.5, 1.5)
        self.wobble = random.uniform(0, 2*math.pi)
    def update(self):
        self.y -= self.speed
        self.x += math.sin(self.wobble) * 0.2
        self.wobble += 0.05
    def draw(self, surface):
        alpha = 180
        pygame.draw.circle(surface, (180, 220, 255, alpha), (int(self.x), int(self.y)), self.radius)

class FloatText:
    def __init__(self, x, y, text, color=GOLD):
        self.x = x
        self.y = y
        self.text = text
        self.color = color
        self.life = 40
    def update(self):
        self.y -= 1
        self.life -= 1
    def draw(self, surface):
        alpha = min(255, self.life * 8)
        surf = font_small.render(self.text, True, self.color)
        surf.set_alpha(alpha)
        surface.blit(surf, (self.x - surf.get_width()//2, self.y))

class Spark:
    def __init__(self, x, y):
        self.x = x
        self.y = y
        self.vx = random.uniform(-1, 1)
        self.vy = random.uniform(-2, -0.5)
        self.life = 20
    def update(self):
        self.x += self.vx
        self.y += self.vy
        self.life -= 1
    def draw(self, surface):
        if self.life > 0:
            alpha = min(255, self.life * 15)
            pygame.draw.circle(surface, (255, 255, 200, alpha), (int(self.x), int(self.y)), 2)

# ========== 鱼（美化版） ==========
class Fish:
    def __init__(self, x, y, species_data):
        self.x = x
        self.y = y
        self.species = species_data
        self.vx = random.uniform(-0.8, 0.8)
        self.vy = random.uniform(-0.4, 0.4)
        self.growth = 0
        self.hunger = 100
        self.alive = True
        self.sick = False
        self.wiggle = random.uniform(0, 2*math.pi)
        self.scale = 0.8 + 0.2 * (self.growth / max(1, self.species["grow_time"]))

    def update(self, tank_rect):
        if not self.alive: return
        self.x += self.vx
        self.y += self.vy
        self.wiggle += 0.1

        if self.x < tank_rect.left + 20 or self.x > tank_rect.right - 20:
            self.vx *= -1
        if self.y < tank_rect.top + 20 or self.y > tank_rect.bottom - 25:
            self.vy *= -1

        if random.random() < 0.02:
            self.vx += random.uniform(-0.3, 0.3)
            self.vy += random.uniform(-0.2, 0.2)
            self.vx = max(-1.5, min(1.5, self.vx))
            self.vy = max(-0.8, min(0.8, self.vy))

        self.hunger -= 0.05
        if self.hunger <= 0:
            self.alive = False
        elif self.hunger < 30:
            self.sick = True
        else:
            self.sick = False

        if self.hunger > 30 and self.growth < self.species["grow_time"]:
            self.growth += 0.5
        self.scale = 0.8 + 0.3 * (self.growth / max(1, self.species["grow_time"]))

    def feed(self, food_type):
        if not self.alive: return False
        if food_type == self.species["food"]:
            self.hunger = min(100, self.hunger + 60)
            self.sick = False
            return True
        else:
            self.hunger = min(100, self.hunger + 20)
            return False

    def is_fully_grown(self):
        return self.growth >= self.species["grow_time"]

    def draw(self, surface):
        if not self.alive: return
        color = self.species["color"]
        if self.sick:
            color = tuple(max(20, c-70) for c in color)

        body_len = int(20 * self.scale)
        body_h = int(10 * self.scale)

        offset = math.sin(self.wiggle * 4) * 2
        body_rect = pygame.Rect(self.x - body_len//2, self.y - body_h//2, body_len, body_h)
        pygame.draw.ellipse(surface, color, body_rect)
        light_color = tuple(min(255, c+50) for c in color)
        light_rect = pygame.Rect(self.x - body_len//2 + 2, self.y - body_h//2, body_len//2, body_h//2)
        pygame.draw.ellipse(surface, light_color, light_rect, 2)

        tail_dir = -1 if self.vx > 0 else 1
        tail_x = self.x + (body_len//2) * tail_dir
        tail_points = [
            (tail_x, self.y),
            (tail_x + 8 * tail_dir, self.y - 6 + offset),
            (tail_x + 8 * tail_dir, self.y + 6 - offset)
        ]
        pygame.draw.polygon(surface, color, tail_points)

        eye_dir = 1 if self.vx > 0 else -1
        eye_x = self.x + (body_len//4) * eye_dir
        eye_y = self.y - 2
        pygame.draw.circle(surface, WHITE, (int(eye_x), eye_y), 4)
        pygame.draw.circle(surface, BLACK, (int(eye_x), eye_y), 2)

        fin_y = self.y - body_h//2
        fin_points = [
            (self.x - 4, fin_y),
            (self.x, fin_y - 7 + offset),
            (self.x + 4, fin_y)
        ]
        pygame.draw.polygon(surface, tuple(min(255, c+40) for c in color), fin_points)

# ========== 主游戏 ==========
class AquariumGame:
    def __init__(self):
        self.state = "menu"
        self.gold = 100
        self.day = 1
        self.tick = 0
        self.tank_rect = pygame.Rect(70, 90, 660, 360)
        self.current_tank = tank_upgrades[0]
        self.fishes = []
        self.decor_placed = []
        self.bubbles = []
        self.float_texts = []
        self.sparks = []

        self.food = {"flakes": 5, "pellets": 3, "worms": 2}
        self.selected_food = None

        self.collected_species = set()

        self.btn_play = pygame.Rect(WIDTH//2-60, 350, 120, 50)
        self.btn_shop = pygame.Rect(20, 40, 80, 40)
        self.btn_food = pygame.Rect(620, 40, 100, 40)
        self.btn_sell = pygame.Rect(620, 90, 100, 40)
        self.btn_album = pygame.Rect(500, 40, 100, 40)

        self.msg = ""
        self.msg_timer = 0

        self.add_fish(fish_species[0])
        self.add_fish(fish_species[0])

        for _ in range(20):
            self.bubbles.append(Bubble(random.randint(self.tank_rect.left+10, self.tank_rect.right-10),
                                      random.randint(self.tank_rect.top+10, self.tank_rect.bottom-10)))

    def add_fish(self, species):
        if len(self.fishes) < self.current_tank["capacity"]:
            x = random.randint(self.tank_rect.left+40, self.tank_rect.right-40)
            y = random.randint(self.tank_rect.top+40, self.tank_rect.bottom-40)
            self.fishes.append(Fish(x, y, species))
            self.collected_species.add(species["name"])
            self.float_texts.append(FloatText(x, y-10, f"+{species['name']}", GOLD))
            return True
        return False

    def handle_click(self, pos):
        if self.state == "menu":
            if self.btn_play.collidepoint(pos):
                self.state = "playing"
            return
        if self.state == "shop":
            back_btn = pygame.Rect(20, 40, 80, 40)
            if back_btn.collidepoint(pos):
                self.state = "playing"
                return
            self.buy_shop_item(pos)
            return
        if self.state == "album":
            back_btn = pygame.Rect(20, 40, 80, 40)
            if back_btn.collidepoint(pos):
                self.state = "playing"
                return
            return

        if self.btn_shop.collidepoint(pos):
            self.state = "shop"
            return
        if self.btn_album.collidepoint(pos):
            self.state = "album"
            return

        if self.btn_food.collidepoint(pos):
            food_types = list(self.food.keys())
            if self.selected_food in food_types:
                idx = food_types.index(self.selected_food)
                self.selected_food = food_types[(idx+1)%len(food_types)]
            else:
                self.selected_food = food_types[0]
            return

        if self.tank_rect.collidepoint(pos) and self.selected_food and self.food[self.selected_food] > 0:
            self.food[self.selected_food] -= 1
            fed = any(fish.alive and fish.feed(self.selected_food) for fish in self.fishes)
            self.msg = f"Fed {self.selected_food}!" if fed else "No fish to feed..."
            self.msg_timer = 40
            for _ in range(10):
                self.sparks.append(Spark(pos[0], pos[1]))
            return

        for fish in self.fishes:
            if fish.alive and fish.is_fully_grown():
                if math.hypot(fish.x - pos[0], fish.y - pos[1]) < 20:
                    self.gold += fish.species["sell"]
                    self.float_texts.append(FloatText(fish.x, fish.y-10, f"+{fish.species['sell']}g", GOLD))
                    self.fishes.remove(fish)
                    self.msg = f"Sold {fish.species['name']}!"
                    self.msg_timer = 50
                    return

    def buy_shop_item(self, pos):
        y = 120
        for sp in fish_species:
            rect = pygame.Rect(80, y, 300, 40)
            if rect.collidepoint(pos) and self.gold >= sp["buy"] and len(self.fishes) < self.current_tank["capacity"]:
                if self.add_fish(sp):
                    self.gold -= sp["buy"]
                    self.msg = f"Bought {sp['name']}!"
                    self.msg_timer = 40
            y += 45

        y += 15
        for dec in decorations:
            rect = pygame.Rect(80, y, 300, 40)
            if rect.collidepoint(pos) and self.gold >= dec["cost"]:
                self.gold -= dec["cost"]
                dx = random.randint(self.tank_rect.left+30, self.tank_rect.right-30)
                dy = random.randint(self.tank_rect.top+30, self.tank_rect.bottom-30)
                self.decor_placed.append((dec, dx, dy))
                self.msg = f"Placed {dec['name']}!"
                self.msg_timer = 40
            y += 45

        y += 15
        for upg in tank_upgrades:
            if not upg["owned"]:
                rect = pygame.Rect(80, y, 300, 40)
                if rect.collidepoint(pos) and self.gold >= upg["cost"]:
                    self.gold -= upg["cost"]
                    upg["owned"] = True
                    self.current_tank = upg
                    self.msg = f"Upgraded to {upg['name']}!"
                    self.msg_timer = 60
                y += 45

    def update(self):
        if self.state != "playing": return
        self.tick += 1

        for fish in self.fishes[:]:
            fish.update(self.tank_rect)
            if not fish.alive:
                self.fishes.remove(fish)
                self.msg = "A fish died... 😢"
                self.msg_timer = 50

        for bub in self.bubbles[:]:
            bub.update()
            if bub.y < self.tank_rect.top or bub.x < self.tank_rect.left or bub.x > self.tank_rect.right:
                self.bubbles.remove(bub)
        for ft in self.float_texts[:]:
            ft.update()
            if ft.life <= 0: self.float_texts.remove(ft)
        for sp in self.sparks[:]:
            sp.update()
            if sp.life <= 0: self.sparks.remove(sp)

        if len(self.bubbles) < 25 and random.random() < 0.2:
            self.bubbles.append(Bubble(random.randint(self.tank_rect.left+5, self.tank_rect.right-5),
                                      self.tank_rect.bottom-10))

        if self.tick % 1800 == 0:
            self.day += 1
            for fish in self.fishes:
                fish.hunger -= 10
            self.msg = f"☀️ Day {self.day}"
            self.msg_timer = 50

        if self.msg_timer > 0:
            self.msg_timer -= 1
        else:
            self.msg = ""

    def draw_tank(self):
        # ★ 清除全屏背景，彻底解决拖尾 ★
        screen.fill((10, 30, 55))

        # 水渐变
        for y in range(self.tank_rect.top, self.tank_rect.bottom, 2):
            ratio = (y - self.tank_rect.top) / self.tank_rect.height
            r = int(15 + 10*ratio)
            g = int(55 + 20*ratio)
            b = int(95 + 30*ratio)
            pygame.draw.line(screen, (r, g, b), (self.tank_rect.left, y), (self.tank_rect.right, y))

        # 光线效果
        for i in range(3):
            lx = self.tank_rect.left + 40 + i*200 + math.sin(self.tick*0.02 + i)*20
            for j in range(5):
                alpha = 10 + j*2
                pygame.draw.line(screen, (255, 255, 200, alpha), (lx, self.tank_rect.top), (lx, self.tank_rect.bottom), 3)

        # 沙地
        for x in range(self.tank_rect.left, self.tank_rect.right, 10):
            h = random.randint(12, 18)
            pygame.draw.rect(screen, SAND if random.random()>0.3 else SAND_DARK,
                            (x, self.tank_rect.bottom-h, 10, h))

        # 装饰品
        for dec, dx, dy in self.decor_placed:
            if dec["type"] == "plant":
                for i in range(3):
                    sway = math.sin(self.tick*0.03 + i + dx) * 3
                    pygame.draw.line(screen, PLANT_GREEN, (dx+sway, dy), (dx+sway, dy-30), 5)
                    pygame.draw.line(screen, PLANT_DARK, (dx+sway, dy), (dx+sway, dy-30), 2)
            elif dec["type"] == "structure":
                if dec["name"] == "Castle":
                    pygame.draw.rect(screen, ROCK, (dx-12, dy-18, 24, 24), border_radius=4)
                    pygame.draw.rect(screen, (80,80,80), (dx-12, dy-18, 24, 24), 3, border_radius=4)
                elif dec["name"] == "Treasure":
                    pygame.draw.rect(screen, GOLD, (dx-10, dy-12, 20, 16), border_radius=3)
                    pygame.draw.rect(screen, BLACK, (dx-10, dy-12, 20, 16), 2, border_radius=3)
            elif dec["type"] == "device":
                if random.random() < 0.1:
                    self.bubbles.append(Bubble(dx, dy-8))

        # 气泡
        for bub in self.bubbles:
            bub.draw(screen)

        # 鱼
        for fish in self.fishes:
            fish.draw(screen)

        # 特效粒子
        for sp in self.sparks:
            sp.draw(screen)
        for ft in self.float_texts:
            ft.draw(screen)

        # 水箱边框
        pygame.draw.rect(screen, (60, 80, 100), self.tank_rect, 6, border_radius=12)
        pygame.draw.rect(screen, (90, 120, 150), self.tank_rect, 2, border_radius=12)

    def draw_ui(self):
        panel = pygame.Surface((WIDTH, 70), pygame.SRCALPHA)
        panel.fill((10, 25, 45, 220))
        screen.blit(panel, (0, 0))
        gold_t = font_med.render(f"Gold: {self.gold}", True, GOLD)
        screen.blit(gold_t, (WIDTH-150, 45))
        day_t = font_small.render(f"Day {self.day}", True, WHITE)
        screen.blit(day_t, (WIDTH//2-30, 45))

        def draw_button(rect, text, color):
            pygame.draw.rect(screen, color, rect, border_radius=8)
            pygame.draw.rect(screen, WHITE, rect, 2, border_radius=8)
            txt = font_small.render(text, True, WHITE)
            screen.blit(txt, (rect.x+8, rect.y+8))

        draw_button(self.btn_shop, "Shop", (40, 70, 110))
        draw_button(self.btn_album, "Album", (40, 70, 110))
        food_color = (100, 70, 30) if self.selected_food else (40, 60, 90)
        draw_button(self.btn_food, self.selected_food or "Food", food_color)
        if self.selected_food:
            cnt = self.food.get(self.selected_food, 0)
            cnt_t = font_tiny.render(f"x{cnt}", True, WHITE)
            screen.blit(cnt_t, (self.btn_food.x+78, self.btn_food.y+24))

        if self.msg:
            msg_surf = font_med.render(self.msg, True, WHITE)
            bg_rect = msg_surf.get_rect(center=(WIDTH//2, HEIGHT-30))
            pygame.draw.rect(screen, (0,0,0,160), bg_rect.inflate(20,8), border_radius=10)
            screen.blit(msg_surf, bg_rect)

    def draw_menu(self):
        screen.fill((10, 30, 55))
        title = font_large.render("🐠 Aquarium Tycoon", True, WHITE)
        screen.blit(title, (WIDTH//2 - title.get_width()//2, 150))
        pygame.draw.rect(screen, (50, 100, 160), self.btn_play, border_radius=15)
        pygame.draw.rect(screen, WHITE, self.btn_play, 2, border_radius=15)
        play_t = font_med.render("Play", True, WHITE)
        screen.blit(play_t, (WIDTH//2-25, 362))
        hint = font_small.render("Buy fish, feed them, grow, sell, upgrade!", True, WHITE)
        screen.blit(hint, (WIDTH//2 - hint.get_width()//2, 450))

    def draw_shop(self):
        screen.fill((15, 35, 55))
        back_btn = pygame.Rect(20, 40, 80, 40)
        pygame.draw.rect(screen, (40, 70, 110), back_btn, border_radius=8)
        pygame.draw.rect(screen, WHITE, back_btn, 2, border_radius=8)
        screen.blit(font_small.render("Back", True, WHITE), (back_btn.x+15, back_btn.y+8))

        y = 120
        screen.blit(font_med.render("Fish Market", True, WHITE), (50, 90))
        for sp in fish_species:
            rect = pygame.Rect(80, y, 300, 40)
            color = (30, 90, 130) if self.gold >= sp["buy"] else (50,50,70)
            pygame.draw.rect(screen, color, rect, border_radius=6)
            pygame.draw.rect(screen, WHITE, rect, 2, border_radius=6)
            txt = f"{sp['icon']} {sp['name']} - Buy {sp['buy']}g | Sell {sp['sell']}g"
            screen.blit(font_small.render(txt, True, WHITE), (85, y+8))
            y += 45

        y += 15
        screen.blit(font_med.render("Decorations", True, WHITE), (50, y-15))
        for dec in decorations:
            rect = pygame.Rect(80, y, 300, 40)
            color = (30, 90, 130) if self.gold >= dec["cost"] else (50,50,70)
            pygame.draw.rect(screen, color, rect, border_radius=6)
            pygame.draw.rect(screen, WHITE, rect, 2, border_radius=6)
            screen.blit(font_small.render(f"{dec['name']} ({dec['cost']}g)", True, WHITE), (85, y+8))
            y += 45

        y += 15
        screen.blit(font_med.render("Tank Upgrades", True, WHITE), (50, y-15))
        for upg in tank_upgrades:
            if not upg["owned"]:
                rect = pygame.Rect(80, y, 300, 40)
                color = (30, 90, 130) if self.gold >= upg["cost"] else (50,50,70)
                pygame.draw.rect(screen, color, rect, border_radius=6)
                pygame.draw.rect(screen, WHITE, rect, 2, border_radius=6)
                screen.blit(font_small.render(f"{upg['name']} ({upg['cost']}g, cap {upg['capacity']})", True, WHITE), (85, y+8))
                y += 45

    def draw_album(self):
        screen.fill((15, 35, 55))
        back_btn = pygame.Rect(20, 40, 80, 40)
        pygame.draw.rect(screen, (40, 70, 110), back_btn, border_radius=8)
        pygame.draw.rect(screen, WHITE, back_btn, 2, border_radius=8)
        screen.blit(font_small.render("Back", True, WHITE), (back_btn.x+15, back_btn.y+8))

        screen.blit(font_large.render("Collection", True, WHITE), (WIDTH//2-80, 30))
        y = 120
        for sp in fish_species:
            owned = sp["name"] in self.collected_species
            color = sp["color"] if owned else (100,100,100)
            pygame.draw.circle(screen, color, (100, y+15), 12)
            pygame.draw.circle(screen, WHITE, (100, y+15), 12, 2)
            txt = sp["name"] if owned else "???"
            screen.blit(font_med.render(txt, True, WHITE), (130, y+4))
            y += 50

    def draw(self):
        if self.state == "menu":
            self.draw_menu()
        elif self.state == "playing":
            self.draw_tank()
            self.draw_ui()
        elif self.state == "shop":
            self.draw_shop()
        elif self.state == "album":
            self.draw_album()

    def run(self):
        running = True
        while running:
            for event in pygame.event.get():
                if event.type == pygame.QUIT:
                    running = False
                elif event.type == pygame.MOUSEBUTTONDOWN and event.button == 1:
                    self.handle_click(event.pos)

            self.update()
            self.draw()
            pygame.display.flip()
            clock.tick(60)

        pygame.quit()
        sys.exit()

if __name__ == "__main__":
    game = AquariumGame()
    game.run()