"""
Excel to SQL Server Migration Script
Migrates production execution card data from Excel (.xlsm) to SQL Server.
Configuration is loaded from config.yaml.
"""

import argparse
import os
import sys
from collections import OrderedDict
from datetime import datetime

import openpyxl
import pyodbc
import yaml


# ============================================================
#  配置加载
# ============================================================

def load_config(config_path=None):
    """加载 YAML 配置文件"""
    if config_path is None:
        config_path = os.path.join(
            os.path.dirname(os.path.abspath(__file__)), 'config.yaml'
        )
    with open(config_path, 'r', encoding='utf-8') as f:
        return yaml.safe_load(f)


def get_connection_string(cfg):
    """根据配置构建 pyodbc 连接字符串"""
    sql = cfg['sql_server']
    return (
        f"DRIVER={{{sql['driver']}}};"
        f"SERVER={sql['server']};"
        f"DATABASE={sql['database']};"
        f"UID={sql['username']};"
        f"PWD={sql['password']};"
        f"TrustServerCertificate={sql['TrustServerCertificate']};"
    )


# ============================================================
#  Excel 读取与数据合并
# ============================================================

def sanitize_headers(headers):
    """处理原始表头：
    1. 空列名 → col_N
    2. 重复列名 → 第二次出现加 _2 后缀
    返回处理后的列名列表
    """
    seen = {}
    result = []
    for i, h in enumerate(headers):
        h = str(h).strip() if h else ''
        if not h:
            h = f'col_{i}'
        if h in seen:
            result.append(f'{h}_{seen[h] + 1}')
            seen[h] += 1
        else:
            seen[h] = 1
            result.append(h)
    return result


def remap_headers(headers, mapping):
    """根据 config.yaml 中的 column_mapping 重映射列名"""
    if not mapping:
        return headers
    return [mapping.get(h, h) for h in headers]


def resolve_pk_index(headers, primary_key):
    """找到主键列的索引。
    优先精确匹配 primary_key，其次匹配含 'BW' 的列，最后回退到第 0 列。
    """
    for idx, h in enumerate(headers):
        if h == primary_key:
            return idx
    # 回退：首列含 BW（如 20BW序号）
    for idx, h in enumerate(headers):
        if 'BW' in h:
            return idx
    return 0


def read_sheet(file_path, sheet_name, primary_key, cfg):
    """读取单个工作表。
    返回 (列名列表, 实际主键列名, 数据列表)。
    每行数据为 dict {列名: 值}。
    跳过主键列为空的行。
    """
    wb = openpyxl.load_workbook(file_path, read_only=True, data_only=True)
    ws = wb[sheet_name]

    rows_iter = ws.iter_rows(values_only=True)
    # 第一行为表头
    raw_headers = next(rows_iter)
    headers = sanitize_headers(raw_headers)
    headers = remap_headers(headers, cfg.get('column_mapping'))

    pk_idx = resolve_pk_index(headers, primary_key)
    actual_pk_col = headers[pk_idx]

    data = []
    for row in rows_iter:
        pk_val = row[pk_idx] if pk_idx < len(row) else None
        if pk_val is None or str(pk_val).strip() == '':
            continue

        record = {}
        for j, h in enumerate(headers):
            val = row[j] if j < len(row) else None
            record[h] = val
        data.append(record)

    wb.close()
    return headers, actual_pk_col, data


def merge_sources(sources, excel_dir, primary_key, cfg):
    """合并多个工作表的数据。
    - 列名取并集（按出现顺序）
    - 数据按主键去重，后读覆盖先读
    返回 (unified_headers, actual_pk_col, merged_data_ordered_dict)
    """
    all_headers = []       # 保持顺序的列名并集
    header_set = set()
    merged = OrderedDict()  # key=主键值, value=dict
    resolved_pk = primary_key  # 实际使用的主键列名

    for src in sources:
        file_path = os.path.join(excel_dir, src['file'])
        sheet_name = src['sheet']

        headers, actual_pk_col, data = read_sheet(file_path, sheet_name, primary_key, cfg)
        resolved_pk = actual_pk_col

        # 扩展统一列名集
        for h in headers:
            if h not in header_set:
                all_headers.append(h)
                header_set.add(h)

        # 合并数据（后读覆盖先读）
        for record in data:
            pk_val = str(record.get(actual_pk_col, '')).strip()
            if not pk_val:
                continue
            if pk_val in merged:
                merged[pk_val].update(record)
            else:
                merged[pk_val] = record

    return all_headers, resolved_pk, merged


# ============================================================
#  SQL Server 操作
# ============================================================

def get_connection(cfg):
    """根据配置建立 SQL Server 连接"""
    conn_str = get_connection_string(cfg)
    return pyodbc.connect(conn_str, autocommit=False)


