import pandas as pd
import streamlit as st
import matplotlib.pyplot as plt
import matplotlib.font_manager as font_manager
import io
import base64
import os
from datetime import datetime, timedelta
from matplotlib.patches import Rectangle
from pypinyin import lazy_pinyin, Style
from matplotlib.backends.backend_pdf import PdfPages

# --- 字体设置 ---
def get_font(size=14):
    """根据操作系统环境查找并返回中文字体属性"""
    font_path = "simHei.ttc" # 优先使用 simHei.ttc
    if not os.path.exists(font_path):
        font_path = "SimHei.ttf" # SimHei.ttf 作为备选
    # 如果两者都不存在，可以添加更多备选字体路径
    if not os.path.exists(font_path):
        # for Windows
        font_path = "C:/Windows/Fonts/simhei.ttf"
    if not os.path.exists(font_path):
        # for MacOS
         font_path = "/System/Library/Fonts/STHeiti Medium.ttc"
    # 如果仍然找不到，matplotlib会回退到默认字体
    return font_manager.FontProperties(fname=font_path, size=size)

# --- 拼音处理 ---
def get_pinyin_abbr(text):
    """获取文本前两个汉字的拼音首字母"""
    if not text or not isinstance(text, str):
        return ""
    # 提取中文字符
    chars = [c for c in text if '\u4e00' <= c <= '\u9fff']
    # 取前两个汉字
    chars_to_process = chars[:2]
    if not chars_to_process:
        return ""
    # 获取拼音首字母并转为大写
    pinyin_list = lazy_pinyin(chars_to_process, style=Style.FIRST_LETTER)
    return ''.join(pinyin_list).upper()

# --- 数据处理 ---
def process_schedule(file):
    """读取并处理 Excel 文件，返回格式化的 DataFrame 和日期"""
    try:
        # 尝试读取日期
        date_df = pd.read_excel(file, header=None, skiprows=7, nrows=1, usecols=[3])
        date_str = pd.to_datetime(date_df.iloc[0, 0]).strftime('%Y-%m-%d')
        base_date = pd.to_datetime(date_str).date()
    except Exception:
        # 读取失败则使用当天日期
        date_str = datetime.today().strftime('%Y-%m-%d')
        base_date = datetime.today().date()

    try:
        # 读取排片数据
        df = pd.read_excel(file, header=9, usecols=[1, 2, 4, 5])
        df.columns = ['Hall', 'StartTime', 'EndTime', 'Movie']
        
        # 数据清洗
        df['Hall'] = df['Hall'].ffill()
        df.dropna(subset=['StartTime', 'EndTime', 'Movie'], inplace=True)
        df['Hall'] = df['Hall'].astype(str).str.extract(r'(\d+号)')
        df.dropna(subset=['Hall'], inplace=True)

        # 时间转换
        df['StartTime_dt'] = pd.to_datetime(df['StartTime'], format='%H:%M', errors='coerce').apply(
            lambda t: t.replace(year=base_date.year, month=base_date.month, day=base_date.day) if pd.notnull(t) else t
        )
        df['EndTime_dt'] = pd.to_datetime(df['EndTime'], format='%H:%M', errors='coerce').apply(
            lambda t: t.replace(year=base_date.year, month=base_date.month, day=base_date.day) if pd.notnull(t) else t
        )
        df.dropna(subset=['StartTime_dt', 'EndTime_dt'], inplace=True)
        
        # 处理跨天结束时间
        df.loc[df['EndTime_dt'] < df['StartTime_dt'], 'EndTime_dt'] += timedelta(days=1)
        df = df.sort_values(['Hall', 'StartTime_dt'])
        
        # 合并同一影厅的连续相同影片
        merged_rows = []
        for _, group in df.groupby('Hall'):
            group = group.sort_values('StartTime_dt')
            current = None
            for _, row in group.iterrows():
                if current is None:
                    current = row.copy()
                else:
                    if row['Movie'] == current['Movie']:
                        current['EndTime_dt'] = row['EndTime_dt']
                    else:
                        merged_rows.append(current)
                        current = row.copy()
            if current is not None:
                merged_rows.append(current)
        
        if not merged_rows:
            return None, date_str
            
        merged_df = pd.DataFrame(merged_rows)
        
        # 统一调整时间
        merged_df['StartTime_dt'] -= timedelta(minutes=10)
        merged_df['EndTime_dt'] -= timedelta(minutes=5)
        
        # 格式化最终输出的列
        merged_df['Time'] = merged_df['StartTime_dt'].dt.strftime('%H:%M') + ' - ' + merged_df['EndTime_dt'].dt.strftime('%H:%M')
        merged_df['Pinyin'] = merged_df['Movie'].apply(get_pinyin_abbr)
        merged_df['Hall'] = merged_df['Hall'].str.replace('号', '')
        
        return merged_df[['Hall', 'Movie', 'Pinyin', 'Time']], date_str

    except Exception as e:
        st.error(f"处理 Excel 数据时发生错误: {e}")
        return None, date_str


