import pandas as pd
import streamlit as st
from datetime import datetime, timedelta
import matplotlib.pyplot as plt
import io
import base64
from matplotlib.patches import Rectangle
import matplotlib.gridspec as gridspec
import math

# [前面的常量定义和 process_schedule 函数保持不变]
SPLIT_TIME = "17:30"
BUSINESS_START = "09:30"
BUSINESS_END = "01:30"
BORDER_COLOR = '#A9A9A9'

def process_schedule(file):
    """处理上传的 Excel 文件，生成排序和分组后的打印内容"""
    try:
        df = pd.read_excel(file, skiprows=8)
        df = df.iloc[:, [6, 7, 9]]
        df.columns = ['Hall', 'StartTime', 'EndTime']
        df = df.dropna(subset=['Hall', 'StartTime', 'EndTime'])
        df['Hall'] = df['Hall'].str.extract(r'(\d+)号').astype(str) + '#'
        df['StartTime'] = pd.to_datetime(df['StartTime']).dt.strftime('%H:%M')
        df['EndTime'] = pd.to_datetime(df['EndTime']).dt.strftime('%H:%M')
        
        base_date = datetime.today().date()
        df['StartTime'] = pd.to_datetime(base_date.strftime('%Y-%m-%d ') + df['StartTime'])
        df['EndTime'] = pd.to_datetime(base_date.strftime('%Y-%m-%d ') + df['EndTime'])
        
        df.loc[df['EndTime'] < df['StartTime'], 'EndTime'] += timedelta(days=1)
        
        business_start = datetime.strptime(base_date.strftime('%Y-%m-%d ') + BUSINESS_START, '%Y-%m-%d %H:%M')
        business_end = datetime.strptime(base_date.strftime('%Y-%m-%d ') + BUSINESS_END, '%Y-%m-%d %H:%M')
        if business_end < business_start:
            business_end += timedelta(days=1)
            
        mask = (df['StartTime'] >= business_start) | (df['StartTime'] <= business_end)
        df = df[mask]
        
        df = df.sort_values('EndTime')
        
        split_time = datetime.strptime(base_date.strftime('%Y-%m-%d ') + SPLIT_TIME, '%Y-%m-%d %H:%M')
        part1 = df[df['EndTime'] <= split_time].copy()
        part2 = df[df['EndTime'] > split_time].copy()
        
        for part in [part1, part2]:
            part['EndTime'] = part['EndTime'].dt.strftime('%-I:%M')
            
        return part1[['Hall', 'EndTime']], part2[['Hall', 'EndTime']]
    
    except Exception as e:
        st.error(f"处理文件时出错: {str(e)}")
        return None, None

def create_print_layout(data, title):
    """创建打印布局"""
    if data.empty:
        return None
        
    # 设置 A5 纸张竖向尺寸
    fig = plt.figure(figsize=(5.83, 8.27), dpi=300)
    plt.subplots_adjust(left=0.05, right=0.95, top=0.95, bottom=0.05)
    
    # 设置字体
    plt.rcParams['font.family'] = 'sans-serif'
    plt.rcParams['font.sans-serif'] = ['SimHei']
    
    # 计算行数和总数
    total_items = len(data)
    num_cols = 3
    num_rows = math.ceil(total_items / num_cols)
    
    # 创建网格
    gs = gridspec.GridSpec(num_rows, num_cols, hspace=0.1, wspace=0.1)
    
    # 计算字体大小
    base_fontsize = min(50, 280 / num_rows)
    
    # 重要改动：正确的竖向排序逻辑
    data_values = data.values.tolist()
    
    # 确保数据长度是3的倍数
    while len(data_values) % 3 != 0:
        data_values.append(['', ''])
        
    # 计算每列应该包含的行数
    rows_per_col = math.ceil(len(data_values) / 3)
    
    # 创建一个新的数据列表，用于存储重新排序后的数据
    sorted_data = [['', '']] * len(data_values)
    
    # 正确的竖向排序逻辑
    for i, item in enumerate(data_values):
        if item[0] and item[1]:  # 只处理非空数据
            # 计算新的位置
            row = i % rows_per_col
            col = i // rows_per_col
            new_index = row * 3 + col
            if new_index < len(sorted_data):
                sorted_data[new_index] = item
    
    # 填充数据
    for idx, (hall, end_time) in enumerate(sorted_data):
        if hall and end_time:  # 只处理非空数据
            row = idx // 3
            col = idx % 3
            
            ax = plt.subplot(gs[row, col])
            
            # 设置边框
            for spine in ax.spines.values():
                spine.set_color(BORDER_COLOR)
                spine.set_linewidth(0.5)
            
            # 显示文本
            display_text = f"{hall}{end_time}"
            ax.text(0.5, 0.5, display_text,
                    fontsize=base_fontsize,
                    fontweight='bold',
                    ha='center',
                    va='center')
            
            # 设置边距
            ax.set_xlim(-0.02, 1.02)
            ax.set_ylim(-0.02, 1.02)
            
            # 移除坐标轴
            ax.set_xticks([])
            ax.set_yticks([])
    
    # 转换为图片
    buffer = io.BytesIO()
    plt.savefig(buffer, format='png', bbox_inches='tight', pad_inches=0.05)
    buffer.seek(0)
    image_base64 = base64.b64encode(buffer.getvalue()).decode()
    plt.close()
    
    return f'data:image/png;base64,{image_base64}'

# Streamlit 界面
st.set_page_config(page_title="散厅时间快捷打印", layout="wide")
st.title("散厅时间快捷打印")

uploaded_file = st.file_uploader("上传【放映场次核对表.xls】文件", type=["xls", "xlsx"])

if uploaded_file:
    part1, part2 = process_schedule(uploaded_file)
    
    if part1 is not None and part2 is not None:
        # 创建打印布局
        part1_image = create_print_layout(part1, "第一部分")
        part2_image = create_print_layout(part2, "第二部分")
        
        # 显示预览
        col1, col2 = st.columns(2)
        
        with col1:
            st.subheader("白班散场预览（时间 ≤ 17:30）")
            if part1_image:
                st.image(part1_image)
            else:
                st.info("白班部分没有数据")
        
        with col2:
            st.subheader("夜班散场预览（时间 > 17:30）")
            if part2_image:
                st.image(part2_image)
            else:
                st.info("夜班部分没有数据")