import time
from datetime import datetime
from loguru import logger
import requests
import json
import redis
from aidso_geo.models import spider_save_tos
import aidso_geo.utils.bh_utils as bh_utils
from concurrent.futures import (
    ThreadPoolExecutor,
    wait,
    FIRST_COMPLETED,
)
MAX_WORKERS = 1
QUEUE_KEY = "geo:task_commit:JIKE:list1"
_redis8_client = None
_redis8_pool = None
import threading
import random
REQUEST_INTERVAL = 0.6
request_rate_lock = threading.Lock()
next_request_time = 0.0


def wait_request_slot():
    global next_request_time

    with request_rate_lock:
        now = time.monotonic()
        request_time = max(now, next_request_time)
        next_request_time = request_time + REQUEST_INTERVAL

    wait_seconds = request_time - now
    if wait_seconds > 0:
        time.sleep(wait_seconds)


def request_ark(url, headers, payload, max_retries=5):
    response = None

    for attempt in range(max_retries):
        wait_request_slot()

        response = requests.post(
            url,
            headers=headers,
            data=payload,
            timeout=300,
        )

        if response.status_code != 429:
            response.raise_for_status()
            return response

        retry_after = response.headers.get("Retry-After")

        if retry_after:
            try:
                wait_seconds = float(retry_after)
            except ValueError:
                wait_seconds = min(2 ** attempt, 16)
        else:
            wait_seconds = min(2 ** attempt, 16)

        wait_seconds += random.uniform(0, 0.5)
        logger.warning(
            f"触发429，第{attempt + 1}次重试，"
            f"等待{wait_seconds:.2f}秒"
        )
        time.sleep(wait_seconds)

    response.raise_for_status()


def init_redis8():
    global _redis8_client, _redis8_pool

    if _redis8_client is not None:
        return _redis8_client

    try:
        _redis8_pool = redis.ConnectionPool(
            host='redis-cnlfmu7rl14awitrz.redis.ivolces.com',
            port=6379,
            db=8,
            password='aiyingli@@123',
            socket_timeout=120,
            socket_connect_timeout=5,
            decode_responses=True,
            retry_on_timeout=True,
            health_check_interval=30,
            socket_keepalive=True,
            max_connections=100,
        )

        _redis8_client = redis.Redis(connection_pool=_redis8_pool)
        _redis8_client.ping()
        return _redis8_client
    except Exception as e:
        logger.error(f"Redis 初始化失败: {e}")
        _redis8_client = None
        _redis8_pool = None
        return None


def get_result(prompt):
    url = "https://ark.cn-beijing.volces.com/api/v3/responses"

    payload = json.dumps({
       "model": "doubao-seed-2-1-pro-260628",
       "stream": False,
       "tools": [
          {
             "type": "doubao_app",
             "feature": {
                "ai_search": {
                   "type": "enabled",
                   "role_description": prompt
                }
             }
          }
       ],
       "input": [
          {
             "type": "message",
             "role": "user",
             "content": [
                {
                   "type": "input_text",
                   "text": prompt
                }
             ]
          }
       ]
    })
    headers = {
       'Authorization': 'Bearer ark-7afc3be2-37a8-47fd-9f02-996258a3d305-27da0',
       'ark-beta-doubao-app': 'true',
       'Content-Type': 'application/json'
    }

    response_text = request_ark(url, headers, payload)
    response_json = response_text.json()
    result = []
    for out in  response_json.get('output'):
        out_blocks = out.get('blocks')
        result.extend(out_blocks)
    return result


