#!/usr/bin/env python3
"""
从达梦数据库导出三元组数据到 JSON 文件
三元组来源：RELATION_INSTANCE 表 (HEAD_ID, RELATION_TYPE, TAIL_ID)
支持可选清洗：去重、过滤悬空引用、移除空值、规范化关系类型
用法: python scripts/export_triplets.py
需设置环境变量 DM_PASSWORD，且 config 中 source=db
"""
import json
import os
import sys

_script_dir = os.path.dirname(os.path.abspath(__file__))
_project_root = os.path.dirname(_script_dir)
sys.path.insert(0, _project_root)


def _json_serializer(obj):
    """处理 Timestamp、datetime、numpy 等不可直接 JSON 序列化的类型"""
    if hasattr(obj, 'isoformat'):
        return obj.isoformat()
    if hasattr(obj, 'item'):
        return obj.item()
    if isinstance(obj, float) and (obj != obj or obj == float('inf') or obj == float('-inf')):
        return None
    raise TypeError(f'Object of type {type(obj).__name__} is not JSON serializable')


def load_config(project_root: str) -> dict:
    """加载 config"""
    config_dir = os.path.join(project_root, 'config')
    json_path = os.path.join(config_dir, 'config.json')
    yaml_path = os.path.join(config_dir, 'config.yaml')
    if os.path.exists(json_path):
        with open(json_path, 'r', encoding='utf-8') as f:
            return json.load(f)
    if os.path.exists(yaml_path):
        try:
            import yaml
            with open(yaml_path, 'r', encoding='utf-8') as f:
                return yaml.safe_load(f)
        except ImportError:
            pass
    return {'source': 'json', 'db': {}, 'triplet': {}}


def main():
    from src.data_loader import load_triplets_from_dm, load_from_dm
    from src.triplet_cleaner import clean_triplets, transform_triplet_output_format

    config = load_config(_project_root)
    if config.get('source') != 'db':
        print("错误：三元组导出仅支持达梦模式，请将 config 中 source 改为 db")
        sys.exit(1)

    tc = config.get('triplet', {})
    output_file = tc.get('output_file', 'data/triples.json')
    report_file = tc.get('report_file', 'report/triplet_cleaning_report.json')
    output_path = output_file if os.path.isabs(output_file) else os.path.join(_project_root, output_file)
    report_path = report_file if os.path.isabs(report_file) else os.path.join(_project_root, report_file)
    os.makedirs(os.path.dirname(output_path), exist_ok=True)
    os.makedirs(os.path.dirname(report_path), exist_ok=True)

    # 清洗配置（clean: false 则跳过；否则按子项配置）
    clean_cfg = tc.get('clean', True)
    do_clean = False
    if clean_cfg is True:
        do_clean = {'deduplicate': True, 'filter_orphans': False, 'remove_null_core': True, 'normalize_relation_type': True}
    elif isinstance(clean_cfg, dict):
        do_clean = clean_cfg

    deduplicate = do_clean.get('deduplicate', True) if do_clean else False
    filter_orphans = do_clean.get('filter_orphans', False) if do_clean else False
    remove_null_core = do_clean.get('remove_null_core', True) if do_clean else False
    normalize_relation_type = do_clean.get('normalize_relation_type', True) if do_clean else False

    print("正在从达梦读取三元组 (RELATION_INSTANCE)...")
    triplets = load_triplets_from_dm(config, _project_root)
    print(f"读取到 {len(triplets)} 条三元组")

    valid_ids = None
    if filter_orphans:
        print("正在加载实体表以校验悬空引用...")
        entities = load_from_dm(config, _project_root)
        valid_ids = {str(e.get('TARGET_ID', '')).strip() for e in entities if e.get('TARGET_ID')}
        print(f"有效实体数: {len(valid_ids)}")

    if do_clean and (deduplicate or filter_orphans or remove_null_core or normalize_relation_type):
        print("正在执行三元组清洗...")
        triplets, report = clean_triplets(
            triplets,
            valid_ids=valid_ids,
            deduplicate=deduplicate,
            filter_orphans=filter_orphans,
            remove_null_core=remove_null_core,
            normalize_relation_type=normalize_relation_type,
        )
        with open(report_path, 'w', encoding='utf-8') as f:
            json.dump(report, f, ensure_ascii=False, indent=2, default=_json_serializer)
        print(f"清洗完成：{report['original_count']} -> {report['final_count']}，报告: {report_path}")

    # 转为输出格式：head, head_type, relation, tail, tail_type
    entities = load_from_dm(config, _project_root)
    id_to_role = {str(e.get('TARGET_ID', '')).strip(): str(e.get('ROLE_ID', '')).strip() for e in entities if e.get('TARGET_ID')}
    triplets = transform_triplet_output_format(triplets, id_to_role)

    with open(output_path, 'w', encoding='utf-8') as f:
        json.dump(triplets, f, ensure_ascii=False, indent=2, default=_json_serializer)
    print(f"已导出至 {output_path}")


if __name__ == "__main__":
    main()