#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
OSS Inventory 元数据导入脚本
读取 OSS Bucket 清单（inventory）的 manifest.json，
下载并解析 CSV 压缩文件，将文件元数据批量写入 Tablestore。

用法:
  python3 oss_inventory_import.py \
    --manifest oss://metadata-scan-test/metadata-scan-test/inventory-test/2026-03-12T09-46Z/manifest.json \
    --config config.json

config.json 格式:
  {
    "oss_access_key_id": "...",
    "oss_access_key_secret": "...",
    "oss_region": "cn-hangzhou",
    "ots_access_key_id": "...",
    "ots_access_key_secret": "...",
    "ots_endpoint": "https://<instance>.<region>.ots.aliyuncs.com",
    "ots_instance_name": "...",
    "ots_table_name": "oss_inventory"
  }

也支持通过环境变量配置（优先级高于 config.json）:
  OSS_ACCESS_KEY_ID, OSS_ACCESS_KEY_SECRET, OSS_REGION,
  OTS_ACCESS_KEY_ID, OTS_ACCESS_KEY_SECRET, OTS_ENDPOINT,
  OTS_INSTANCE_NAME, OTS_TABLE_NAME
"""

import argparse
import csv
import gzip
import io
import json
import logging
import os
import sys
import tempfile
from datetime import datetime, timezone
from urllib.parse import unquote_plus

import alibabacloud_oss_v2 as oss
from tablestore import OTSClient, Row, PutRowItem, Condition, RowExistenceExpectation, BatchWriteRowRequest, TableInBatchWriteRowItem

# ============================================================
# 常量
# ============================================================
BATCH_SIZE = 200  # OTS 批量写入每批最大行数

# ============================================================
# 日志
# ============================================================
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s %(levelname)s: %(message)s',
    datefmt='%Y-%m-%d %H:%M:%S'
)
logger = logging.getLogger('oss_inventory_import')


# ============================================================
# 配置加载
# ============================================================
def load_config(config_file=None):
    """从 config.json 和环境变量加载配置，环境变量优先"""
    cfg = {}
    if config_file and os.path.exists(config_file):
        with open(config_file, 'r') as f:
            cfg = json.load(f)

    def get(env_key, cfg_key, required=True):
        val = os.environ.get(env_key) or cfg.get(cfg_key)
        if required and not val:
            raise ValueError(f"缺少配置项: {env_key} / {cfg_key}")
        return val

    return {
        'oss_access_key_id':     get('OSS_ACCESS_KEY_ID', 'oss_access_key_id'),
        'oss_access_key_secret': get('OSS_ACCESS_KEY_SECRET', 'oss_access_key_secret'),
        'oss_region':            get('OSS_REGION', 'oss_region'),
        'ots_access_key_id':     get('OTS_ACCESS_KEY_ID', 'ots_access_key_id'),
        'ots_access_key_secret': get('OTS_ACCESS_KEY_SECRET', 'ots_access_key_secret'),
        'ots_endpoint':          get('OTS_ENDPOINT', 'ots_endpoint'),
        'ots_instance_name':     get('OTS_INSTANCE_NAME', 'ots_instance_name'),
        'ots_table_name':        get('OTS_TABLE_NAME', 'ots_table_name', required=False) or 'oss_inventory',
    }


# ============================================================
# OSS 客户端
# ============================================================
def make_oss_client(cfg):
    oss_cfg = oss.config.load_default()
    oss_cfg.credentials_provider = oss.credentials.StaticCredentialsProvider(
        cfg['oss_access_key_id'],
        cfg['oss_access_key_secret']
    )
    oss_cfg.region = cfg['oss_region']
    oss_cfg.endpoint = f"https://oss-{cfg['oss_region']}.aliyuncs.com"
    return oss.Client(oss_cfg)


def parse_oss_uri(uri):
    """解析 oss://bucket/key 格式的 URI"""
    if not uri.startswith('oss://'):
        raise ValueError(f"无效的 OSS URI: {uri}，格式应为 oss://bucket/key")
    without_scheme = uri[6:]
    slash = without_scheme.index('/')
    bucket = without_scheme[:slash]
    key = without_scheme[slash + 1:]
    return bucket, key


