"""Storage providers for per-thread memory."""

from __future__ import annotations

import abc
import json
import logging
import sqlite3
import threading
from datetime import UTC, datetime
from pathlib import Path
from typing import Any

from deerflow.agents.memory.thread_prompt import create_empty_thread_memory
from deerflow.config.paths import get_paths
from deerflow.config.thread_memory_config import get_thread_memory_config

logger = logging.getLogger(__name__)


class ThreadMemoryStorage(abc.ABC):
    @abc.abstractmethod
    def load(self, thread_id: str) -> dict[str, Any] | None:
        pass

    @abc.abstractmethod
    def save(self, thread_id: str, data: dict[str, Any], expected_version: int | None = None) -> bool:
        pass

    @abc.abstractmethod
    def delete(self, thread_id: str) -> bool:
        pass


def _row_to_memory(row: tuple[Any, ...]) -> dict[str, Any]:
    thread_id, owner_id_col, memory_json_raw, memory_version, last_updated = row
    decoded: dict[str, Any] = {}
    if isinstance(memory_json_raw, str) and memory_json_raw.strip():
        try:
            parsed_json = json.loads(memory_json_raw)
            if isinstance(parsed_json, dict):
                decoded = parsed_json
        except Exception:
            decoded = {}

    owner_id = decoded.get("ownerId")
    if owner_id is None:
        owner_id = owner_id_col

    user = decoded.get("user", create_empty_thread_memory()["user"])
    history = decoded.get("history", create_empty_thread_memory()["history"])
    facts = decoded.get("facts", [])
    return {
        "threadId": thread_id,
        "ownerId": owner_id,
        "user": user,
        "history": history,
        "facts": facts,
        "memoryVersion": int(memory_version),
        "lastUpdated": str(last_updated),
    }


class SqliteThreadMemoryStorage(ThreadMemoryStorage):
    def __init__(self, db_path: str):
        path = Path(db_path)
        if not path.is_absolute():
            path = get_paths().base_dir / path
        path.parent.mkdir(parents=True, exist_ok=True)
        self._conn = sqlite3.connect(str(path), check_same_thread=False)
        self._lock = threading.Lock()
        with self._lock:
            self._conn.execute(
                """
                CREATE TABLE IF NOT EXISTS thread_memory (
                    thread_id TEXT PRIMARY KEY,
                    owner_id TEXT NULL,
                    memory_json TEXT NOT NULL DEFAULT '',
                    memory_version INTEGER NOT NULL DEFAULT 0,
                    last_updated TEXT NOT NULL DEFAULT (datetime('now'))
                )
                """
            )
            self._ensure_memory_json_column()
            self._conn.execute("CREATE INDEX IF NOT EXISTS idx_thread_memory_owner_id ON thread_memory(owner_id)")
            self._conn.commit()

    def _ensure_memory_json_column(self) -> None:
        columns = self._conn.execute("PRAGMA table_info(thread_memory)").fetchall()
        has_memory_json = any(col[1] == "memory_json" for col in columns)
        if not has_memory_json:
            self._conn.execute("ALTER TABLE thread_memory ADD COLUMN memory_json TEXT NOT NULL DEFAULT ''")

    def load(self, thread_id: str) -> dict[str, Any] | None:
        with self._lock:
            row = self._conn.execute(
                "SELECT thread_id, owner_id, memory_json, memory_version, last_updated "
                "FROM thread_memory WHERE thread_id = ?",
                (thread_id,),
            ).fetchone()
            if row is None:
                return None
            return _row_to_memory(row)

    def save(self, thread_id: str, data: dict[str, Any], expected_version: int | None = None) -> bool:
        now = datetime.now(UTC).isoformat().replace("+00:00", "Z")
        owner_id = data.get("ownerId")
        if expected_version is None:
            expected_version = 0
        with self._lock:
            cur = self._conn.execute(
                """
                INSERT INTO thread_memory (thread_id, owner_id, memory_json, memory_version, last_updated)
                VALUES (?, ?, ?, 0, ?)
                ON CONFLICT(thread_id) DO NOTHING
                """,
                (
                    thread_id,
                    owner_id,
                    json.dumps(data, ensure_ascii=False),
                    now,
                ),
            )
            if cur.rowcount == 1:
                self._conn.commit()
                return True

            cur = self._conn.execute(
                """
                UPDATE thread_memory
                SET owner_id = ?, memory_json = ?, memory_version = memory_version + 1, last_updated = ?
                WHERE thread_id = ? AND memory_version = ?
                """,
                (
                    owner_id,
                    json.dumps(data, ensure_ascii=False),
                    now,
                    thread_id,
                    expected_version,
                ),
            )
            self._conn.commit()
            return cur.rowcount == 1

    def delete(self, thread_id: str) -> bool:
        with self._lock:
            self._conn.execute("DELETE FROM thread_memory WHERE thread_id = ?", (thread_id,))
            self._conn.commit()
        return True

    def count_legacy_rows(self) -> int:
        return 0

    def backfill_legacy_rows(self, *, limit: int | None = None) -> dict[str, int]:
        _ = limit
        return {"scanned": 0, "updated": 0, "skipped": 0, "failed": 0}


