ocr_verify/table_line_generator/editor/save_controls.py

"""
保存功能控件
"""
import streamlit as st
import io
import json
from pathlib import Path
from typing import Dict

from .drawing import draw_clean_table_lines


def create_save_section(work_mode: str, structure: Dict, image, line_width: int, output_config: Dict):
    """
    保存设置（目录/命名来自配置）
    
    Args:
        work_mode: 工作模式
        structure: 表格结构
        image: 图片对象
        line_width: 线条宽度
        output_config: 输出配置（兜底用）
    """
    st.divider()

    # 🔑 优先使用当前数据源的输出配置
    if 'current_output_config' in st.session_state:
        active_output_config = st.session_state.current_output_config
        st.info(f"📂 保存位置：{active_output_config.get('directory', 'N/A')}")
    else:
        active_output_config = output_config

    defaults = active_output_config.get("defaults", {})
    line_colors = active_output_config.get("line_colors") or [
        {"name": "黑色", "rgb": [0, 0, 0]},
        {"name": "蓝色", "rgb": [0, 0, 255]},
        {"name": "红色", "rgb": [255, 0, 0]},
    ]

    save_col1, save_col2, save_col3 = st.columns(3)

    with save_col1:
        save_structure = st.checkbox(
            "保存表格结构配置",
            value=bool(defaults.get("save_structure", True)),
        )

    with save_col2:
        save_image = st.checkbox(
            "保存表格线图片",
            value=bool(defaults.get("save_image", True)),
        )

    color_names = [c["name"] for c in line_colors]
    default_color = defaults.get("line_color", color_names[0])
    default_index = (
        color_names.index(default_color) 
        if default_color in color_names 
        else 0
    )

    with save_col3:
        line_color_option = st.selectbox(
            "线条颜色",
            color_names,
            index=default_index,
			label_visibility="collapsed",
            key="save_line_color"
        )

    if st.button("💾 保存", type="primary"):
        output_dir = Path(active_output_config.get("directory", "output/table_structures"))
        output_dir.mkdir(parents=True, exist_ok=True)

        structure_suffix = active_output_config.get("structure_suffix", "_structure.json")
        image_suffix = active_output_config.get("image_suffix", "_with_lines.png")

        # 确定文件名
        base_name = _determine_base_name(work_mode)
        
        saved_files = []
        
        if save_structure:
            _save_structure_file(
                structure, 
                output_dir, 
                base_name, 
                structure_suffix, 
                saved_files
            )
        
        if save_image:
            _save_image_file(
                image,
                structure,
                line_width,
                line_color_option,
                line_colors,
                output_dir,
                base_name,
                image_suffix,
                saved_files
            )
        
        if saved_files:
            st.success(f"✅ 已保存 {len(saved_files)} 个文件:")
            for file_type, file_path in saved_files:
                st.info(f"  • {file_type}: {file_path}")
            
            # 显示当前数据源信息
            if 'current_data_source' in st.session_state:
                ds = st.session_state.current_data_source
                with st.expander("📋 数据源信息"):
                    st.json({
                        "名称": ds.get("name"),
                        "JSON目录": str(ds.get("json_dir")),
                        "图片目录": str(ds.get("image_dir")),
                        "输出目录": str(output_dir),
                    })


def _determine_base_name(work_mode: str) -> str:
    """确定保存文件的基础名称"""
    if work_mode == "🆕 新建标注" or work_mode == "new":
        if st.session_state.loaded_json_name:
            return Path(st.session_state.loaded_json_name).stem
        else:
            return "table_structure"
    else:
        if st.session_state.loaded_config_name:
            base_name = Path(st.session_state.loaded_config_name).stem
            if base_name.endswith('_structure'):
                base_name = base_name[:-10]
            return base_name
        elif st.session_state.loaded_image_name:
            return Path(st.session_state.loaded_image_name).stem
        else:
            return "table_structure"


def _save_structure_file(structure, output_dir, base_name, suffix, saved_files):
    """保存结构配置文件"""
    structure_filename = f"{base_name}{suffix}"
    structure_path = output_dir / structure_filename
    # save_structure_to_config(structure, structure_path)
    with open(structure_path, 'w', encoding='utf-8') as f:
        json.dump(structure, f, indent=2, ensure_ascii=False)
    saved_files.append(("配置文件", structure_path))
    
    with open(structure_path, 'r') as f:
        st.download_button(
            "📥 下载配置文件",
            f.read(),
            file_name=f"{base_name}_structure.json",
            mime="application/json"
        )


def _save_image_file(image, structure, line_width, color_option, line_colors, 
                     output_dir, base_name, suffix, saved_files):
    """保存表格线图片"""
    if image is None:
        st.warning("⚠️ 无法保存图片：未加载图片文件")
        return
    
    selected_color_rgb = next(
        (tuple(c["rgb"]) for c in line_colors if c["name"] == color_option),
        (0, 0, 0),
    )
    
    clean_img = draw_clean_table_lines(
        image,
        structure,
        line_width=line_width,
        line_color=selected_color_rgb,
    )
    
    image_filename = f"{base_name}{suffix}"
    output_image_path = output_dir / image_filename
    clean_img.save(output_image_path)
    saved_files.append(("表格线图片", output_image_path))
    
    buf = io.BytesIO()
    clean_img.save(buf, format='PNG')
    buf.seek(0)
    
    st.download_button(
        "📥 下载表格线图片",
        buf,
        file_name=f"{base_name}_with_lines.png",
        mime="image/png"
    )