def create_schema(cursor, schema_name):
    """创建 schema（如不存在）"""
    sql = f"""
    IF NOT EXISTS (
        SELECT 1 FROM sys.schemas WHERE name = ?
    )
    BEGIN
        EXEC('CREATE SCHEMA [{schema_name}]')
    END
    """
    cursor.execute(sql, (schema_name,))
    cursor.commit()


def infer_sql_type(col_name, type_rules, primary_key):
    """根据 config.yaml 中的 type_rules 推断 SQL Server 数据类型"""
    name = col_name

    # 主键列（含 TM_2022 的 20BW序号）
    if name == primary_key or '20BW' in name:
        return type_rules['primary_key_type']

    # DATETIME
    if any(kw in name for kw in type_rules.get('datetime', [])):
        return 'DATETIME'

    # INT（精确匹配）
    for rule in type_rules.get('int', []):
        if rule.get('exact') and rule['name'] == name:
            return 'INT'
        elif not rule.get('exact') and rule['name'] in name:
            return 'INT'

    # DECIMAL（精确匹配）
    for rule in type_rules.get('decimal', []):
        if rule.get('exact') and rule['name'] == name:
            return 'DECIMAL(18,2)'
        elif not rule.get('exact') and rule['name'] in name:
            return 'DECIMAL(18,2)'

    # NVARCHAR(MAX) 长文本
    if any(kw in name for kw in type_rules.get('long_text', [])):
        return 'NVARCHAR(MAX)'

    return type_rules.get('default_type', 'NVARCHAR(500)')


def create_table(cursor, table_name, headers, cfg, actual_pk):
    """动态创建表。如已存在则先 DROP。"""
    schema = cfg['schema']
    type_rules = cfg['type_rules']

    full_name = f'[{schema}].[{table_name}]'

    # DROP if exists
    cursor.execute(f"IF OBJECT_ID('{full_name}', 'U') IS NOT NULL DROP TABLE {full_name}")
    cursor.commit()

    # 构建列定义
    col_defs = ['[id] INT IDENTITY(1,1)']
    for h in headers:
        sql_type = infer_sql_type(h, type_rules, actual_pk)
        col_defs.append(f'[{h}] {sql_type}')

    # 主键约束
    col_defs.append(f'CONSTRAINT [PK_{table_name}] PRIMARY KEY ([{actual_pk}])')

    create_sql = f"CREATE TABLE {full_name} ({', '.join(col_defs)})"
    cursor.execute(create_sql)
    cursor.commit()

    return create_sql


def build_typed_cols(headers, type_rules, actual_pk):
    """构建需要做类型转换的列映射 {列名: 'datetime'|'int'|'decimal'}"""
    typed = {}
    for h in headers:
        if h == actual_pk or '20BW' in h:
            continue
        if any(kw in h for kw in type_rules.get('datetime', [])):
            typed[h] = 'datetime'
        for rule in type_rules.get('int', []):
            match = (rule['name'] == h) if rule.get('exact') else (rule['name'] in h)
            if match:
                typed[h] = 'int'
        for rule in type_rules.get('decimal', []):
            match = (rule['name'] == h) if rule.get('exact') else (rule['name'] in h)
            if match:
                typed[h] = 'decimal'
    return typed


def safe_datetime(val):
    """将值转换为 datetime，失败则返回 None"""
    if val is None:
        return None
    if isinstance(val, str) and val.strip() == '':
        return None
    if isinstance(val, datetime):
        if val.year < 1753:
            return None
        return val
    if isinstance(val, str):
        for fmt in ('%Y-%m-%d %H:%M:%S', '%Y-%m-%d %H:%M', '%Y-%m-%d',
                     '%Y/%m/%d %H:%M:%S', '%Y/%m/%d %H:%M', '%Y/%m/%d'):
            try:
                dt = datetime.strptime(val.strip(), fmt)
                if dt.year >= 1753:
                    return dt
                return None
            except ValueError:
                continue
        return None
    return None


def safe_int(val):
    """将值转换为 int，失败则返回 None"""
    if val is None:
        return None
    if isinstance(val, str) and val.strip() == '':
        return None
    if isinstance(val, (int, float)):
        return int(val)
    try:
        return int(float(str(val).strip()))
    except (ValueError, TypeError):
        return None


def safe_decimal(val):
    """将值转换为 float（pyodbc 会映射到 DECIMAL），失败则返回 None"""
    if val is None:
        return None
    if isinstance(val, str) and val.strip() == '':
        return None
    if isinstance(val, (int, float)):
        return float(val)
    try:
        return float(str(val).strip())
    except (ValueError, TypeError):
        return None


TYPE_CONVERTERS = {
    'datetime': safe_datetime,
    'int': safe_int,
    'decimal': safe_decimal,
}