class MysqlThreadMemoryStorage(ThreadMemoryStorage):
    def __init__(self, host: str, port: int, user: str, password: str, database: str):
        import pymysql

        self._conn = pymysql.connect(host=host, port=port, user=user, password=password, database=database, charset="utf8mb4")
        with self._conn.cursor() as cur:
            cur.execute(
                """
                CREATE TABLE IF NOT EXISTS thread_memory (
                    thread_id VARCHAR(64) PRIMARY KEY,
                    owner_id VARCHAR(64) NULL,
                    memory_json LONGTEXT NOT NULL,
                    memory_version INT NOT NULL DEFAULT 0,
                    last_updated TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,
                    INDEX idx_owner_id (owner_id)
                )
                """
            )
            cur.execute("SHOW COLUMNS FROM thread_memory LIKE 'memory_json'")
            if cur.fetchone() is None:
                cur.execute("ALTER TABLE thread_memory ADD COLUMN memory_json LONGTEXT NOT NULL DEFAULT ''")
        self._conn.commit()

    def load(self, thread_id: str) -> dict[str, Any] | None:
        with self._conn.cursor() as cur:
            cur.execute(
                "SELECT thread_id, owner_id, memory_json, memory_version, last_updated FROM thread_memory WHERE thread_id = %s",
                (thread_id,),
            )
            row = cur.fetchone()
            if row is None:
                return None
            return _row_to_memory(row)

    def save(self, thread_id: str, data: dict[str, Any], expected_version: int | None = None) -> bool:
        if expected_version is None:
            expected_version = 0
        owner_id = data.get("ownerId")
        with self._conn.cursor() as cur:
            cur.execute(
                """
                INSERT INTO thread_memory (thread_id, owner_id, memory_json, memory_version)
                VALUES (%s, %s, %s, 0)
                ON DUPLICATE KEY UPDATE thread_id = thread_id
                """,
                (
                    thread_id,
                    owner_id,
                    json.dumps(data, ensure_ascii=False),
                ),
            )
            if cur.rowcount == 1:
                self._conn.commit()
                return True
            cur.execute(
                """
                UPDATE thread_memory
                SET owner_id = %s, memory_json = %s, memory_version = memory_version + 1
                WHERE thread_id = %s AND memory_version = %s
                """,
                (
                    owner_id,
                    json.dumps(data, ensure_ascii=False),
                    thread_id,
                    expected_version,
                ),
            )
        self._conn.commit()
        return cur.rowcount == 1

    def delete(self, thread_id: str) -> bool:
        with self._conn.cursor() as cur:
            cur.execute("DELETE FROM thread_memory WHERE thread_id = %s", (thread_id,))
        self._conn.commit()
        return True

    def count_legacy_rows(self) -> int:
        return 0

    def backfill_legacy_rows(self, *, limit: int | None = None) -> dict[str, int]:
        _ = limit
        return {"scanned": 0, "updated": 0, "skipped": 0, "failed": 0}


_thread_storage: ThreadMemoryStorage | None = None
_thread_storage_lock = threading.Lock()


def get_thread_memory_storage() -> ThreadMemoryStorage:
    global _thread_storage
    if _thread_storage is not None:
        return _thread_storage

    with _thread_storage_lock:
        if _thread_storage is not None:
            return _thread_storage
        config = get_thread_memory_config()
        if config.database.type == "mysql":
            mysql = config.database.mysql
            _thread_storage = MysqlThreadMemoryStorage(
                host=mysql.host,
                port=mysql.port,
                user=mysql.user,
                password=mysql.password,
                database=mysql.database,
            )
        else:
            _thread_storage = SqliteThreadMemoryStorage(config.database.sqlite.path)
    return _thread_storage


def get_thread_memory_data(thread_id: str) -> dict[str, Any] | None:
    return get_thread_memory_storage().load(thread_id)


def delete_thread_memory_data(thread_id: str) -> bool:
    return get_thread_memory_storage().delete(thread_id)


def initial_thread_memory_record() -> dict[str, Any]:
    return {"ownerId": None, **create_empty_thread_memory()}
