import pygame
import sys
import datetime
import json
import os
from pygame.locals import *

# 初始化pygame
pygame.init()

# 颜色定义
WHITE = (255, 255, 255)
BLACK = (0, 0, 0)
RED = (255, 0, 0)
GREEN = (0, 255, 0)
BLUE = (0, 120, 255)
GRAY = (200, 200, 200)
LIGHT_GRAY = (240, 240, 240)
DARK_GRAY = (100, 100, 100)
YELLOW = (255, 255, 0)
PURPLE = (128, 0, 128)

# 窗口设置
WIDTH, HEIGHT = 1000, 700
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("日历提醒与计划表")

# 字体
font = pygame.font.SysFont(None, 24)
title_font = pygame.font.SysFont(None, 36)
small_font = pygame.font.SysFont(None, 20)

# 数据存储
DATA_FILE = "calendar_data.json"

def load_data():
    """加载保存的数据"""
    if os.path.exists(DATA_FILE):
        with open(DATA_FILE, 'r', encoding='utf-8') as f:
            return json.load(f)
    return {}

def save_data(data):
    """保存数据到文件"""
    with open(DATA_FILE, 'w', encoding='utf-8') as f:
        json.dump(data, f, ensure_ascii=False, indent=2)

class Calendar:
    def __init__(self):
        self.current_date = datetime.date.today()
        self.selected_date = datetime.date.today()
        self.plans = load_data()
        self.input_text = ""
        self.editing_plan_id = None
        self.input_active = False
        
    def get_month_days(self, year, month):
        """获取某年某月的所有日期"""
        first_day = datetime.date(year, month, 1)
        last_day = datetime.date(year, month + 1, 1) - datetime.timedelta(days=1) if month < 12 else datetime.date(year + 1, 1, 1) - datetime.timedelta(days=1)
        
        # 计算第一天是周几（0-6，0代表周一）
        first_weekday = first_day.weekday()
        
        days = []
        # 添加前面的空白
        for _ in range(first_weekday):
            days.append(None)
        
        # 添加当月的日期
        for day in range(1, last_day.day + 1):
            days.append(datetime.date(year, month, day))
        
        return days
    
    def draw_calendar(self, screen):
        """绘制日历"""
        year = self.current_date.year
        month = self.current_date.month
        
        # 绘制标题
        month_names = ["一月", "二月", "三月", "四月", "五月", "六月", 
                      "七月", "八月", "九月", "十月", "十一月", "十二月"]
        title = f"{year}年 {month_names[month-1]}"
        title_surface = title_font.render(title, True, BLUE)
        screen.blit(title_surface, (400, 20))
        
        # 绘制星期标题
        weekdays = ["一", "二", "三", "四", "五", "六", "日"]
        for i, day in enumerate(weekdays):
            x = 150 + i * 100
            day_surface = font.render(day, True, DARK_GRAY)
            screen.blit(day_surface, (x + 35, 80))
        
        # 获取当月所有日期
        days = self.get_month_days(year, month)
        
        # 绘制日期方格
        for i, day in enumerate(days):
            row = i // 7
            col = i % 7
            
            x = 150 + col * 100
            y = 120 + row * 80
            
            if day:
                # 绘制日期方格
                color = BLUE if day == self.selected_date else LIGHT_GRAY
                if day == datetime.date.today():
                    color = YELLOW
                
                pygame.draw.rect(screen, color, (x, y, 80, 60))
                pygame.draw.rect(screen, BLACK, (x, y, 80, 60), 1)
                
                # 绘制日期数字
                day_surface = font.render(str(day.day), True, BLACK)
                screen.blit(day_surface, (x + 5, y + 5))
                
                # 如果有计划，显示小红点
                date_str = day.strftime("%Y-%m-%d")
                if date_str in self.plans and self.plans[date_str]:
                    pygame.draw.circle(screen, RED, (x + 70, y + 10), 5)
                
                # 如果是选中的日期，绘制选中边框
                if day == self.selected_date:
                    pygame.draw.rect(screen, RED, (x, y, 80, 60), 3)
    
    def draw_plans(self, screen):
        """绘制计划列表"""
        # 绘制计划区域标题
        date_str = self.selected_date.strftime("%Y年%m月%d日")
        title = f"{date_str} 的计划"
        title_surface = font.render(title, True, BLUE)
        screen.blit(title_surface, (750, 20))
        
        # 绘制计划列表区域
        pygame.draw.rect(screen, LIGHT_GRAY, (750, 50, 220, 300))
        pygame.draw.rect(screen, BLACK, (750, 50, 220, 300), 1)
        
        # 获取当天的计划
        date_str = self.selected_date.strftime("%Y-%m-%d")
        day_plans = self.plans.get(date_str, [])
        
        # 显示计划
        for i, plan in enumerate(day_plans):
            y = 60 + i * 30
            if y < 340:  # 只显示可见区域内的计划
                # 计划文本
                plan_text = f"{i+1}. {plan}"
                if len(plan_text) > 20:
                    plan_text = plan_text[:20] + "..."
                plan_surface = small_font.render(plan_text, True, BLACK)
                screen.blit(plan_surface, (760, y))
                
                # 删除按钮
                pygame.draw.rect(screen, RED, (950, y, 15, 15))
                del_surface = small_font.render("×", True, WHITE)
                screen.blit(del_surface, (953, y))
        
        # 如果没有计划
        if not day_plans:
            no_plan = small_font.render("暂无计划", True, DARK_GRAY)
            screen.blit(no_plan, (810, 150))
    
    def draw_input_area(self, screen):
        """绘制输入区域"""
        # 输入框
        input_rect = pygame.Rect(750, 400, 220, 30)
        color = BLUE if self.input_active else BLACK
        pygame.draw.rect(screen, WHITE, input_rect)
        pygame.draw.rect(screen, color, input_rect, 2)
        
        # 输入文本
        text_surface = small_font.render(self.input_text, True, BLACK)
        screen.blit(text_surface, (input_rect.x + 5, input_rect.y + 5))
        
        # 添加按钮
        add_btn = pygame.Rect(750, 450, 100, 40)
        pygame.draw.rect(screen, GREEN, add_btn)
        add_text = font.render("添加计划", True, WHITE)
        screen.blit(add_text, (add_btn.x + 10, add_btn.y + 10))
        
        # 清空按钮
        clear_btn = pygame.Rect(870, 450, 100, 40)
        pygame.draw.rect(screen, RED, clear_btn)
        clear_text = font.render("清空", True, WHITE)
        screen.blit(clear_text, (clear_btn.x + 25, clear_btn.y + 10))
        
        # 切换月份按钮
        prev_btn = pygame.Rect(300, 20, 40, 40)
        next_btn = pygame.Rect(700, 20, 40, 40)
        
        pygame.draw.rect(screen, BLUE, prev_btn)
        pygame.draw.rect(screen, BLUE, next_btn)
        
        prev_text = font.render("<", True, WHITE)
        next_text = font.render(">", True, WHITE)
        screen.blit(prev_text, (prev_btn.x + 15, prev_btn.y + 10))
        screen.blit(next_text, (next_btn.x + 15, next_btn.y + 10))
        
        # 今天按钮
        today_btn = pygame.Rect(450, 550, 100, 40)
        pygame.draw.rect(screen, PURPLE, today_btn)
        today_text = font.render("返回今天", True, WHITE)
        screen.blit(today_text, (today_btn.x + 10, today_btn.y + 10))
    
    def draw(self, screen):
        """绘制整个界面"""
        screen.fill(WHITE)
        
        # 绘制分隔线
        pygame.draw.line(screen, GRAY, (730, 0), (730, HEIGHT), 2)
        pygame.draw.line(screen, GRAY, (0, 500), (WIDTH, 500), 2)
        
        # 绘制各个部分
        self.draw_calendar(screen)
        self.draw_plans(screen)
        self.draw_input_area(screen)
        
        # 绘制说明文字
        instructions = [
            "操作说明:",
            "1. 点击日历中的日期选择日期",
            "2. 在右侧输入计划，点击添加",
            "3. 点击×按钮删除计划",
            "4. 使用< >按钮切换月份",
            "5. 带红点的日期表示有计划"
        ]
        
        for i, text in enumerate(instructions):
            text_surface = small_font.render(text, True, DARK_GRAY)
            screen.blit(text_surface, (20, 520 + i * 25))
    
    def handle_click(self, pos):
        """处理鼠标点击"""
        x, y = pos
        
        # 检查是否点击了日历日期
        if 150 <= x <= 830 and 120 <= y <= 480:
            col = (x - 150) // 100
            row = (y - 120) // 80
            idx = row * 7 + col
            
            year = self.current_date.year
            month = self.current_date.month
            days = self.get_month_days(year, month)
            
            if idx < len(days) and days[idx]:
                self.selected_date = days[idx]
        
        # 检查是否点击了月份切换按钮
        elif 300 <= x <= 340 and 20 <= y <= 60:  # 上个月
            self.current_date = self.current_date.replace(day=1)
            if self.current_date.month == 1:
                self.current_date = self.current_date.replace(year=self.current_date.year-1, month=12)
            else:
                self.current_date = self.current_date.replace(month=self.current_date.month-1)
        
        elif 700 <= x <= 740 and 20 <= y <= 60:  # 下个月
            self.current_date = self.current_date.replace(day=1)
            if self.current_date.month == 12:
                self.current_date = self.current_date.replace(year=self.current_date.year+1, month=1)
            else:
                self.current_date = self.current_date.replace(month=self.current_date.month+1)
        
        # 检查是否点击了今天按钮
        elif 450 <= x <= 550 and 550 <= y <= 590:
            self.current_date = datetime.date.today()
            self.selected_date = datetime.date.today()
        
        # 检查是否点击了输入框
        elif 750 <= x <= 970 and 400 <= y <= 430:
            self.input_active = True
        
        # 检查是否点击了添加按钮
        elif 750 <= x <= 850 and 450 <= y <= 490:
            if self.input_text.strip():
                date_str = self.selected_date.strftime("%Y-%m-%d")
                if date_str not in self.plans:
                    self.plans[date_str] = []
                self.plans[date_str].append(self.input_text)
                self.input_text = ""
                save_data(self.plans)
        
        # 检查是否点击了清空按钮
        elif 870 <= x <= 970 and 450 <= y <= 490:
            self.input_text = ""
            self.input_active = False
        
        # 检查是否点击了删除计划按钮
        elif 750 <= x <= 970 and 50 <= y <= 350:
            date_str = self.selected_date.strftime("%Y-%m-%d")
            if date_str in self.plans:
                plan_index = (y - 60) // 30
                if 0 <= plan_index < len(self.plans[date_str]):
                    if 950 <= x <= 965:  # 点击了删除按钮
                        del self.plans[date_str][plan_index]
                        if not self.plans[date_str]:
                            del self.plans[date_str]
                        save_data(self.plans)
        
        else:
            self.input_active = False
    
    def handle_keypress(self, event):
        """处理键盘输入"""
        if self.input_active:
            if event.key == pygame.K_BACKSPACE:
                self.input_text = self.input_text[:-1]
            elif event.key == pygame.K_RETURN:
                if self.input_text.strip():
                    date_str = self.selected_date.strftime("%Y-%m-%d")
                    if date_str not in self.plans:
                        self.plans[date_str] = []
                    self.plans[date_str].append(self.input_text)
                    self.input_text = ""
                    save_data(self.plans)
            else:
                if len(self.input_text) < 50:  # 限制输入长度
                    self.input_text += event.unicode

def main():
    calendar = Calendar()
    clock = pygame.time.Clock()
    
    running = True
    while running:
        for event in pygame.event.get():
            if event.type == QUIT:
                running = False
            
            elif event.type == MOUSEBUTTONDOWN:
                if event.button == 1:  # 左键点击
                    calendar.handle_click(event.pos)
            
            elif event.type == KEYDOWN:
                calendar.handle_keypress(event)
        
        # 绘制界面
        calendar.draw(screen)
        pygame.display.flip()
        clock.tick(60)
    
    # 退出前保存数据
    save_data(calendar.plans)
    pygame.quit()
    sys.exit()

if __name__ == "__main__":
    main()