s2f/backend/app/services/ai_recognizer.py

"""
AI 字段识别服务

使用 AI 自动识别 Excel 表头对应的标准字段
"""

import json
import re
from typing import List, Dict, Optional, Tuple

from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy import select

from app.models.standard_field import StandardField, FIELD_TYPES


class AIFieldRecognizer:
    """AI 字段识别器"""

    # 字段关键词映射（用于规则匹配兜底）
    FIELD_KEYWORDS = {
        StandardField.EMPLOYEE_NAME.value: ["姓名", "名字", "员工名", "name", "员工姓名"],
        StandardField.EMPLOYEE_ID.value: ["工号", "员工号", "编号", "id", "员工编号"],
        StandardField.DEPARTMENT.value: ["部门", "科室", "事业部", "department", "所属部门"],
        StandardField.POSITION.value: ["岗位", "职位", "职务", "position", "job"],
        StandardField.BASE_SALARY.value: ["基本工资", "岗位工资", "底薪", "base", "基本薪资"],
        StandardField.BONUS.value: ["奖金", "绩效", "bonus", "绩效工资", "奖励"],
        StandardField.ALLOWANCE.value: ["补贴", "津贴", "allowance", "餐补", "交通补贴"],
        StandardField.OVERTIME_PAY.value: ["加班费", "加班工资", "overtime"],
        StandardField.DEDUCTION.value: ["扣款", "扣除", "deduction", "罚款", "迟到扣款"],
        StandardField.GROSS_SALARY.value: ["应发工资", "应发", "税前工资", "gross", "工资总额"],
        StandardField.NET_SALARY.value: ["实发工资", "实发", "净工资", "net", "实发金额", "银行实发"],
        StandardField.BANK_CARD.value: ["银行账号", "卡号", "账号", "bank", "银行卡"],
        StandardField.ID_CARD.value: ["身份证", "证件号", "id_card", "身份证号"],
        StandardField.SOCIAL_SECURITY_BASE.value: ["社保基数", "缴费基数", "基数"],
        StandardField.PENSION_INSURANCE.value: ["养老保险", "养老", "pension", "养保"],
        StandardField.MEDICAL_INSURANCE.value: ["医疗保险", "医保", "medical"],
        StandardField.UNEMPLOYMENT_INSURANCE.value: ["失业保险", "失业", "unemployment"],
        StandardField.HOUSING_FUND.value: ["公积金", "住房基金", "housing", "住房公金"],
        StandardField.PENSION_INSURANCE_COMPANY.value: ["养老保险(公司)", "养老保险公司", "养保公司"],
        StandardField.MEDICAL_INSURANCE_COMPANY.value: ["医疗保险(公司)", "医疗公司"],
        StandardField.UNEMPLOYMENT_INSURANCE_COMPANY.value: ["失业保险(公司)", "失业公司"],
        StandardField.HOUSING_FUND_COMPANY.value: ["公积金(公司)", "公积金公司"],
        StandardField.SOCIAL_SECURITY_TOTAL.value: ["社保合计", "社保总计", "社保总额"],
        StandardField.TAXABLE_INCOME.value: ["应税收入", "应税工资", "税前收入", "taxable"],
        StandardField.PRE_TAX_DEDUCTION.value: ["税前扣除", "三险一金", "个人缴费"],
        StandardField.TAX_FREE_INCOME.value: ["免税收入", "免税", "tax_free"],
        StandardField.TAX_EXEMPT_INCOME.value: ["税前减免", "减免"],
        StandardField.QUICK_DEDUCTION.value: ["速算扣除", "速算"],
        StandardField.TAX_AMOUNT.value: ["应缴个税", "个人所得税", "个税", "tax"],
        StandardField.TAX_PAID.value: ["已缴个税", "已扣税", "已缴税"],
        StandardField.AFTER_TAX_INCOME.value: ["税后收入", "税后工资", "after_tax"],
    }

    def __init__(self, db: AsyncSession, company_id: int):
        self.db = db
        self.company_id = company_id

    async def recognize_fields(
        self,
        headers: List[str],
        sample_data: List[Dict],
        file_type: str = "工资表"
    ) -> List[Dict]:
        """
        识别字段映射

        Args:
            headers: 表头列表
            sample_data: 样例数据（前10行）
            file_type: 文件类型

        Returns:
            字段映射列表
        """
        mappings = []

        for header in headers:
            # 获取该列的样例值
            samples = [row.get(header) for row in sample_data if row.get(header)]
            sample_values = samples[:5] if samples else []

            # 先尝试规则匹配
            matched_field, confidence, reasoning = self._rule_match(header, sample_values)

            if matched_field:
                mappings.append({
                    "source_field": header,
                    "standard_field": matched_field,
                    "confidence": confidence,
                    "reasoning": reasoning,
                    "sample_values": sample_values,
                })
            else:
                # 无法匹配
                mappings.append({
                    "source_field": header,
                    "standard_field": "",
                    "confidence": 0.0,
                    "reasoning": "无法识别字段类型",
                    "sample_values": sample_values,
                })

        return mappings

    def _rule_match(
        self,
        header: str,
        sample_values: List
    ) -> Tuple[Optional[str], float, str]:
        """
        规则匹配

        Args:
            header: 字段名
            sample_values: 样例值

        Returns:
            (标准字段, 置信度, 判断依据)
        """
        header_lower = header.lower().strip()
        header_normalized = header.strip()

        best_match = None
        best_confidence = 0.0
        best_reasoning = ""

        for standard_field, keywords in self.FIELD_KEYWORDS.items():
            for keyword in keywords:
                keyword_lower = keyword.lower()

                # 精确匹配（完全相同）
                if header_normalized == keyword or header_lower == keyword_lower:
                    return (
                        standard_field,
                        1.0,
                        f"字段名完全匹配：'{keyword}'"
                    )

                # 包含匹配
                if keyword_lower in header_lower or header_lower in keyword_lower:
                    confidence = 0.9
                    reasoning = f"字段名包含关键词：'{keyword}'"

                    # 数值字段检查样例值
                    field_type = FIELD_TYPES.get(standard_field, "string")
                    if field_type == "number" and sample_values:
                        if self._validate_numeric_samples(sample_values):
                            confidence = 0.95
                            reasoning += "，样例值验证为数值类型"

                    if confidence > best_confidence:
                        best_match = standard_field
                        best_confidence = confidence
                        best_reasoning = reasoning

                # 模糊匹配（编辑距离）
                distance = self._levenshtein_distance(header_lower, keyword_lower)
                max_len = max(len(header_lower), len(keyword_lower))
                similarity = 1 - (distance / max_len) if max_len > 0 else 0

                if similarity > 0.7 and similarity > best_confidence:
                    best_match = standard_field
                    best_confidence = similarity * 0.8  # 模糊匹配降权
                    best_reasoning = f"字段名相似度：{similarity:.0%}，参考词：'{keyword}'"

        return best_match, best_confidence, best_reasoning

    def _validate_numeric_samples(self, samples: List) -> bool:
        """验证样例值是否为数值"""
        numeric_count = 0
        for sample in samples[:5]:
            if sample is None:
                continue
            sample_str = str(sample).strip()
            # 移除常见的货币符号和逗号
            sample_str = sample_str.replace("¥", "").replace(",", "").replace("元", "")
            try:
                float(sample_str)
                numeric_count += 1
            except ValueError:
                pass
        return numeric_count >= len(samples) * 0.8

    @staticmethod
    def _levenshtein_distance(s1: str, s2: str) -> int:
        """计算编辑距离"""
        if len(s1) < len(s2):
            return AIFieldRecognizer._levenshtein_distance(s2, s1)

        if len(s2) == 0:
            return len(s1)

        previous_row = range(len(s2) + 1)
        for i, c1 in enumerate(s1):
            current_row = [i + 1]
            for j, c2 in enumerate(s2):
                insertions = previous_row[j + 1] + 1
                deletions = current_row[j] + 1
                substitutions = previous_row[j] + (c1 != c2)
                current_row.append(min(insertions, deletions, substitutions))
            previous_row = current_row

        return previous_row[-1]


async def recognize_fields_with_ai(
    db: AsyncSession,
    company_id: int,
    headers: List[str],
    sample_data: List[Dict],
    file_type: str = "工资表"
) -> List[Dict]:
    """
    使用 AI 识别字段（带重试的版本）

    优先使用 OpenAI API，失败时降级到规则匹配
    """
    recognizer = AIFieldRecognizer(db, company_id)

    # 优先使用规则匹配（当前实现）
    # TODO: 后续集成 OpenAI API
    mappings = await recognizer.recognize_fields(headers, sample_data, file_type)

    return mappings