def download_object(oss_client, bucket, key):
    """从 OSS 下载对象，返回 bytes"""
    logger.info(f"下载: oss://{bucket}/{key}")
    result = oss_client.get_object(oss.GetObjectRequest(bucket=bucket, key=key))
    return result.body.read()


# ============================================================
# 数据解析工具
# ============================================================
def parse_last_modified(date_str):
    """
    将 OSS inventory 的时间格式 "2026-03-12T07-34-24Z" 转换为 epoch 毫秒。
    注意：时间部分用 - 而不是 : 分隔。
    """
    try:
        # 分离日期和时间
        date_part, time_part = date_str.split('T')
        # 去掉末尾 Z，将时间部分的 - 替换为 :
        time_part = time_part.rstrip('Z').replace('-', ':')
        iso_str = f"{date_part}T{time_part}+00:00"
        dt = datetime.fromisoformat(iso_str)
        return int(dt.timestamp() * 1000)
    except Exception:
        return 0


def parse_bool(val):
    """将 "true"/"false" 字符串转换为 bool"""
    return val.strip().lower() == 'true'


def get_prefix(key):
    """提取对象 key 的目录前缀（最后一个 / 之前的部分）"""
    idx = key.rfind('/')
    if idx == -1:
        return ''
    return key[:idx + 1]


def parse_csv_row(row_dict):
    """
    将 CSV 行（按 fileSchema 解析的 dict）转换为 OTS 属性列列表。
    返回 (pk_value, attribute_columns)

    Bucket 和 Key 是 inventory 的必选字段，其余字段按 fileSchema 实际存在列写入，
    不在 fileSchema 中的列不会写入 OTS（不写空值占位）。
    """
    # 解码 URL 编码的 key（如 photograph%2Ftest.jpg → photograph/test.jpg）
    decoded_key = unquote_plus(row_dict['Key'])

    # 必选列
    attribute_columns = [
        ('bucket', row_dict['Bucket']),
        ('prefix', get_prefix(decoded_key)),
    ]

    # 可选列：仅当 fileSchema 中包含时才写入
    if 'Crc64' in row_dict:
        attribute_columns.append(('crc64', row_dict['Crc64']))
    if 'ObjectType' in row_dict:
        attribute_columns.append(('object_type', row_dict['ObjectType']))
    if 'ObjectAcl' in row_dict:
        attribute_columns.append(('object_acl', row_dict['ObjectAcl']))
    if 'EncryptionStatus' in row_dict:
        attribute_columns.append(('encryption_status', parse_bool(row_dict['EncryptionStatus'])))
    if 'IsMultipartUploaded' in row_dict:
        attribute_columns.append(('is_multipart_uploaded', parse_bool(row_dict['IsMultipartUploaded'])))
    if 'ETag' in row_dict:
        attribute_columns.append(('etag', row_dict['ETag']))
    if 'LastModifiedDate' in row_dict:
        attribute_columns.append(('last_modified_ms', parse_last_modified(row_dict['LastModifiedDate'])))
    if 'StorageClass' in row_dict:
        attribute_columns.append(('storage_class', row_dict['StorageClass']))
    if 'Size' in row_dict:
        attribute_columns.append(('size', int(row_dict['Size'])))

    return decoded_key, attribute_columns


# ============================================================
# OTS 写入
# ============================================================
def make_ots_client(cfg):
    return OTSClient(
        cfg['ots_endpoint'],
        cfg['ots_access_key_id'],
        cfg['ots_access_key_secret'],
        cfg['ots_instance_name']
    )


