import os
import requests
from fastapi import FastAPI
from pydantic import BaseModel
from typing import List, Optional, Dict
from concurrent.futures import ThreadPoolExecutor
import argparse
import uvicorn
from dotenv import load_dotenv, find_dotenv
import re
from Tea.exceptions import TeaException
from alibabacloud_iqs20241111 import models
from alibabacloud_iqs20241111.client import Client
from alibabacloud_tea_openapi import models as open_api_models
from typing import List, Dict, Any

parser = argparse.ArgumentParser(description="Launch online search server.")

parser.add_argument('--topk', type=int, default=3, 
                    help="Number of results to return per query")
args = parser.parse_args()

def quark_search(query) -> List[Dict[str, Any]]:
    """
    使用阿里云 IQS 接口进行统一搜索，返回搜索结果列表。
    :param query: 搜索关键词，例如"杭州美食"
    :return: 包含搜索结果的字典列表
    """
    def create_client() -> Client:
        load_dotenv(find_dotenv())
        config = open_api_models.Config(
            # 推荐使用环境变量方式管理 AK/SK
            access_key_id="xxx",
            access_key_secret="xxx"
        )
        config.endpoint = 'iqs.cn-zhangjiakou.aliyuncs.com'
        return Client(config)

    client = create_client()
    run_instances_request = models.UnifiedSearchRequest(
        body=models.UnifiedSearchInput(
            query=query,
            time_range='NoLimit',
            contents=models.RequestContents(
                summary=True,
                main_text=True,
            )
        )
    )

    try:
        response = client.unified_search(run_instances_request)
        # logger.info(
        #     "API 成功, Request ID: %s, 结果数量: %d, 搜索耗时: %s",
        #     response.body.request_id,
        #     len(response.body.page_items),
        #     response.body.search_information.search_time
        # )
        results = []
        # 遍历搜索结果
        for item in response.body.page_items:
            # 处理每个搜索结果
            results.append(item.__dict__)
        return results

    except TeaException as e:
        print(e)
        return []

# --- Config ---
class OnlineSearchConfig:
    def __init__(
        self,
        topk: int = 3
    ):
        self.topk = topk


# --- Online Search Wrapper ---
class OnlineSearchEngine:
    def __init__(self, config: OnlineSearchConfig):
        self.config = config

    def _search_query(self, query: str):
        response = quark_search(query)
        return response

    def batch_search(self, queries: List[str]):
        results = []
        with ThreadPoolExecutor() as executor:
            for result in executor.map(self._search_query, queries):
                results.append(self._process_result(result))
        return results

    def _process_result(self, search_result):
        results = []
        search_result = search_result[:self.config.topk]
        for i, result in enumerate(search_result):
            doc_content = ""
            # doc_content += f"[网页{i+1}]\n"
            doc_content += f"标题：{result['title']}\n"
            if result['main_text']!="" and result['main_text']!=None and result['main_text']!='None':
                doc_content += f"内容：{result['main_text']}\n"
            else:
                if result['summary']!="":
                    doc_content += f"内容：{result['summary']}\n"
                elif result['snippet']!="":
                    doc_content += f"内容：{result['snippet']}\n"
                else:
                    doc_content += f"内容：无\n"
            doc_content += f"网页来源：{result['hostname']}\n\n"

            results.append({
                'document': {"contents": doc_content},
            })

        return results


# --- FastAPI Setup ---
app = FastAPI(title="Online Search Proxy Server")

class SearchRequest(BaseModel):
    queries: List[str]

config = OnlineSearchConfig(
    topk=3
)

engine = OnlineSearchEngine(config)

# --- Routes ---
@app.post("/retrieve")
def search_endpoint(request: SearchRequest):
    results = engine.batch_search(request.queries)
    return {"result": results}

## return {"result": List[List[{'document': {"id": xx, "content": "title" + \n + "content"}, 'score': xx}]]}

# if __name__ == "__main__":
    # 3) Launch the server. By default, it listens on http://127.0.0.1:8000
    # uvicorn.run(app, host="0.0.0.0", port=8000)

# print(quark_search("杭州美食"))