import json

from .base import BaseMemoryRepo


class UserMemoryRepo(BaseMemoryRepo):
    """管理 user_memories + vec_user_memories 表（跨项目用户偏好）"""
    TABLE = "user_memories"
    VEC_TABLE = "vec_user_memories"
    TAG_TABLE = "user_memory_tags"

    def __init__(self, conn):
        super().__init__(conn, project_dir="")

    def search(self, embedding: list[float], top_k: int = 5, tags: list[str] | None = None,
               source: str | None = None) -> list[dict]:
        """向量搜索（用户记忆无 project_dir 过滤）"""
        emb_blob = json.dumps(embedding).encode("utf-8")
        sql = f"""
            SELECT m.*, vec_distance_cosine(v.embedding, ?) as distance
            FROM {self.TABLE} m
            JOIN {self.VEC_TABLE} v ON m.rowid = v.rowid
        """
        params: list = [emb_blob]
        conditions = []

        if tags:
            tag_placeholders = ",".join("?" * len(tags))
            conditions.append(f"m.id IN (SELECT memory_id FROM {self.TAG_TABLE} WHERE tag IN ({tag_placeholders}))")
            params.extend(tags)

        if source:
            conditions.append("m.source=?")
            params.append(source)

        if conditions:
            sql += " WHERE " + " AND ".join(conditions)

        sql += " ORDER BY distance LIMIT ?"
        params.append(top_k)

        return [dict(r) for r in self.conn.execute(sql, params).fetchall()]

    def get_tag_counts(self) -> dict[str, int]:
        sql = f"SELECT tag, COUNT(*) as count FROM {self.TAG_TABLE} GROUP BY tag ORDER BY count DESC"
        return {row["tag"]: row["count"] for row in self.conn.execute(sql).fetchall()}

    def get_ids_with_tag(self, tag: str) -> list[dict]:
        return [dict(r) for r in self.conn.execute(
            f"SELECT um.id, um.tags FROM {self.TABLE} um JOIN {self.TAG_TABLE} umt ON um.id = umt.memory_id WHERE umt.tag=?",
            (tag,)
        ).fetchall()]
