#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
补扣款数据模型
单表: shein_replenish_data
使用SQLAlchemy定义表结构，支持批量upsert
"""

from sqlalchemy import create_engine, Column, Integer, String, Text, DateTime, Numeric, UniqueConstraint, Index
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.dialects.mysql import insert
from sqlalchemy.exc import OperationalError
from sqlalchemy.orm import sessionmaker
from datetime import datetime
import time

import openpyxl

Base = declarative_base()

# Excel列名 → DB列名 映射 (顺序与Excel列一致)
EXCEL_COLUMN_MAPPING = [
    ('补扣款单号', 'replenish_no'),
    ('款项类型', 'replenish_type_name'),
    ('补扣款分类', 'category_name'),
    ('对单类型', 'to_order_type_name'),
    ('关联单据', 'relation_no'),
    ('单价', 'unit_price'),
    ('数量', 'quantity'),
    ('总金额', 'amount'),
    ('币种', 'currency_code'),
    ('创建时间', 'add_time'),
    ('单据状态', 'replenish_status_name'),
    ('关联报账单', 'report_order_no'),
    ('拒绝原因', 'refuse_reason'),
    ('确认/拒绝时间', 'decision_time'),
    ('操作人', 'operator'),
    ('会计日期', 'account_date'),
    ('是否可报账', 'reportable_name'),
    ('申诉单号', 'appeal_no'),
    ('公司主体', 'company_name'),
    ('出口模式', 'exporting_mode_name'),
    ('费用类型', 'expense_type_name'),
    ('备注', 'remark'),
    ('代收服务费', 'collection_fee'),
]


class SheinReplenishData(Base):
    """补扣款表"""
    __tablename__ = 'shein_replenish_data'

    id = Column(Integer, primary_key=True, autoincrement=True, comment='主键ID')
    replenish_no = Column(String(100), nullable=True, comment='补扣款单号')
    replenish_type_name = Column(String(50), nullable=True, comment='款项类型')
    category_name = Column(String(100), nullable=True, comment='补扣款分类')
    to_order_type_name = Column(String(50), nullable=True, comment='对单类型')
    relation_no = Column(String(200), nullable=True, comment='关联单据')
    unit_price = Column(Numeric(12, 2), nullable=True, comment='单价')
    quantity = Column(Integer, nullable=True, comment='数量')
    amount = Column(Numeric(12, 2), nullable=True, comment='总金额')
    currency_code = Column(String(10), nullable=True, comment='币种')
    add_time = Column(String(30), nullable=True, comment='创建时间')
    replenish_status_name = Column(String(50), nullable=True, comment='单据状态')
    report_order_no = Column(String(100), nullable=True, comment='关联报账单')
    refuse_reason = Column(Text, nullable=True, comment='拒绝原因')
    decision_time = Column(String(30), nullable=True, comment='确认/拒绝时间')
    operator = Column(String(50), nullable=True, comment='操作人')
    account_date = Column(String(30), nullable=True, comment='会计日期')
    reportable_name = Column(String(10), nullable=True, comment='是否可报账')
    appeal_no = Column(String(50), nullable=True, comment='申诉单号')
    company_name = Column(String(200), nullable=True, comment='公司主体')
    exporting_mode_name = Column(String(50), nullable=True, comment='出口模式')
    expense_type_name = Column(String(50), nullable=True, comment='费用类型')
    remark = Column(Text, nullable=True, comment='备注')
    collection_fee = Column(Numeric(12, 2), nullable=True, comment='代收服务费')
    store_username = Column(String(50), nullable=False, comment='店铺账号')
    store_name = Column(String(100), nullable=True, comment='店铺别名')
    created_at = Column(DateTime, default=datetime.now, comment='创建时间')

    __table_args__ = (
        UniqueConstraint('replenish_no', 'store_username', name='uk_replenish_unique'),
        Index('idx_replenish_store', 'store_username'),
        Index('idx_replenish_add_time', 'add_time'),
    )


class ReplenishDataManager:
    """
    补扣款数据管理器
    """

    def __init__(self, database_url):
        self.engine = create_engine(
            database_url,
            echo=False,
            pool_size=20,
            max_overflow=10,
            pool_recycle=3600,
            pool_pre_ping=True,
        )
        self.Session = sessionmaker(bind=self.engine, autoflush=False)

    def create_table(self):
        """创建数据表"""
        Base.metadata.create_all(self.engine, tables=[SheinReplenishData.__table__], checkfirst=True)
        print("[OK] 表 shein_replenish_data 创建成功")

    def read_excel_data(self, excel_path, store_username, store_name):
        """
        从Excel文件读取补扣款数据

        Args:
            excel_path: Excel文件路径
            store_username: 店铺账号
            store_name: 店铺别名
        Returns:
            list: 数据字典列表
        """
        wb = openpyxl.load_workbook(excel_path)
        ws = wb.active

        if ws.max_row is None or ws.max_row < 2:
            wb.close()
            return []

        rows = []
        for row in ws.iter_rows(values_only=True):
            rows.append(row)

        headers = [str(h).strip() if h else '' for h in rows[0]]

        col_map = dict(EXCEL_COLUMN_MAPPING)
        header_to_db = {}
        for idx, header in enumerate(headers):
            if header in col_map:
                header_to_db[idx] = col_map[header]

        records = []
        for row in rows[1:]:
            record = {}
            for idx, db_col in header_to_db.items():
                val = row[idx] if idx < len(row) else None
                if val is not None and str(val).strip() in ('', '/', '-'):
                    val = None
                record[db_col] = val
            record['store_username'] = store_username
            record['store_name'] = store_name
            record['created_at'] = datetime.now()
            records.append(record)

        wb.close()
        return records

    def upsert_data(self, records, batch_size=1000):
        """
        批量upsert补扣款数据

        Args:
            records: 数据字典列表
            batch_size: 每批次大小
        """
        if not records:
            print("没有数据需要插入")
            return

        self.create_table()

        max_retries = 3
        with self.engine.connect() as conn:
            for i in range(0, len(records), batch_size):
                batch = records[i:i + batch_size]
                stmt = insert(SheinReplenishData).values(batch)
                stmt = stmt.on_duplicate_key_update(
                    replenish_type_name=stmt.inserted.replenish_type_name,
                    category_name=stmt.inserted.category_name,
                    amount=stmt.inserted.amount,
                    replenish_status_name=stmt.inserted.replenish_status_name,
                    report_order_no=stmt.inserted.report_order_no,
                    refuse_reason=stmt.inserted.refuse_reason,
                    decision_time=stmt.inserted.decision_time,
                    remark=stmt.inserted.remark,
                )
                for retry in range(max_retries):
                    try:
                        conn.execute(stmt)
                        break
                    except OperationalError as e:
                        if 'Deadlock' in str(e) and retry < max_retries - 1:
                            print(f"[WARN] 死锁，重试 {retry + 1}/{max_retries}...")
                            time.sleep(1)
                            continue
                        raise
            conn.commit()

        print(f"[OK] 成功 upsert {len(records)} 条补扣款数据")

    def import_from_excel_files(self, excel_paths, store_username, store_name):
        """
        从多个Excel文件导入数据

        Args:
            excel_paths: Excel文件路径列表
            store_username: 店铺账号
            store_name: 店铺别名
        """
        all_records = []
        for path in excel_paths:
            records = self.read_excel_data(path, store_username, store_name)
            all_records.extend(records)
            print(f"[INFO] 读取 {path}: {len(records)} 条记录")

        if all_records:
            self.upsert_data(all_records)
        else:
            print(f"[INFO] 没有补扣款数据需要导入")
