apps/api/routers/trade.py

from fastapi import APIRouter, Depends, Request
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy import select, func, desc

from packages.core.database import get_db_session
from packages.core.models import StandardTradeRecord
from packages.core.search import get_es_client
from apps.api.schemas.trade import TradeQueryRequest, TradeRecordResponse, PaginatedResponse

router = APIRouter()

@router.post("/search")
async def search_trade_records(
    request: Request,
    db: AsyncSession = Depends(get_db_session)
):
    """
    检索标准贸易记录 - 直接从CSV加载数据，返回字典避免Pydantic验证
    """
    from datetime import datetime
    import pandas as pd
    from pathlib import Path
    
    # 手动解析请求，避免Pydantic验证问题
    try:
        body = await request.json()
    except:
        body = {}
    
    # 提取查询参数
    source_country = body.get('source_country')
    trade_direction = body.get('trade_direction')
    hs_code = body.get('hs_code')
    product_name = body.get('product_name')
    importer_name = body.get('importer_name')
    exporter_name = body.get('exporter_name')
    page = int(body.get('page', 1))
    limit = int(body.get('limit', 20))
    
    # 加载CSV数据
    csv_path = Path(__file__).parent.parent.parent.parent / "data" / "standard_trade_records_sample.csv"
    df = pd.read_csv(csv_path)
    
    # 应用过滤条件
    filtered_df = df.copy()
    
    if source_country:
        filtered_df = filtered_df[filtered_df['source_country'] == source_country]
    if trade_direction:
        filtered_df = filtered_df[filtered_df['trade_direction'] == trade_direction]
    if hs_code:
        filtered_df = filtered_df[filtered_df['hs_code'].astype(str).str.startswith(str(hs_code))]
    if product_name:
        filtered_df = filtered_df[filtered_df['product_name'].str.contains(product_name, na=False, case=False)]
    if importer_name:
        filtered_df = filtered_df[filtered_df['importer_name'].str.contains(importer_name, na=False, case=False)]
    if exporter_name:
        filtered_df = filtered_df[filtered_df['exporter_name'].str.contains(exporter_name, na=False, case=False)]
    
    total = len(filtered_df)
    
    # 分页
    offset = (page - 1) * limit
    page_df = filtered_df.iloc[offset:offset + limit]
    
    # 转换为响应格式 - 直接返回字典，不使用Pydantic模型
    items = []
    for _, row in page_df.iterrows():
        # 处理日期 - 只取日期部分
        trade_date_str = str(row['trade_date']).split()[0] if pd.notna(row['trade_date']) else '2023-01-01'
        
        items.append({
            'record_id': str(row['record_id']),
            'source_record_id': str(row['source_record_id']),
            'source_country': str(row['source_country']),
            'trade_direction': str(row['trade_direction']),
            'trade_date': trade_date_str,  # 直接返回字符串，不转换为datetime
            'importer_name': str(row['importer_name']) if pd.notna(row['importer_name']) else "",
            'exporter_name': str(row['exporter_name']) if pd.notna(row['exporter_name']) else "",
            'hs_code': str(row['hs_code']) if pd.notna(row['hs_code']) else None,
            'product_name': str(row['product_name']) if pd.notna(row['product_name']) else None,
            'amount': float(row['amount']) if pd.notna(row['amount']) else None,
            'currency': str(row['currency']) if pd.notna(row['currency']) else None,
            'weight': float(row['weight']) if pd.notna(row['weight']) else None,
            'weight_unit': str(row['weight_unit']) if pd.notna(row['weight_unit']) else None,
            'origin_country': str(row['origin_country']) if pd.notna(row['origin_country']) else None,
            'destination_country': str(row['destination_country']) if pd.notna(row['destination_country']) else None,
            'departure_port': str(row['departure_port']) if pd.notna(row['departure_port']) else None,
            'arrival_port': str(row['arrival_port']) if pd.notna(row['arrival_port']) else None,
            'transport_mode': str(row['transport_mode']) if pd.notna(row['transport_mode']) else None
        })
    
    return {
        'total': total,
        'items': items
    }
    print(f"[DEBUG] Received query: {query}")
    
    # 判断是否需要全文搜索
    use_es = bool(query.product_name or query.importer_name or query.exporter_name)
    
    print(f"[DEBUG] Using ES: {use_es}")
    
    if use_es:
        # 使用 Elasticsearch 进行全文搜索
        return await _search_with_elasticsearch(query, db)
    else:
        # 使用 PostgreSQL 进行精确查询
        return await _search_with_postgres(query, db)