def batch_insert(cursor, table_name, headers, data_dict, cfg, actual_pk):
    """批量插入数据到 SQL Server。"""
    schema = cfg['schema']
    batch_size = cfg.get('batch_size', 500)
    type_rules = cfg['type_rules']
    full_name = f'[{schema}].[{table_name}]'

    typed_cols = build_typed_cols(headers, type_rules, actual_pk)

    col_list = ', '.join(f'[{h}]' for h in headers)
    placeholders = ', '.join(['?'] * len(headers))
    insert_sql = f"INSERT INTO {full_name} ({col_list}) VALUES ({placeholders})"

    batch = []
    total = 0
    for pk, record in data_dict.items():
        row = []
        for h in headers:
            val = record.get(h)
            if isinstance(val, str) and val.strip() == '':
                val = None
            # 强类型列清洗
            if h in typed_cols and val is not None:
                converter = TYPE_CONVERTERS[typed_cols[h]]
                val = converter(val)
            row.append(val)
        batch.append(row)
        total += 1

        if len(batch) >= batch_size:
            cursor.executemany(insert_sql, batch)
            cursor.commit()
            batch = []

    if batch:
        cursor.executemany(insert_sql, batch)
        cursor.commit()

    return total


# ============================================================
#  主流程
# ============================================================

def migrate_table(cursor, table_config, cfg, dry_run=False):
    """迁移单个表"""
    table_name = table_config['table']
    sources = table_config['sources']
    primary_key = cfg['primary_key']
    excel_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'Excel')

    print(f"\n  迁移 {table_name} ...")

    # 1. 读取并合并数据
    headers, actual_pk, merged = merge_sources(sources, excel_dir, primary_key, cfg)

    # 打印各源行数
    for src in sources:
        fpath = os.path.join(excel_dir, src['file'])
        _, _, data = read_sheet(fpath, src['sheet'], primary_key, cfg)
        print(f"    读取: {src['file']} > {src['sheet']} → {len(data):,} 行")

    print(f"    合并去重后: {len(merged):,} 行, {len(headers)} 列 (主键列: {actual_pk})")

    if dry_run:
        print(f"    [DRY-RUN] 建表 SQL:")
        type_rules = cfg['type_rules']
        col_defs = ['[id] INT IDENTITY(1,1)']
        for h in headers:
            sql_type = infer_sql_type(h, type_rules, actual_pk)
            col_defs.append(f'  [{h}] {sql_type}')
        col_defs.append(f'  CONSTRAINT [PK_{table_name}] PRIMARY KEY ([{actual_pk}])')
        print(f"    CREATE TABLE [{cfg['schema']}].[{table_name}] (")
        print(',\n'.join(col_defs))
        print("    );")
        return len(merged)

    # 2. 建表
    create_sql = create_table(cursor, table_name, headers, cfg, actual_pk)
    print(f"    建表: [{cfg['schema']}].[{table_name}] ({len(headers)} 列)")

    # 3. 插入数据
    total = batch_insert(cursor, table_name, headers, merged, cfg, actual_pk)
    print(f"    插入: {total:,} 行 done")

    return total


def main():
    parser = argparse.ArgumentParser(description='Excel to SQL Server Migration')
    parser.add_argument('--config', default=None, help='配置文件路径 (默认: config.yaml)')
    parser.add_argument('--dry-run', action='store_true', help='仅预览，不写入数据库')
    parser.add_argument('--table', nargs='+', default=None,
                        help='只迁移指定的表，例如 --table PG_2022 TM_2025')
    args = parser.parse_args()

    # 加载配置
    cfg = load_config(args.config)

    # 按参数筛选表
    if args.table:
        specified = set(args.table)
        tables = [t for t in cfg['tables'] if t['table'] in specified]
        missing = specified - {t['table'] for t in tables}
        if missing:
            print(f"  错误: 未找到表 {', '.join(missing)}")
            available = ', '.join(t['table'] for t in cfg['tables'])
            print(f"  可用表: {available}")
            sys.exit(1)
    else:
        tables = cfg['tables']

    total_tables = len(tables)
    primary_key = cfg['primary_key']

    print(f"{'=' * 60}")
    print(f"  Excel → SQL Server 迁移工具")
    print(f"  目标: [{cfg['schema']}], {total_tables} 张表")
    print(f"  模式: {'DRY-RUN' if args.dry_run else 'LIVE'}")
    print(f"{'=' * 60}")

    conn = None
    cursor = None

    if not args.dry_run:
        conn = get_connection(cfg)
        cursor = conn.cursor()
        create_schema(cursor, cfg['schema'])
        print(f"\n  Schema [{cfg['schema']}] 已就绪")

    total_rows = 0
    for i, table_cfg in enumerate(tables, 1):
        print(f"\n[{i}/{total_tables}]", end='')
        rows = migrate_table(cursor, table_cfg, cfg, dry_run=args.dry_run)
        total_rows += rows

    if not args.dry_run:
        cursor.close()
        conn.close()

    print(f"\n{'=' * 60}")
    print(f"  完成! 共迁移 {total_tables} 张表, {total_rows:,} 行")
    print(f"{'=' * 60}")


if __name__ == '__main__':
    main()