ocr_platform/table_line_generator/backend/services/template_service.py

"""
模板服务
管理表格线条模板的创建、应用、预览和删除
"""
import json
import os
from pathlib import Path
from typing import List, Dict, Optional, Tuple
from datetime import datetime


class TemplateService:
    """表格线条模板服务"""
    
    # 模板存储目录
    _file_path = Path(__file__).resolve()
    TEMPLATE_DIR = _file_path.parent.parent / "templates"
    TEMPLATE_SUFFIX = ".template.json"
    
    def __init__(self):
        """初始化模板服务，确保模板目录存在"""
        self.TEMPLATE_DIR.mkdir(parents=True, exist_ok=True)
    
    def _get_template_path(self, name: str) -> Path:
        """获取模板文件路径"""
        # 清理名称，移除不安全字符
        safe_name = "".join(c for c in name if c.isalnum() or c in "._- ")
        return self.TEMPLATE_DIR / f"{safe_name}{self.TEMPLATE_SUFFIX}"
    
    def create_template(
        self, 
        name: str, 
        structure: Dict,
        image_size: Dict,
        description: str = ""
    ) -> Dict:
        """
        从当前结构创建模板
        
        Args:
            name: 模板名称（用户自定义）
            structure: 当前表格结构 {horizontal_lines, vertical_lines, table_bbox, ...}
            image_size: 图片尺寸 {width, height}
            description: 模板描述
            
        Returns:
            创建的模板信息
        """
        if not name or not name.strip():
            raise ValueError("模板名称不能为空")
        
        name = name.strip()
        template_path = self._get_template_path(name)
        
        if template_path.exists():
            raise ValueError(f"模板 '{name}' 已存在")
        
        # 提取模板信息
        horizontal_lines = structure.get("horizontal_lines", [])
        vertical_lines = structure.get("vertical_lines", [])
        table_bbox = structure.get("table_bbox")
        
        if not horizontal_lines and not vertical_lines:
            raise ValueError("结构中没有线条信息")
        
        # 计算相对坐标（相对于 table_bbox，这样应用时可以直接映射到目标 bbox）
        img_width = image_size.get("width", 1)
        img_height = image_size.get("height", 1)
        
        # 如果有 table_bbox，计算相对于 bbox 的坐标；否则相对于整个图片
        if table_bbox and len(table_bbox) == 4:
            bbox_x1, bbox_y1, bbox_x2, bbox_y2 = table_bbox
            bbox_width = bbox_x2 - bbox_x1
            bbox_height = bbox_y2 - bbox_y1
            
            # 相对于 table_bbox 的坐标 (0-1)
            rel_horizontal_lines = [(y - bbox_y1) / bbox_height for y in horizontal_lines] if bbox_height > 0 else []
            rel_vertical_lines = [(x - bbox_x1) / bbox_width for x in vertical_lines] if bbox_width > 0 else []
            # table_bbox 本身相对于图片尺寸
            rel_table_bbox = [
                bbox_x1 / img_width,
                bbox_y1 / img_height,
                bbox_x2 / img_width,
                bbox_y2 / img_height,
            ]
        else:
            # 没有 table_bbox 时，相对于整个图片
            rel_horizontal_lines = [y / img_height for y in horizontal_lines]
            rel_vertical_lines = [x / img_width for x in vertical_lines]
            rel_table_bbox = None
        
        template = {
            "name": name,
            "description": description,
            "created_at": datetime.now().isoformat(),
            "source_image_size": image_size,
            # 存储绝对坐标
            "horizontal_lines": horizontal_lines,
            "vertical_lines": vertical_lines,
            "table_bbox": table_bbox,
            # 相对坐标 (0-1范围)，相对于 table_bbox，用于不同尺寸图片的适配
            "relative": {
                "horizontal_lines": rel_horizontal_lines,
                "vertical_lines": rel_vertical_lines,
                "table_bbox": rel_table_bbox
            },
            # 统计信息
            "stats": {
                "row_count": len(horizontal_lines) - 1 if len(horizontal_lines) > 1 else 0,
                "col_count": len(vertical_lines) - 1 if len(vertical_lines) > 1 else 0,
            }
        }
        
        # 保存模板
        with open(template_path, 'w', encoding='utf-8') as f:
            json.dump(template, f, ensure_ascii=False, indent=2)
        
        return {
            "name": name,
            "path": str(template_path),
            "stats": template["stats"],
            "created_at": template["created_at"]
        }
    
    def list_templates(self) -> List[Dict]:
        """
        列出所有模板
        
        Returns:
            模板列表 [{name, description, created_at, stats}, ...]
        """
        templates = []
        
        for file_path in self.TEMPLATE_DIR.glob(f"*{self.TEMPLATE_SUFFIX}"):
            try:
                with open(file_path, 'r', encoding='utf-8') as f:
                    template = json.load(f)
                templates.append({
                    "name": template.get("name", file_path.stem),
                    "description": template.get("description", ""),
                    "created_at": template.get("created_at", ""),
                    "stats": template.get("stats", {}),
                    "source_image_size": template.get("source_image_size", {})
                })
            except Exception as e:
                print(f"读取模板 {file_path} 失败: {e}")
                continue
        
        # 按创建时间排序（最新在前）
        templates.sort(key=lambda x: x.get("created_at", ""), reverse=True)
        return templates
    
    def get_template(self, name: str) -> Optional[Dict]:
        """
        获取模板详情
        
        Args:
            name: 模板名称
            
        Returns:
            模板详情或 None
        """
        template_path = self._get_template_path(name)
        if not template_path.exists():
            return None
        
        with open(template_path, 'r', encoding='utf-8') as f:
            return json.load(f)
    
    def delete_template(self, name: str) -> bool:
        """
        删除模板
        
        Args:
            name: 模板名称
            
        Returns:
            是否删除成功
        """
        template_path = self._get_template_path(name)
        if template_path.exists():
            template_path.unlink()
            return True
        return False
    
    def preview_apply(
        self, 
        template_name: str, 
        target_image_size: Dict,
        target_table_bbox: Optional[List] = None,
        mode: str = "relative"
    ) -> Dict:
        """
        预览模板应用效果（不保存）
        
        Args:
            template_name: 模板名称
            target_image_size: 目标图片尺寸 {width, height}
            target_table_bbox: 目标页面的表格边界框 [x1, y1, x2, y2]
                              如果提供，则将模板线条映射到此区域内
                              如果不提供，则使用模板的 table_bbox（缩放后）
            mode: 应用模式
                - "relative": 按相对坐标缩放（适用于不同尺寸图片）
                - "absolute": 使用绝对坐标（适用于相同尺寸图片）
                
        Returns:
            预览结构 {horizontal_lines, vertical_lines, table_bbox, ...}
        """
        template = self.get_template(template_name)
        if not template:
            raise ValueError(f"模板 '{template_name}' 不存在")
        
        target_width = target_image_size.get("width", 1)
        target_height = target_image_size.get("height", 1)
        
        # 获取模板的相对坐标信息（相对于 table_bbox 的 0-1 坐标）
        relative = template.get("relative", {})
        rel_h_lines = relative.get("horizontal_lines", [])
        rel_v_lines = relative.get("vertical_lines", [])
        rel_bbox = relative.get("table_bbox")
        
        if mode == "absolute":
            # 使用绝对坐标（直接使用模板原始值）
            horizontal_lines = template.get("horizontal_lines", []).copy()
            vertical_lines = template.get("vertical_lines", []).copy()
            table_bbox = template.get("table_bbox")
        else:
            # relative 模式
            if target_table_bbox and len(target_table_bbox) == 4:
                # 使用目标页面的 table_bbox
                bbox_x1, bbox_y1, bbox_x2, bbox_y2 = target_table_bbox
                table_bbox = list(target_table_bbox)
            elif rel_bbox and len(rel_bbox) == 4:
                # 使用模板的 table_bbox（缩放到目标图片尺寸）
                bbox_x1 = int(rel_bbox[0] * target_width)
                bbox_y1 = int(rel_bbox[1] * target_height)
                bbox_x2 = int(rel_bbox[2] * target_width)
                bbox_y2 = int(rel_bbox[3] * target_height)
                table_bbox = [bbox_x1, bbox_y1, bbox_x2, bbox_y2]
            else:
                # 没有 bbox 信息，使用整个图片
                bbox_x1, bbox_y1 = 0, 0
                bbox_x2, bbox_y2 = target_width, target_height
                table_bbox = [bbox_x1, bbox_y1, bbox_x2, bbox_y2]
            
            bbox_width = bbox_x2 - bbox_x1
            bbox_height = bbox_y2 - bbox_y1
            
            # 将模板的相对位置(0-1)映射到目标 bbox 内
            horizontal_lines = [
                int(bbox_y1 + y * bbox_height) 
                for y in rel_h_lines
            ]
            vertical_lines = [
                int(bbox_x1 + x * bbox_width) 
                for x in rel_v_lines
            ]
        
        # 强制对齐：确保第一条和最后一条线与 table_bbox 边界一致
        if table_bbox and horizontal_lines:
            if len(horizontal_lines) > 0:
                horizontal_lines[0] = table_bbox[1]  # 第一条横线 = top
            if len(horizontal_lines) > 1:
                horizontal_lines[-1] = table_bbox[3]  # 最后一条横线 = bottom
        
        if table_bbox and vertical_lines:
            if len(vertical_lines) > 0:
                vertical_lines[0] = table_bbox[0]  # 第一条竖线 = left
            if len(vertical_lines) > 1:
                vertical_lines[-1] = table_bbox[2]  # 最后一条竖线 = right
        
        # 计算行列数
        total_rows = len(horizontal_lines) - 1 if len(horizontal_lines) > 1 else 0
        total_cols = len(vertical_lines) - 1 if len(vertical_lines) > 1 else 0
        
        return {
            "horizontal_lines": horizontal_lines,
            "vertical_lines": vertical_lines,
            "table_bbox": table_bbox,
            "total_rows": total_rows,
            "total_cols": total_cols,
            "modified_h_lines": [],
            "modified_v_lines": [],
            "applied_template": template_name,
            "apply_mode": mode
        }
    
    def apply_template(
        self, 
        template_name: str, 
        target_image_size: Dict,
        target_table_bbox: Optional[List] = None,
        mode: str = "relative"
    ) -> Dict:
        """
        应用模板（返回结构，由调用者决定是否保存）
        
        与 preview_apply 相同，但标记为已确认应用
        """
        structure = self.preview_apply(template_name, target_image_size, target_table_bbox, mode)
        structure["confirmed"] = True
        return structure


# 单例
_template_service: Optional[TemplateService] = None

def get_template_service() -> TemplateService:
    """获取模板服务单例"""
    global _template_service
    if _template_service is None:
        _template_service = TemplateService()
    return _template_service