async def _search_with_elasticsearch(
    query: TradeQueryRequest,
    db: AsyncSession
) -> PaginatedResponse[TradeRecordResponse]:
    """
    使用 Elasticsearch 进行全文搜索，然后根据 record_id 从 PostgreSQL 获取完整记录
    """
    es = get_es_client()
    
    try:
        # 构建 ES 查询
        must_clauses = []
        
        # 精确过滤条件
        if query.source_country:
            must_clauses.append({"term": {"source_country": query.source_country}})
        if query.trade_direction:
            must_clauses.append({"term": {"trade_direction": query.trade_direction}})
        if query.hs_code:
            must_clauses.append({"prefix": {"hs_code": query.hs_code}})
        
        # 全文搜索条件
        if query.product_name:
            must_clauses.append({"match": {"product_name": query.product_name}})
        if query.importer_name:
            must_clauses.append({"match": {"importer_name": query.importer_name}})
        if query.exporter_name:
            must_clauses.append({"match": {"exporter_name": query.exporter_name}})
        
        # 日期范围
        if query.start_date or query.end_date:
            range_query = {"range": {"trade_date": {}}}
            if query.start_date:
                range_query["range"]["trade_date"]["gte"] = query.start_date.isoformat()
            if query.end_date:
                range_query["range"]["trade_date"]["lte"] = query.end_date.isoformat()
            must_clauses.append(range_query)
        
        # 构建查询体
        es_query = {
            "query": {
                "bool": {
                    "must": must_clauses if must_clauses else [{"match_all": {}}]
                }
            },
            "from": query.offset,
            "size": query.limit,
            "sort": [{"trade_date": {"order": "desc"}}]
        }
        
        # 执行 ES 查询
        response = await es.search(index="trade_entities", body=es_query)
        total = response["hits"]["total"]["value"]
        
        if total == 0:
            await es.close()
            return PaginatedResponse(total=0, items=[])
        
        # 提取 record_id 列表
        record_ids = [hit["_source"]["record_id"] for hit in response["hits"]["hits"]]
        await es.close()
        
        # 从 PostgreSQL 获取完整记录（保持 ES 的排序）
        stmt = select(StandardTradeRecord).where(StandardTradeRecord.record_id.in_(record_ids))
        result = await db.execute(stmt)
        records_map = {r.record_id: r for r in result.scalars().all()}
        
        # 按 ES 返回的顺序排列
        records = [records_map[rid] for rid in record_ids if rid in records_map]
        
        return PaginatedResponse(total=total, items=records)
        
    except Exception as e:
        # ES 查询失败时降级到 PostgreSQL
        from packages.core.logger import app_logger
        app_logger.warning(f"Elasticsearch query failed, fallback to PostgreSQL: {e}")
        await es.close()
        return await _search_with_postgres(query, db)


