from fastapi import FastAPI, Query
import uvicorn
from typing import List, Dict, Any  # 导入类型提示，让代码更清晰
from search_engine import SearchEngine
from tqdm import tqdm
import os
import argparse

# 默认配置
DEFAULT_DATASET_DIR = './search_engine/corpus'
DEFAULT_EMBED_MODEL_NAME = 'vidore/colqwen2-v1.0'
DEFAULT_NODE_DIR_PREFIX = 'colqwen_ingestion'
DEFAULT_HOST = "0.0.0.0"
DEFAULT_PORT = 8002

# 全局变量
dataset_dir = DEFAULT_DATASET_DIR
embed_model_name = DEFAULT_EMBED_MODEL_NAME
node_dir_prefix = DEFAULT_NODE_DIR_PREFIX
img_test_dir = None

# 创建 FastAPI 应用实例
app = FastAPI(
    title="Hybrid Search Engine API",
    description="Provides search functionality using a pre-initialized HybridSearchEngine.",
    version="1.0.0",
)

# 全局变量，用于存储 HybridSearchEngine 实例
search_engine = None

# 在应用启动时初始化 HybridSearchEngine
@app.on_event("startup")
async def startup_event():
    global search_engine
    print(f"Initializing SearchEngine with dataset_dir: {dataset_dir}, embed_model: {embed_model_name}, node_dir_prefix: {node_dir_prefix}")
    search_engine = SearchEngine(dataset_dir, node_dir_prefix=node_dir_prefix, embed_model_name=embed_model_name)

# 定义搜索 API 端点
@app.get(
    "/search",
    summary="Perform a search query.",
    description="Executes a search using the initialized SearchEngine and returns the results.",
    response_model=List[List[Dict[str, Any]]]  # 定义响应模型，提高 API 文档的清晰度
)
async def search(queries: List[str] = Query(...)):
    """
    执行搜索操作。

    Args:
        query: 搜索查询字符串。

    Returns:
        搜索结果列表。
    """
    results_batch = search_engine.batch_search(queries)
    # 使用动态的图片目录路径
    img_dir = img_test_dir or os.path.join(dataset_dir, 'img')
    results_batch = [[dict(idx=idx, image_file=os.path.join(img_dir, file)) for idx, file in enumerate(query_results)] for query_results in results_batch]
    return results_batch

def parse_arguments():
    """解析命令行参数"""
    parser = argparse.ArgumentParser(description='启动搜索引擎API服务')
    parser.add_argument('--dataset-dir', type=str, default=DEFAULT_DATASET_DIR,
                       help=f'数据集目录路径 (默认: {DEFAULT_DATASET_DIR})')
    parser.add_argument('--embed-model', type=str, default=DEFAULT_EMBED_MODEL_NAME,
                       help=f'嵌入模型名称 (默认: {DEFAULT_EMBED_MODEL_NAME})')
    parser.add_argument('--node-dir-prefix', type=str, default=DEFAULT_NODE_DIR_PREFIX,
                       help=f'节点目录前缀 (默认: {DEFAULT_NODE_DIR_PREFIX})')
    parser.add_argument('--img-test-dir', type=str, default=None,
                       help='图片测试目录路径 (默认: 使用 dataset_dir/img-test)')
    parser.add_argument('--host', type=str, default=DEFAULT_HOST,
                       help=f'服务器主机地址 (默认: {DEFAULT_HOST})')
    parser.add_argument('--port', type=int, default=DEFAULT_PORT,
                       help=f'服务器端口 (默认: {DEFAULT_PORT})')
    return parser.parse_args()

if __name__ == "__main__":
    # 解析命令行参数
    args = parse_arguments()
    
    # 设置全局变量
    dataset_dir = args.dataset_dir
    embed_model_name = args.embed_model
    node_dir_prefix = args.node_dir_prefix
    img_test_dir = args.img_test_dir
    
    print(f"启动搜索引擎API服务:")
    print(f"  数据集目录: {dataset_dir}")
    print(f"  嵌入模型: {embed_model_name}")
    print(f"  节点目录前缀: {node_dir_prefix}")
    print(f"  图片测试目录: {img_test_dir or os.path.join(dataset_dir, 'img')}")
    print(f"  服务器地址: {args.host}:{args.port}")
    
    uvicorn.run(app, host=args.host, port=args.port)