from __future__ import annotations

import json
import uuid
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

import aiosqlite

from aicd.db.schema import SCHEMA


def _utc_now() -> str:
    return datetime.now(timezone.utc).isoformat()


@dataclass
class ChatMessageRow:
    id: int
    session_id: str
    role: str
    content: str
    tool_calls_json: str | None
    created_at: str
    tool_call_id: str | None
    tool_name: str | None


@dataclass
class ChatSummaryRow:
    id: int
    session_id: str
    content: str
    covers_up_to_message_id: int | None
    created_at: str


@dataclass
class TaskRow:
    id: str
    type: str
    status: str
    parent_id: str | None
    created_at: str
    updated_at: str
    payload: dict[str, Any]
    result_summary: str | None
    error: str | None


@dataclass
class TaskMessageRow:
    id: int
    thread_root_id: str | None
    to_task_id: str
    from_task_id: str
    kind: str
    body: dict[str, Any]
    created_at: str
    read_at: str | None


class TaskStore:
    def __init__(self, db_path: Path) -> None:
        self._db_path = db_path

    async def connect(self) -> aiosqlite.Connection:
        self._db_path.parent.mkdir(parents=True, exist_ok=True)
        conn = await aiosqlite.connect(self._db_path)
        await conn.executescript(SCHEMA)
        await self._migrate_chat_schema(conn)
        await conn.commit()
        return conn

    async def _migrate_chat_schema(self, conn: aiosqlite.Connection) -> None:
        cur = await conn.execute(
            "SELECT name FROM sqlite_master WHERE type='table' AND name='chat_messages'"
        )
        if not await cur.fetchone():
            return
        cur = await conn.execute("PRAGMA table_info(chat_messages)")
        cols = {str(r[1]) for r in await cur.fetchall()}
        alters: list[str] = []
        if "tool_call_id" not in cols:
            alters.append(
                "ALTER TABLE chat_messages ADD COLUMN tool_call_id TEXT"
            )
        if "tool_name" not in cols:
            alters.append("ALTER TABLE chat_messages ADD COLUMN tool_name TEXT")
        for sql in alters:
            await conn.execute(sql)

    async def insert_task(
        self,
        conn: aiosqlite.Connection,
        *,
        task_type: str,
        status: str,
        parent_id: str | None,
        payload: dict[str, Any],
        task_id: str | None = None,
    ) -> str:
        tid = task_id or str(uuid.uuid4())
        now = _utc_now()
        await conn.execute(
            """
            INSERT INTO tasks (id, type, status, parent_id, created_at, updated_at,
                payload_json, result_summary, error)
            VALUES (?, ?, ?, ?, ?, ?, ?, NULL, NULL)
            """,
            (
                tid,
                task_type,
                status,
                parent_id,
                now,
                now,
                json.dumps(payload, ensure_ascii=False),
            ),
        )
        await conn.commit()
        return tid

    async def update_task(
        self,
        conn: aiosqlite.Connection,
        task_id: str,
        *,
        status: str | None = None,
        result_summary: str | None = None,
        error: str | None = None,
        payload: dict[str, Any] | None = None,
    ) -> None:
        row = await self.get_task(conn, task_id)
        if not row:
            return
        now = _utc_now()
        st = status or row.status
        rs = result_summary if result_summary is not None else row.result_summary
        err = error if error is not None else row.error
        pl = json.dumps(payload, ensure_ascii=False) if payload is not None else json.dumps(
            row.payload, ensure_ascii=False
        )
        await conn.execute(
            """
            UPDATE tasks SET status=?, updated_at=?, result_summary=?, error=?,
                payload_json=? WHERE id=?
            """,
            (st, now, rs, err, pl, task_id),
        )
        await conn.commit()

    def _row_to_task(self, row: tuple[Any, ...]) -> TaskRow:
        (
            tid,
            ttype,
            status,
            parent_id,
            created_at,
            updated_at,
            payload_json,
            result_summary,
            error,
        ) = row
        try:
            payload = json.loads(payload_json or "{}")
        except json.JSONDecodeError:
            payload = {}
        return TaskRow(
            id=tid,
            type=ttype,
            status=status,
            parent_id=parent_id,
            created_at=created_at,
            updated_at=updated_at,
            payload=payload,
            result_summary=result_summary,
            error=error,
        )

    async def get_task(
        self, conn: aiosqlite.Connection, task_id: str
    ) -> TaskRow | None:
        cur = await conn.execute(
            "SELECT id, type, status, parent_id, created_at, updated_at, "
            "payload_json, result_summary, error FROM tasks WHERE id=?",
            (task_id,),
        )
        row = await cur.fetchone()
        if not row:
            return None
        return self._row_to_task(row)

    async def list_tasks(
        self, conn: aiosqlite.Connection, limit: int = 100
    ) -> list[TaskRow]:
        cur = await conn.execute(
            "SELECT id, type, status, parent_id, created_at, updated_at, "
            "payload_json, result_summary, error FROM tasks "
            "ORDER BY updated_at DESC LIMIT ?",
            (limit,),
        )
        rows = await cur.fetchall()
        return [self._row_to_task(r) for r in rows]

    async def list_tasks_by_parent(
        self, conn: aiosqlite.Connection, parent_id: str, *, limit: int = 200
    ) -> list[TaskRow]:
        cur = await conn.execute(
            "SELECT id, type, status, parent_id, created_at, updated_at, "
            "payload_json, result_summary, error FROM tasks "
            "WHERE parent_id = ? ORDER BY updated_at DESC LIMIT ?",
            (parent_id, limit),
        )
        rows = await cur.fetchall()
        return [self._row_to_task(r) for r in rows]

    async def list_top_level_tasks_for_session_thread(
        self,
        conn: aiosqlite.Connection,
        session_thread_id: str,
        *,
        limit: int = 80,
        since_updated_at: str | None = None,
    ) -> list[TaskRow]:
        """顶层任务且 payload.thread_root_id 等于主会话 id（由主对话发起的后台任务）。"""
        sid = (session_thread_id or "").strip()
        if not sid:
            return []
        if since_updated_at and since_updated_at.strip():
            cur = await conn.execute(
                """
                SELECT id, type, status, parent_id, created_at, updated_at,
                    payload_json, result_summary, error FROM tasks
                WHERE parent_id IS NULL
                  AND json_extract(payload_json, '$.thread_root_id') = ?
                  AND updated_at > ?
                ORDER BY updated_at DESC
                LIMIT ?
                """,
                (sid, since_updated_at.strip(), limit),
            )
        else:
            cur = await conn.execute(
                """
                SELECT id, type, status, parent_id, created_at, updated_at,
                    payload_json, result_summary, error FROM tasks
                WHERE parent_id IS NULL
                  AND json_extract(payload_json, '$.thread_root_id') = ?
                ORDER BY updated_at DESC
                LIMIT ?
                """,
                (sid, limit),
            )
        rows = await cur.fetchall()
        return [self._row_to_task(r) for r in rows]

    async def clear_task_history(self, conn: aiosqlite.Connection) -> dict[str, int]:
        """删除全部 task_messages 与 tasks 行（不删 chat_* 表）。"""
        cur = await conn.execute("SELECT COUNT(*) FROM task_messages")
        row_m = await cur.fetchone()
        n_msg = int(row_m[0]) if row_m else 0
        cur = await conn.execute("SELECT COUNT(*) FROM tasks")
        row_t = await cur.fetchone()
        n_tasks = int(row_t[0]) if row_t else 0
        await conn.execute("DELETE FROM task_messages")
        await conn.execute("DELETE FROM tasks")
        await conn.commit()
        return {"task_messages_deleted": n_msg, "tasks_deleted": n_tasks}

    async def append_chat_message(
        self,
        conn: aiosqlite.Connection,
        *,
        session_id: str,
        role: str,
        content: str | None,
        tool_calls: list[dict[str, Any]] | None = None,
        tool_call_id: str | None = None,
        tool_name: str | None = None,
    ) -> int:
        now = _utc_now()
        tc = json.dumps(tool_calls, ensure_ascii=False) if tool_calls else None
        cur = await conn.execute(
            """
            INSERT INTO chat_messages (
                session_id, role, content, tool_calls_json, created_at,
                tool_call_id, tool_name)
            VALUES (?, ?, ?, ?, ?, ?, ?)
            """,
            (
                session_id,
                role,
                content or "",
                tc,
                now,
                tool_call_id,
                tool_name,
            ),
        )
        await conn.commit()
        return int(cur.lastrowid)

    def _row_to_chat_message(self, row: tuple[Any, ...]) -> ChatMessageRow:
        return ChatMessageRow(
            id=int(row[0]),
            session_id=str(row[1]),
            role=str(row[2]),
            content=str(row[3] or ""),
            tool_calls_json=row[4] if row[4] else None,
            created_at=str(row[5]),
            tool_call_id=str(row[6]) if row[6] is not None else None,
            tool_name=str(row[7]) if row[7] is not None else None,
        )

    async def list_chat_messages(
        self,
        conn: aiosqlite.Connection,
        session_id: str,
        *,
        limit: int = 500,
        offset: int = 0,
    ) -> list[ChatMessageRow]:
        cur = await conn.execute(
            """
            SELECT id, session_id, role, content, tool_calls_json, created_at,
                   tool_call_id, tool_name
            FROM chat_messages
            WHERE session_id = ?
            ORDER BY id ASC
            LIMIT ? OFFSET ?
            """,
            (session_id, limit, offset),
        )
        rows = await cur.fetchall()
        return [self._row_to_chat_message(r) for r in rows]

    async def count_chat_messages(
        self, conn: aiosqlite.Connection, session_id: str
    ) -> int:
        cur = await conn.execute(
            "SELECT COUNT(*) FROM chat_messages WHERE session_id=?",
            (session_id,),
        )
        r = await cur.fetchone()
        return int(r[0]) if r else 0

    async def list_chat_messages_tail(
        self,
        conn: aiosqlite.Connection,
        session_id: str,
        *,
        limit: int,
    ) -> list[ChatMessageRow]:
        cur = await conn.execute(
            """
            SELECT id, session_id, role, content, tool_calls_json, created_at,
                   tool_call_id, tool_name
            FROM chat_messages
            WHERE session_id = ?
            ORDER BY id DESC
            LIMIT ?
            """,
            (session_id, limit),
        )
        rows = await cur.fetchall()
        rev = list(reversed(rows))
        return [self._row_to_chat_message(r) for r in rev]

    async def get_chat_message(
        self, conn: aiosqlite.Connection, message_id: int
    ) -> ChatMessageRow | None:
        cur = await conn.execute(
            """
            SELECT id, session_id, role, content, tool_calls_json, created_at,
                   tool_call_id, tool_name
            FROM chat_messages WHERE id=?
            """,
            (message_id,),
        )
        row = await cur.fetchone()
        if not row:
            return None
        return self._row_to_chat_message(row)

    async def append_chat_summary(
        self,
        conn: aiosqlite.Connection,
        *,
        session_id: str,
        content: str,
        covers_up_to_message_id: int | None = None,
    ) -> int:
        now = _utc_now()
        cur = await conn.execute(
            """
            INSERT INTO chat_summaries (session_id, content, covers_up_to_message_id, created_at)
            VALUES (?, ?, ?, ?)
            """,
            (session_id, content, covers_up_to_message_id, now),
        )
        await conn.commit()
        return int(cur.lastrowid)

    async def list_chat_summaries(
        self, conn: aiosqlite.Connection, session_id: str, limit: int = 20
    ) -> list[ChatSummaryRow]:
        cur = await conn.execute(
            """
            SELECT id, session_id, content, covers_up_to_message_id, created_at
            FROM chat_summaries
            WHERE session_id = ?
            ORDER BY id ASC
            LIMIT ?
            """,
            (session_id, limit),
        )
        rows = await cur.fetchall()
        out: list[ChatSummaryRow] = []
        for r in rows:
            out.append(
                ChatSummaryRow(
                    id=int(r[0]),
                    session_id=str(r[1]),
                    content=str(r[2]),
                    covers_up_to_message_id=int(r[3]) if r[3] is not None else None,
                    created_at=str(r[4]),
                )
            )
        return out

    def _row_to_task_message(self, row: tuple[Any, ...]) -> TaskMessageRow:
        (
            mid,
            thread_root_id,
            to_task_id,
            from_task_id,
            kind,
            body_json,
            created_at,
            read_at,
        ) = row
        try:
            body = json.loads(body_json or "{}")
        except json.JSONDecodeError:
            body = {}
        if not isinstance(body, dict):
            body = {"value": body}
        return TaskMessageRow(
            id=int(mid),
            thread_root_id=str(thread_root_id) if thread_root_id else None,
            to_task_id=str(to_task_id),
            from_task_id=str(from_task_id),
            kind=str(kind),
            body=body,
            created_at=str(created_at),
            read_at=str(read_at) if read_at else None,
        )

    async def insert_task_message(
        self,
        conn: aiosqlite.Connection,
        *,
        to_task_id: str,
        from_task_id: str,
        kind: str,
        body: dict[str, Any],
        thread_root_id: str | None = None,
    ) -> int:
        now = _utc_now()
        cur = await conn.execute(
            """
            INSERT INTO task_messages (
                thread_root_id, to_task_id, from_task_id, kind, body_json, created_at, read_at)
            VALUES (?, ?, ?, ?, ?, ?, NULL)
            """,
            (
                thread_root_id,
                to_task_id,
                from_task_id,
                kind,
                json.dumps(body, ensure_ascii=False),
                now,
            ),
        )
        await conn.commit()
        return int(cur.lastrowid)

    async def list_task_messages(
        self,
        conn: aiosqlite.Connection,
        *,
        to_task_id: str,
        since_id: int = 0,
        limit: int = 50,
    ) -> list[TaskMessageRow]:
        cur = await conn.execute(
            """
            SELECT id, thread_root_id, to_task_id, from_task_id, kind, body_json, created_at, read_at
            FROM task_messages
            WHERE to_task_id = ? AND id > ?
            ORDER BY id ASC
            LIMIT ?
            """,
            (to_task_id, since_id, limit),
        )
        rows = await cur.fetchall()
        return [self._row_to_task_message(r) for r in rows]

    async def list_task_messages_unread(
        self,
        conn: aiosqlite.Connection,
        *,
        to_task_id: str,
        limit: int = 50,
    ) -> list[TaskMessageRow]:
        cur = await conn.execute(
            """
            SELECT id, thread_root_id, to_task_id, from_task_id, kind, body_json, created_at, read_at
            FROM task_messages
            WHERE to_task_id = ? AND read_at IS NULL
            ORDER BY id ASC
            LIMIT ?
            """,
            (to_task_id, limit),
        )
        rows = await cur.fetchall()
        return [self._row_to_task_message(r) for r in rows]

    async def mark_task_messages_read(
        self, conn: aiosqlite.Connection, message_ids: list[int]
    ) -> None:
        if not message_ids:
            return
        now = _utc_now()
        qmarks = ",".join("?" * len(message_ids))
        await conn.execute(
            f"UPDATE task_messages SET read_at=? WHERE id IN ({qmarks})",
            [now, *message_ids],
        )
        await conn.commit()