def batch_write(ots_client, table_name, rows):
    """
    批量写入 OTS，rows 为 (pk_value, attribute_columns) 列表。
    自动分批（每批最多 BATCH_SIZE 行），失败行打印警告后继续。
    返回成功写入行数。
    """
    success_count = 0
    for i in range(0, len(rows), BATCH_SIZE):
        batch = rows[i:i + BATCH_SIZE]
        put_items = []
        for pk_value, attrs in batch:
            primary_key = [('key', pk_value)]
            condition = Condition(RowExistenceExpectation.IGNORE)
            put_items.append(PutRowItem(Row(primary_key, attrs), condition))

        request = BatchWriteRowRequest()
        request.add(TableInBatchWriteRowItem(table_name, put_items))

        try:
            response = ots_client.batch_write_row(request)
            failed = response.get_failed_of_put()
            if failed:
                for item in failed:
                    logger.warning(f"写入失败: {item.error_code} - {item.error_message}")
                success_count += len(batch) - len(failed)
            else:
                success_count += len(batch)
        except Exception as e:
            logger.error(f"批量写入异常: {e}")

    return success_count


# ============================================================
# 核心导入逻辑
# ============================================================
def import_from_manifest(manifest_bucket, manifest_key, cfg):
    """
    主入口：读取 manifest.json，下载并解析所有 CSV 文件，写入 OTS。
    返回写入的总行数。
    """
    oss_client = make_oss_client(cfg)
    ots_client = make_ots_client(cfg)
    table_name = cfg['ots_table_name']

    # 1. 读取 manifest.json
    manifest_data = download_object(oss_client, manifest_bucket, manifest_key)
    manifest = json.loads(manifest_data.decode('utf-8'))

    file_schema_str = manifest.get('fileSchema', '')
    columns = [c.strip() for c in file_schema_str.split(',')]
    logger.info(f"fileSchema: {columns}")

    csv_files = manifest.get('files', [])
    logger.info(f"共 {len(csv_files)} 个 CSV 文件")

    total_written = 0

    # 2. 逐个处理 CSV 文件
    for file_info in csv_files:
        csv_key = file_info['key']
        logger.info(f"处理: {csv_key}")

        gz_data = download_object(oss_client, manifest_bucket, csv_key)

        # 解压并按 fileSchema 列顺序解析
        rows = []
        with gzip.open(io.BytesIO(gz_data), 'rt', encoding='utf-8') as f:
            reader = csv.reader(f)
            for csv_row in reader:
                if len(csv_row) != len(columns):
                    logger.warning(f"列数不匹配，跳过: {csv_row}")
                    continue
                row_dict = dict(zip(columns, csv_row))
                pk_value, attrs = parse_csv_row(row_dict)
                rows.append((pk_value, attrs))

        logger.info(f"解析 {len(rows)} 行，开始写入 OTS...")
        written = batch_write(ots_client, table_name, rows)
        logger.info(f"写入完成: {written}/{len(rows)} 行")
        total_written += written

    return total_written


# ============================================================
# CLI 入口
# ============================================================
def main():
    parser = argparse.ArgumentParser(description='OSS Inventory 元数据导入 Tablestore')
    parser.add_argument(
        '--manifest', required=True,
        help='manifest.json 的 OSS URI，格式: oss://bucket/path/to/manifest.json'
    )
    parser.add_argument(
        '--config', default='config.json',
        help='配置文件路径（默认: config.json）'
    )
    args = parser.parse_args()

    try:
        cfg = load_config(args.config)
    except ValueError as e:
        logger.error(f"配置错误: {e}")
        sys.exit(1)

    manifest_bucket, manifest_key = parse_oss_uri(args.manifest)
    logger.info(f"Manifest: oss://{manifest_bucket}/{manifest_key}")
    logger.info(f"目标表: {cfg['ots_endpoint']} / {cfg['ots_table_name']}")

    total = import_from_manifest(manifest_bucket, manifest_key, cfg)
    logger.info(f"全部完成，共写入 {total} 行")


if __name__ == '__main__':
    main()