# --- 打印布局生成 ---
def create_print_layout(data, date_str):
    """根据处理好的数据，生成用于打印的 PNG 和 PDF 布局"""
    if data is None or data.empty:
        return None

    A4_SIZE_INCHES = (8.27, 11.69)
    DPI = 300
    
    # 准备一个临时的 figure 用于计算文本渲染尺寸
    temp_fig = plt.figure(figsize=A4_SIZE_INCHES, dpi=DPI)
    renderer = temp_fig.canvas.get_renderer()

    # 1. 计算行高
    num_movies = len(data)
    num_halls = len(data['Hall'].unique())
    # 总行数 = 电影条目数 + 厅间分隔数 + 上下留白(2)
    # total_layout_rows = num_movies + (num_halls - 1) + 2
    total_layout_rows = num_movies + 2 # 简化为条目数+2，使行高更宽松
    row_height_inch = A4_SIZE_INCHES[1] / total_layout_rows
    
    # 2. 计算基准字体大小 (点)
    # 1 point = 1/72 inch.
    # 字体高度为行高的 90%
    font_size_pt = row_height_inch * 0.9 * 72 
    base_font = get_font(font_size_pt)

    # 3. 计算各列宽度 (除电影名外)
    def get_text_width_inch(text, font):
        t = plt.text(0, 0, text, fontproperties=font)
        bbox = t.get_window_extent(renderer=renderer)
        width_pixels = bbox.width
        t.remove()
        return width_pixels / DPI

    # 找到每列最长的内容
    data['Index'] = data.groupby('Hall').cumcount() + 1
    max_hall_str = data['Hall'].max() + "#" # e.g. "10#"
    max_index_str = str(data['Index'].max()) + "." # e.g. "5."
    max_pinyin_str = data['Pinyin'].apply(len).max() * "A" # e.g. "PY" -> "AA"
    max_time_str = data['Time'].apply(len).idxmax()
    max_time_str = data.loc[max_time_str, 'Time']

    col_widths = {}
    col_widths['Hall'] = get_text_width_inch(max_hall_str, base_font) * 1.1
    col_widths['Index'] = get_text_width_inch(max_index_str, base_font) * 1.1
    col_widths['Pinyin'] = get_text_width_inch(max_pinyin_str, base_font) * 1.1
    col_widths['Time'] = get_text_width_inch(max_time_str, base_font) * 1.1
    
    # 电影名列的宽度为剩余宽度
    non_movie_width = sum(col_widths.values())
    col_widths['Movie'] = A4_SIZE_INCHES[0] - non_movie_width

    plt.close(temp_fig) # 关闭临时figure

    # --- 开始正式绘图 ---
    figs = {}
    for fmt in ['png', 'pdf']:
        fig = plt.figure(figsize=A4_SIZE_INCHES, dpi=DPI)
        ax = fig.add_subplot(111)
        ax.set_axis_off()
        fig.subplots_adjust(left=0, right=1, top=1, bottom=0)
        ax.set_xlim(0, A4_SIZE_INCHES[0])
        ax.set_ylim(0, A4_SIZE_INCHES[1])

        # 计算列的 X 轴起始位置
        x_pos = {}
        current_x = 0
        # 新顺序: Hall, Index, Movie, Pinyin, Time
        col_order = ['Hall', 'Index', 'Movie', 'Pinyin', 'Time']
        for col in col_order:
            x_pos[col] = current_x
            current_x += col_widths[col]

        # 从顶部开始绘制 (顶部留出一行空白)
        current_y = A4_SIZE_INCHES[1] - row_height_inch 
        
        halls = sorted(data['Hall'].unique(), key=lambda h: int(h))

        for hall in halls:
            hall_data = data[data['Hall'] == hall].sort_values('Index')
            
            for _, row in hall_data.iterrows():
                y_bottom = current_y - row_height_inch
                
                # 绘制单元格
                for col_name in col_order:
                    cell_x = x_pos[col_name]
                    cell_width = col_widths[col_name]
                    
                    # 绘制灰色虚线边框
                    rect = Rectangle((cell_x, y_bottom), cell_width, row_height_inch,
                                     edgecolor='lightgray', facecolor='none',
                                     linestyle=(0, (1, 2)), linewidth=0.8, zorder=1)
                    ax.add_patch(rect)
                    
                    # 准备文本内容
                    text_content = {
                        'Hall': f"{row['Hall']}#",
                        'Index': f"{row['Index']}.",
                        'Movie': row['Movie'],
                        'Pinyin': row['Pinyin'],
                        'Time': row['Time']
                    }[col_name]

                    # 文本垂直居中
                    text_y = y_bottom + row_height_inch / 2
                    
                    # 电影名列特殊处理
                    if col_name == 'Movie':
                        font_to_use = base_font.copy()
                        # 检查宽度并调整字体
                        text_w_inch = get_text_width_inch(text_content, font_to_use)
                        max_w_inch = cell_width * 0.9 # 目标宽度为单元格宽度的90%
                        if text_w_inch > max_w_inch:
                            scale_factor = max_w_inch / text_w_inch
                            font_to_use.set_size(font_size_pt * scale_factor)

                        ax.text(cell_x + cell_width * 0.05, text_y, text_content, # 左对齐
                                fontproperties=font_to_use, ha='left', va='center', clip_on=True)
                    else: # 其他列
                        ax.text(cell_x + cell_width / 2, text_y, text_content, # 居中对齐
                                fontproperties=base_font, ha='center', va='center', clip_on=True)

                current_y -= row_height_inch
            
            # 在每个影厅块结束后绘制黑色分隔线
            ax.plot([0, A4_SIZE_INCHES[0]], [current_y, current_y], color='black', linewidth=1.5, zorder=2)

        # 在左上角添加日期
        ax.text(0.1, A4_SIZE_INCHES[1] - 0.3, date_str, 
                fontproperties=get_font(12), color='gray', ha='left', va='top')

        figs[fmt] = fig

    # 保存到内存
    png_buffer = io.BytesIO()
    figs['png'].savefig(png_buffer, format='png', dpi=DPI)
    png_buffer.seek(0)
    image_base64 = base64.b64encode(png_buffer.getvalue()).decode()
    plt.close(figs['png'])

    pdf_buffer = io.BytesIO()
    figs['pdf'].savefig(pdf_buffer, format='pdf', dpi=DPI)
    pdf_buffer.seek(0)
    pdf_base64 = base64.b64encode(pdf_buffer.getvalue()).decode()
    plt.close(figs['pdf'])

    return {
        'png': f"data:image/png;base64,{image_base64}",
        'pdf': f"data:application/pdf;base64,{pdf_base64}"
    }


