"""Canvas A2UI — persistent service + bidirectional WebSocket handler."""

from __future__ import annotations

import json
import logging
import uuid
from datetime import UTC, datetime

import aiosqlite
from fastapi import WebSocket, WebSocketDisconnect

log = logging.getLogger(__name__)


class CanvasService:
    """Persistent canvas state backed by SQLite with live WebSocket broadcast."""

    def __init__(self, db: aiosqlite.Connection) -> None:
        self._db = db
        self._connections: set[WebSocket] = set()

    # ── helpers ──────────────────────────────────────────────────────────

    def _now(self) -> str:
        return datetime.now(UTC).isoformat()

    async def _row_to_dict(self, row: aiosqlite.Row) -> dict:
        return {
            "id": row["id"],
            "title": row["title"],
            "content": row["content"],
            "content_type": row["content_type"],
            "pinned": bool(row["pinned"]),
            "collapsed": bool(row["collapsed"]),
            "sort_order": row["sort_order"],
            "user_note": row["user_note"],
            "pushed_at": row["pushed_at"],
            "updated_at": row["updated_at"],
        }

    async def _fetch_item(self, item_id: str) -> dict | None:
        cur = await self._db.execute("SELECT * FROM canvas_items WHERE id = ?", (item_id,))
        row = await cur.fetchone()
        return await self._row_to_dict(row) if row else None

    # ── CRUD ─────────────────────────────────────────────────────────────

    async def get_items(self) -> list[dict]:
        cur = await self._db.execute(
            "SELECT * FROM canvas_items ORDER BY pinned DESC, sort_order ASC, pushed_at ASC"
        )
        rows = await cur.fetchall()
        return [await self._row_to_dict(r) for r in rows]

    async def push(self, content: str, content_type: str = "markdown", title: str = "") -> dict:
        now = self._now()
        # next sort_order
        cur = await self._db.execute("SELECT COALESCE(MAX(sort_order), -1) + 1 FROM canvas_items")
        row = await cur.fetchone()
        sort_order = row[0] if row else 0

        item_id = str(uuid.uuid4())
        await self._db.execute(
            """INSERT INTO canvas_items (id, title, content, content_type, sort_order, pushed_at, updated_at)
               VALUES (?, ?, ?, ?, ?, ?, ?)""",
            (item_id, title, content, content_type, sort_order, now, now),
        )
        await self._db.commit()
        item = await self._fetch_item(item_id)
        await self.broadcast({"type": "push", "item": item})
        return item  # type: ignore[return-value]

    async def update(self, item_id: str, **fields: str) -> dict | None:
        allowed = {"content", "title", "content_type"}
        updates = {k: v for k, v in fields.items() if k in allowed and v is not None}
        if not updates:
            return await self._fetch_item(item_id)
        updates["updated_at"] = self._now()
        set_clause = ", ".join(f"{k} = ?" for k in updates)
        vals = list(updates.values()) + [item_id]
        await self._db.execute(f"UPDATE canvas_items SET {set_clause} WHERE id = ?", vals)
        await self._db.commit()
        item = await self._fetch_item(item_id)
        if item:
            await self.broadcast({"type": "update", "item": item})
        return item

    async def remove(self, item_id: str) -> bool:
        cur = await self._db.execute("DELETE FROM canvas_items WHERE id = ?", (item_id,))
        await self._db.commit()
        if cur.rowcount:
            await self.broadcast({"type": "remove", "id": item_id})
            return True
        return False

    async def set_pinned(self, item_id: str, pinned: bool) -> dict | None:
        await self._db.execute(
            "UPDATE canvas_items SET pinned = ?, updated_at = ? WHERE id = ?",
            (int(pinned), self._now(), item_id),
        )
        await self._db.commit()
        item = await self._fetch_item(item_id)
        if item:
            await self.broadcast({"type": "update", "item": item})
        return item

    async def set_collapsed(self, item_id: str, collapsed: bool) -> dict | None:
        await self._db.execute(
            "UPDATE canvas_items SET collapsed = ?, updated_at = ? WHERE id = ?",
            (int(collapsed), self._now(), item_id),
        )
        await self._db.commit()
        item = await self._fetch_item(item_id)
        if item:
            await self.broadcast({"type": "update", "item": item})
        return item

    async def set_user_note(self, item_id: str, note: str) -> dict | None:
        await self._db.execute(
            "UPDATE canvas_items SET user_note = ?, updated_at = ? WHERE id = ?",
            (note, self._now(), item_id),
        )
        await self._db.commit()
        item = await self._fetch_item(item_id)
        if item:
            await self.broadcast({"type": "update", "item": item})
        return item

    async def reorder(self, ids: list[str]) -> None:
        for idx, item_id in enumerate(ids):
            await self._db.execute(
                "UPDATE canvas_items SET sort_order = ? WHERE id = ?", (idx, item_id)
            )
        await self._db.commit()
        await self.broadcast({"type": "reorder", "ids": ids})

    async def reset(self) -> None:
        await self._db.execute("DELETE FROM canvas_items")
        await self._db.commit()
        await self.broadcast({"type": "reset"})

    # ── snapshots ────────────────────────────────────────────────────────

    async def snapshot_save(self, name: str) -> dict:
        items = await self.get_items()
        snap_id = str(uuid.uuid4())
        now = self._now()
        await self._db.execute(
            "INSERT INTO canvas_snapshots (id, name, items_json, created_at) VALUES (?, ?, ?, ?)",
            (snap_id, name, json.dumps(items), now),
        )
        await self._db.commit()
        return {"id": snap_id, "name": name, "created_at": now}

    async def snapshot_list(self) -> list[dict]:
        cur = await self._db.execute(
            "SELECT id, name, created_at FROM canvas_snapshots ORDER BY created_at DESC"
        )
        rows = await cur.fetchall()
        return [{"id": r["id"], "name": r["name"], "created_at": r["created_at"]} for r in rows]

    async def snapshot_restore(self, snapshot_id: str) -> bool:
        cur = await self._db.execute(
            "SELECT items_json FROM canvas_snapshots WHERE id = ?", (snapshot_id,)
        )
        row = await cur.fetchone()
        if not row:
            return False
        items = json.loads(row["items_json"])
        await self._db.execute("DELETE FROM canvas_items")
        for it in items:
            await self._db.execute(
                """INSERT INTO canvas_items
                   (id, title, content, content_type, pinned, collapsed, sort_order, user_note, pushed_at, updated_at)
                   VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""",
                (
                    it["id"], it.get("title", ""), it["content"], it["content_type"],
                    int(it.get("pinned", False)), int(it.get("collapsed", False)),
                    it.get("sort_order", 0), it.get("user_note", ""),
                    it["pushed_at"], it.get("updated_at", it["pushed_at"]),
                ),
            )
        await self._db.commit()
        new_items = await self.get_items()
        await self.broadcast({"type": "state", "items": new_items})
        return True

    async def snapshot_delete(self, snapshot_id: str) -> bool:
        cur = await self._db.execute(
            "DELETE FROM canvas_snapshots WHERE id = ?", (snapshot_id,)
        )
        await self._db.commit()
        return bool(cur.rowcount)

    # ── WebSocket management ─────────────────────────────────────────────

    def add_connection(self, ws: WebSocket) -> None:
        self._connections.add(ws)

    def remove_connection(self, ws: WebSocket) -> None:
        self._connections.discard(ws)

    async def broadcast(self, event: dict) -> None:
        dead: list[WebSocket] = []
        for ws in self._connections:
            try:
                await ws.send_json(event)
            except Exception:
                dead.append(ws)
        for ws in dead:
            self._connections.discard(ws)


