from elasticsearch import AsyncElasticsearch
from packages.core.config import settings

def get_es_client() -> AsyncElasticsearch:
    """
    获取 Elasticsearch 异步客户端。
    """
    return AsyncElasticsearch(
        settings.ELASTICSEARCH_URL,
        # 如果没有配置证书，可以关闭验证（本地开发时）
        verify_certs=False
    )

async def init_es_schema():
    """
    初始化 ES 索引和映射
    建立商品描述、企业名称的高效索引
    """
    es = get_es_client()
    
    index_name = "trade_entities"
    
    mapping = {
        "mappings": {
            "properties": {
                "record_id": {"type": "keyword"},
                "source_country": {"type": "keyword"},
                "trade_direction": {"type": "keyword"},
                "trade_date": {"type": "date"},
                
                # 企业名称使用 text 配合 keyword 子字段
                "importer_name": {
                    "type": "text",
                    "fields": {
                        "keyword": {"type": "keyword", "ignore_above": 256}
                    }
                },
                "exporter_name": {
                    "type": "text",
                    "fields": {
                        "keyword": {"type": "keyword", "ignore_above": 256}
                    }
                },
                
                # 商品描述使用 text 用于全文检索
                "product_name": {
                    "type": "text",
                    "analyzer": "standard" # 后续可换成分词器如 ik_max_word
                },
                "hs_code": {"type": "keyword"}
            }
        }
    }
    
    try:
        exists = await es.indices.exists(index=index_name)
        if not exists:
            await es.indices.create(index=index_name, body=mapping)
            print(f"Elasticsearch index '{index_name}' created successfully.")
        else:
            print(f"Elasticsearch index '{index_name}' already exists.")
    except Exception as e:
        print(f"Failed to initialize ES schema: {e}")
    finally:
        await es.close()