# --- Streamlit UI ---
st.set_page_config(page_title="LED 屏幕时间表打印", layout="wide")
st.title("LED 屏幕时间表打印")
st.markdown("请上传影院系统导出的 `放映时间核对表.xls` 文件。系统将自动处理数据并生成专业、美观的A4打印布局。")

uploaded_file = st.file_uploader("选择文件", accept_multiple_files=False, type=["xls", "xlsx"])

if uploaded_file:
    with st.spinner("文件处理与布局生成中，请稍候..."):
        schedule, date_str = process_schedule(uploaded_file)
        if schedule is not None and not schedule.empty:
            output = create_print_layout(schedule, date_str)
            
            st.success(f"成功生成 **{date_str}** 的排片表！")

            # 创建下载按钮
            col1, col2 = st.columns(2)
            with col1:
                 st.download_button(
                    label="📥 下载 PNG 图像",
                    data=base64.b64decode(output['png'].split(',')[1]),
                    file_name=f"排片表_{date_str}.png",
                    mime="image/png"
                )
            with col2:
                st.download_button(
                    label="📄 下载 PDF 文档",
                    data=base64.b64decode(output['pdf'].split(',')[1]),
                    file_name=f"排片表_{date_str}.pdf",
                    mime="application/pdf"
                )

            # 创建选项卡进行预览
            tab1, tab2 = st.tabs(["📄 PDF 预览", "🖼️ PNG 预览"])
            
            with tab1:
                st.markdown(f'<iframe src="{output["pdf"]}" width="100%" height="800" type="application/pdf"></iframe>', unsafe_allow_html=True)
            
            with tab2:
                st.image(output['png'], use_container_width=True)
        else:
            st.error("无法处理文件。请检查文件内容是否为空或格式是否正确。确保文件中包含有效的排片数据。")