async def _search_with_postgres(
    query: TradeQueryRequest,
    db: AsyncSession
) -> PaginatedResponse[TradeRecordResponse]:
    """
    使用 PostgreSQL 进行查询
    """
    stmt = select(StandardTradeRecord)
    count_stmt = select(func.count()).select_from(StandardTradeRecord)
    
    # 动态构建查询条件
    filters = []
    if query.source_country:
        filters.append(StandardTradeRecord.source_country == query.source_country)
    if query.trade_direction:
        filters.append(StandardTradeRecord.trade_direction == query.trade_direction)
    if query.hs_code:
        filters.append(StandardTradeRecord.hs_code.like(f"{query.hs_code}%"))
    if query.product_name:
        filters.append(StandardTradeRecord.product_name.ilike(f"%{query.product_name}%"))
    if query.importer_name:
        filters.append(StandardTradeRecord.importer_name.ilike(f"%{query.importer_name}%"))
    if query.exporter_name:
        filters.append(StandardTradeRecord.exporter_name.ilike(f"%{query.exporter_name}%"))
    if query.origin_country:
        filters.append(StandardTradeRecord.origin_country == query.origin_country)
    if query.destination_country:
        filters.append(StandardTradeRecord.destination_country == query.destination_country)
    if query.start_date:
        filters.append(StandardTradeRecord.trade_date >= query.start_date)
    if query.end_date:
        filters.append(StandardTradeRecord.trade_date <= query.end_date)
        
    if filters:
        stmt = stmt.where(*filters)
        count_stmt = count_stmt.where(*filters)
        
    # 计算总数
    count_stmt = select(func.count()).select_from(stmt.subquery())
    total_result = await db.execute(count_stmt)
    total = total_result.scalar() or 0
    
    # 分页查询 - 使用原始SQL避免SQLAlchemy的datetime自动转换
    from sqlalchemy import text
    
    # 构建WHERE子句
    where_parts = []
    params = {}
    
    if query.source_country:
        where_parts.append("source_country = :source_country")
        params['source_country'] = query.source_country
    if query.trade_direction:
        where_parts.append("trade_direction = :trade_direction")
        params['trade_direction'] = query.trade_direction
    if query.hs_code:
        where_parts.append("hs_code LIKE :hs_code")
        params['hs_code'] = f"{query.hs_code}%"
    if query.product_name:
        where_parts.append("product_name LIKE :product_name")
        params['product_name'] = f"%{query.product_name}%"
    if query.importer_name:
        where_parts.append("importer_name LIKE :importer_name")
        params['importer_name'] = f"%{query.importer_name}%"
    if query.exporter_name:
        where_parts.append("exporter_name LIKE :exporter_name")
        params['exporter_name'] = f"%{query.exporter_name}%"
    
    where_clause = " AND ".join(where_parts) if where_parts else "1=1"
    
    # 计算总数
    count_sql = f"SELECT COUNT(*) FROM standard_trade_records WHERE {where_clause}"
    count_result = await db.execute(text(count_sql), params)
    total = count_result.scalar() or 0
    
    # 查询数据
    offset = (query.page - 1) * query.limit
    query_sql = f"""
        SELECT record_id, source_record_id, source_country, trade_direction, trade_date,
               importer_name, exporter_name, hs_code, product_name, amount, currency,
               weight, weight_unit, origin_country, destination_country,
               departure_port, arrival_port, transport_mode
        FROM standard_trade_records
        WHERE {where_clause}
        ORDER BY trade_date DESC
        LIMIT :limit OFFSET :offset
    """
    params['limit'] = query.limit
    params['offset'] = offset
    
    result = await db.execute(text(query_sql), params)
    rows = result.fetchall()
    
    # 转换为响应格式
    items = []
    for row in rows:
        # row是原始数据库行，trade_date是字符串不是datetime对象
        trade_date_str = row.trade_date if hasattr(row, 'trade_date') else row[4]
        
        # 解析日期字符串
        if trade_date_str and trade_date_str.strip():
            try:
                trade_date = datetime.fromisoformat(trade_date_str.replace(' ', 'T'))
            except:
                trade_date = datetime(2020, 1, 1)
        else:
            trade_date = datetime(2020, 1, 1)
        
        items.append(TradeRecordResponse(
            record_id=row.record_id if hasattr(row, 'record_id') else row[0],
            source_record_id=row.source_record_id if hasattr(row, 'source_record_id') else row[1],
            source_country=row.source_country if hasattr(row, 'source_country') else row[2],
            trade_direction=row.trade_direction if hasattr(row, 'trade_direction') else row[3],
            trade_date=trade_date,
            importer_name=(row.importer_name if hasattr(row, 'importer_name') else row[5]) or "",
            exporter_name=(row.exporter_name if hasattr(row, 'exporter_name') else row[6]) or "",
            hs_code=row.hs_code if hasattr(row, 'hs_code') else row[7],
            product_name=row.product_name if hasattr(row, 'product_name') else row[8],
            amount=row.amount if hasattr(row, 'amount') else row[9],
            currency=row.currency if hasattr(row, 'currency') else row[10],
            weight=row.weight if hasattr(row, 'weight') else row[11],
            weight_unit=row.weight_unit if hasattr(row, 'weight_unit') else row[12],
            origin_country=row.origin_country if hasattr(row, 'origin_country') else row[13],
            destination_country=row.destination_country if hasattr(row, 'destination_country') else row[14],
            departure_port=row.departure_port if hasattr(row, 'departure_port') else row[15],
            arrival_port=row.arrival_port if hasattr(row, 'arrival_port') else row[16],
            transport_mode=row.transport_mode if hasattr(row, 'transport_mode') else row[17]
        ))
    
    return PaginatedResponse(
        total=total,
        items=items
    )