"""research_tasks 数据库操作（psycopg3）。"""

import json

import psycopg

from config import config


def get_connection():
    return psycopg.connect(config.DATABASE_URL)


def get_task(task_id: str) -> dict | None:
    with get_connection() as conn:
        with conn.cursor() as cur:
            cur.execute(
                "SELECT id, user_id, app_id, topic, config, status, progress, "
                "status_message, error_message, report, sources, tokens_used, created_at "
                "FROM research_tasks WHERE id = %s",
                (task_id,),
            )
            row = cur.fetchone()
            if not row:
                return None
            cols = [d[0] for d in cur.description]
            return dict(zip(cols, row))


def update_task_status(
    task_id: str,
    status: str,
    progress: int = None,
    status_message: str = None,
    error_message: str = None,
    report: str = None,
    sources: list = None,
    tokens_used: int = None,
):
    fields = ["status = %(status)s", "updated_at = NOW()"]
    params = {"task_id": task_id, "status": status}

    if progress is not None:
        fields.append("progress = %(progress)s")
        params["progress"] = progress
    if status_message is not None:
        fields.append("status_message = %(status_message)s")
        params["status_message"] = status_message
    if error_message is not None:
        fields.append("error_message = %(error_message)s")
        params["error_message"] = error_message
    if report is not None:
        fields.append("report = %(report)s")
        params["report"] = report
    if sources is not None:
        fields.append("sources = %(sources)s")
        params["sources"] = json.dumps(sources, ensure_ascii=False)
    if tokens_used is not None:
        fields.append("tokens_used = %(tokens_used)s")
        params["tokens_used"] = tokens_used

    if status == "planning":
        fields.append("started_at = COALESCE(started_at, NOW())")
    elif status in ("completed", "failed", "canceled"):
        fields.append("completed_at = NOW()")

    sql = f"UPDATE research_tasks SET {', '.join(fields)} WHERE id = %(task_id)s"
    with get_connection() as conn:
        with conn.cursor() as cur:
            cur.execute(sql, params)
        conn.commit()


def is_canceled(task_id: str) -> bool:
    task = get_task(task_id)
    return bool(task and task.get("status") == "canceled")