"""任务消费者 — 从 Redis 队列取研究任务并执行流水线。仿 ppt-worker 模式。"""

import json
import time
import signal
from concurrent.futures import ThreadPoolExecutor

import redis

from config import config
from db import get_task, update_task_status, is_canceled
import llm_client
import search
from pipeline import ResearchPipeline, CanceledError


class ResearchWorker:
    def __init__(self):
        self.redis = redis.from_url(config.REDIS_URL, decode_responses=True)
        self.executor = ThreadPoolExecutor(max_workers=config.CONCURRENCY)
        self.running = True
        try:
            signal.signal(signal.SIGINT, self._shutdown)
            signal.signal(signal.SIGTERM, self._shutdown)
        except ValueError:
            pass

    def _shutdown(self, signum, frame):
        print(f"\n[Research] 收到信号 {signum}，正在优雅关闭…")
        self.running = False

    def start(self):
        print(f"[Research] 启动，并发数: {config.CONCURRENCY}，监听队列: {config.TASK_QUEUE}")
        while self.running:
            try:
                result = self.redis.brpop(config.TASK_QUEUE, timeout=5)
                if result is None:
                    continue
                _, raw = result
                task_id = json.loads(raw).get("task_id")
                if task_id:
                    self.executor.submit(self._process, task_id)
            except redis.ConnectionError as e:
                print(f"[Research] Redis 连接失败: {e}，5 秒后重试…")
                time.sleep(5)
            except Exception as e:
                print(f"[Research] 未知错误: {e}")
                time.sleep(1)
        self.executor.shutdown(wait=True)

    def _redis_status(self, task_id, status, progress, message):
        key = f"{config.TASK_STATUS_PREFIX}{task_id}"
        self.redis.hset(key, mapping={"status": status, "progress": str(progress), "message": message})
        self.redis.expire(key, 3600)

    def _process(self, task_id: str):
        try:
            task = get_task(task_id)
            if not task:
                print(f"[Research] 任务不存在: {task_id}")
                return
            if task["status"] != "pending":
                print(f"[Research] 跳过非 pending 任务: {task_id} ({task['status']})")
                return

            def progress_cb(status, progress, message):
                # 同时更新 Redis（快速轮询）与 DB（Go 端读取）
                self._redis_status(task_id, status, progress, message)
                update_task_status(task_id, status, progress=progress, status_message=message)

            pipeline = ResearchPipeline(
                task_id, task,
                llm_chat=llm_client.chat,
                search_fn=search.search,
                fetch_fn=search.fetch_text,
                update_fn=lambda **kw: update_task_status(task_id, **kw),
                progress_cb=progress_cb,
                is_canceled=lambda: is_canceled(task_id),
                max_subqueries=config.MAX_SUBQUERIES,
                max_sources=config.MAX_SOURCES,
            )
            pipeline.run()
            self._redis_status(task_id, "completed", 100, "完成")
            print(f"[Research] 任务完成: {task_id}")

        except CanceledError:
            update_task_status(task_id, "canceled", status_message="任务已取消")
            self._redis_status(task_id, "canceled", 0, "已取消")
            print(f"[Research] 任务取消: {task_id}")
        except Exception as e:
            update_task_status(task_id, "failed", error_message=str(e))
            self._redis_status(task_id, "failed", 0, f"失败: {str(e)[:200]}")
            print(f"[Research] 任务失败: {task_id} - {e}")


if __name__ == "__main__":
    ResearchWorker().start()