def get_think_result(prompt):

    url = "https://ark.cn-beijing.volces.com/api/v3/responses"

    payload = json.dumps({
        "model": "doubao-seed-2-1-pro-260628",
        "stream": False,
        "tools": [
            {
                "type": "doubao_app",
                "feature": {
                    "reasoning_search": {
                        "type": "enabled"
                    }
                }
            }
        ],
        "input": [
            {
                "type": "message",
                "role": "user",
                "content": [
                    {
                        "type": "input_text",
                        "text": prompt
                    }
                ]
            }
        ]
    })
    headers = {
        'Authorization': 'Bearer ark-7afc3be2-37a8-47fd-9f02-996258a3d305-27da0',
        'ark-beta-doubao-app': 'true',
        'Content-Type': 'application/json'
    }

    response_text = request_ark(url, headers, payload)
    response_json = response_text.json()
    result = []
    for out in  response_json.get('output'):
        out_blocks = out.get('blocks')
        result.extend(out_blocks)
    return result

def parse_think(think_result):
    think_content = ""
    response_content = ""
    search_word = []
    quto = []
    suggestions = []
    for i in think_result:
        if i.get('type') == 'reasoning_text':
            think_content += i.get('reasoning_text')
        if i.get('type') == 'reasoning_search':
            think_content += "\n\n"
            think_content += i.get('summary')
            think_content += "\n\n"
            search_word.extend(i.get('queries'))

            results = i.get('results')
            quto.extend(results)
        if i.get('type') == 'output_text':
            response_content += i.get('text')
    return (search_word, quto, think_content, response_content,suggestions)

def parse_result(think_result):
    think_content = ""
    response_content = ""
    search_word = []
    quto = []
    suggestions = []
    for i in think_result:
        if i.get('type') == 'search':
            think_content += i.get('summary')
            search_word.extend(i.get('queries'))

            results = i.get('results')
            quto.extend(results)
        if i.get('type') == 'output_text':
            response_content += i.get('text')
    return (search_word, quto, think_content, response_content,suggestions)


def process_task(tasks):
    task_json = json.loads(tasks)

    prompt = task_json.get("prompt")
    reqId = task_json.get('reqId')
    logger.success(f"{reqId}开始处理")
    file_path = (
        f'geo/{task_json["taskId"]}/'
        f'{task_json["platform"]}/original.text'
    )

    task_json["pt"] = datetime.now().strftime("%Y%m%d")
    task_json["source"] = "assistant"

    if str(task_json.get("thinkingEnabled")) == "1":
        (
            search_keyword,
            url_list,
            think_content,
            response_content,
            suggestions,
        ) = parse_think(
            get_think_result(prompt)
        )
    else:
        (
            search_keyword,
            url_list,
            think_content,
            response_content,
            suggestions,
        ) = parse_result(
            get_result(prompt)
        )
    if not response_content:
        raise RuntimeError("接口返回内容为空")
    spider_save_tos.process_and_save_files(
        file_path,
        search_keyword,
        url_list,
        think_content,
        response_content,
        suggestions,
    )

    task_json["status"] = "SUCCESS"
    logger.success(f"{reqId}处理完成")
    return task_json



if __name__ == "__main__":
    redis8 = init_redis8()
    future_task_map = {}

    with ThreadPoolExecutor(
        max_workers=MAX_WORKERS
    ) as executor:
        while True:
            while len(future_task_map) < MAX_WORKERS:
                original_task = redis8.rpop(QUEUE_KEY)

                if not original_task:
                    break

                future = executor.submit(
                    process_task,
                    original_task,
                )
                future_task_map[future] = original_task

            if not future_task_map:
                logger.success("当前无任务")
                time.sleep(1)
                continue

            done, _ = wait(
                tuple(future_task_map),
                timeout=1,
                return_when=FIRST_COMPLETED,
            )

            for future in done:
                original_task = future_task_map.pop(future)

                try:
                    success_task = future.result()

                    insert_ok = bh_utils.insert_data(
                        "geo_commit_task",
                        [success_task],
                    )

                    if not insert_ok:
                        raise RuntimeError("数据库写入失败")

                except Exception as error:
                    logger.exception(
                        f"任务处理失败，重新放回队列：{error}"
                    )
                    redis8.lpush(
                        QUEUE_KEY,
                        original_task,
                    )