# ── WebSocket handler ────────────────────────────────────────────────────

async def ws_canvas_handler(websocket: WebSocket) -> None:
    """Bidirectional WebSocket for the live canvas.

    Server → Client:
        {"type": "state", "items": [...]}
        {"type": "push", "item": {...}}
        {"type": "update", "item": {...}}
        {"type": "remove", "id": "..."}
        {"type": "reorder", "ids": [...]}
        {"type": "reset"}

    Client → Server:
        {"action": "remove", "id": "..."}
        {"action": "pin", "id": "...", "pinned": true}
        {"action": "collapse", "id": "...", "collapsed": true}
        {"action": "note", "id": "...", "note": "..."}
        {"action": "edit", "id": "...", "content": "...", "title": "..."}
        {"action": "reorder", "ids": [...]}
        {"action": "snapshot_save", "name": "..."}
        {"action": "snapshot_restore", "snapshot_id": "..."}
    """
    await websocket.accept()
    canvas: CanvasService = websocket.app.state.canvas
    canvas.add_connection(websocket)

    try:
        # send current state
        items = await canvas.get_items()
        await websocket.send_json({"type": "state", "items": items})

        while True:
            raw = await websocket.receive_text()
            try:
                msg = json.loads(raw)
            except json.JSONDecodeError:
                continue

            action = msg.get("action")
            item_id = msg.get("id", "")

            if action == "remove":
                await canvas.remove(item_id)
            elif action == "pin":
                await canvas.set_pinned(item_id, bool(msg.get("pinned", True)))
            elif action == "collapse":
                await canvas.set_collapsed(item_id, bool(msg.get("collapsed", True)))
            elif action == "note":
                await canvas.set_user_note(item_id, msg.get("note", ""))
            elif action == "edit":
                fields: dict = {}
                if "content" in msg:
                    fields["content"] = msg["content"]
                if "title" in msg:
                    fields["title"] = msg["title"]
                if fields:
                    await canvas.update(item_id, **fields)
            elif action == "reorder":
                ids = msg.get("ids", [])
                if ids:
                    await canvas.reorder(ids)
            elif action == "snapshot_save":
                name = msg.get("name", "Untitled")
                await canvas.snapshot_save(name)
            elif action == "snapshot_restore":
                snap_id = msg.get("snapshot_id", "")
                if snap_id:
                    await canvas.snapshot_restore(snap_id)

    except WebSocketDisconnect:
        canvas.remove_connection(websocket)
    except Exception:
        canvas.remove_connection(websocket)
