"""cortex/storage/queries/entity_queries.py — Entity storage mixin."""

from __future__ import annotations

from typing import Any, Dict, List, Optional

from cortex.storage._helpers import _now, _row_to_entity, _row_to_entity_alias


class EntityQueryMixin:
    """Mixin providing entityquery methods."""

    async def create_entity(
        self,
        canonical_name: str,
        entity_type: str = "unknown",
        owner_id: str = "default",
        **kwargs,
    ) -> Entity:
        """Insert a new entity row and return the created Entity."""
        conn = await self._get_conn()
        now = _now()
        description = kwargs.get("description")
        metadata_json = kwargs.get("metadata_json")

        async with self._write_lock:
            async with conn.execute(
                """
                INSERT INTO entity (owner_id, canonical_name, entity_type,
                                    description, metadata_json,
                                    first_seen, last_seen, created_at, updated_at)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
                RETURNING *
                """,
                (owner_id, canonical_name, entity_type, description, metadata_json,
                 now, now, now, now),
            ) as cur:
                row = await cur.fetchone()
            await conn.commit()
        return _row_to_entity(row)


    async def get_entity(self, entity_id: str) -> Optional[Entity]:
        """Return an Entity by its primary-key id, or None if not found."""
        conn = await self._get_conn()
        async with conn.execute(
            "SELECT * FROM entity WHERE id = ?", (entity_id,)
        ) as cur:
            row = await cur.fetchone()
        return _row_to_entity(row) if row else None


    async def get_entity_by_id(self, entity_id: str) -> Optional[Entity]:
        """Alias for get_entity — satisfies StorageAdapter ABC."""
        return await self.get_entity(entity_id)

    async def update_entity(
        self,
        entity_id: str,
        *,
        entity_type: str | None = None,
        owner_id: str | None = None,
    ) -> None:
        """Update mutable fields on an existing entity."""
        sets: list[str] = []
        params: list[str] = []
        if entity_type is not None:
            sets.append("entity_type = ?")
            params.append(entity_type)
        if not sets:
            return
        params.append(entity_id)
        conn = await self._get_conn()
        await conn.execute(
            f"UPDATE entity SET {', '.join(sets)} WHERE id = ?",
            tuple(params),
        )
        await conn.commit()


    async def get_all_entity_aliases(
        self, owner_id: str = "default"
    ) -> List[EntityAlias]:
        """Return all entity aliases for owner_id."""
        conn = await self._get_conn()
        async with conn.execute(
            """
            SELECT a.* FROM entity_alias a
            JOIN entity e ON e.id = a.entity_id
            WHERE e.owner_id = ?
            ORDER BY a.entity_id, a.alias
            """,
            (owner_id,),
        ) as cur:
            rows = await cur.fetchall()
        return [_row_to_entity_alias(row) for row in rows]


    async def find_entity_by_alias(
        self, alias: str, owner_id: str = "default"
    ) -> Optional[Entity]:
        """Look up an entity by alias within an owner scope; returns None if not found."""
        conn = await self._get_conn()
        async with conn.execute(
            """
            SELECT e.* FROM entity e
            JOIN entity_alias a ON a.entity_id = e.id
            WHERE a.alias = ? COLLATE NOCASE
              AND e.owner_id = ?
            LIMIT 1
            """,
            (alias, owner_id),
        ) as cur:
            row = await cur.fetchone()
        return _row_to_entity(row) if row else None


    async def add_entity_alias(
        self,
        entity_id: str,
        alias: str,
        alias_type: str = "name",
        **kwargs,
    ) -> EntityAlias:
        async with self._write_lock:
            conn = await self._get_conn()
            now = _now()
            source = kwargs.get("source", "extracted")
            confidence = kwargs.get("confidence", 1.0)

            async with conn.execute(
                """
                INSERT INTO entity_alias (entity_id, alias, alias_type,
                                          source, confidence, created_at)
                VALUES (?, ?, ?, ?, ?, ?)
                RETURNING *
                """,
                (entity_id, alias, alias_type, source, confidence, now),
            ) as cur:
                row = await cur.fetchone()
            await conn.commit()
            return _row_to_entity_alias(row)


    async def merge_entities(self, keep_id: str, merge_id: str) -> Entity:
        """Merge merge_id into keep_id: re-points all aliases/claims/edges, sums mention_count, deletes merge_id."""
        async with self._write_lock:
            conn = await self._get_conn()
            now = _now()

            # Re-point aliases
            await conn.execute(
                "UPDATE entity_alias SET entity_id = ? WHERE entity_id = ?",
                (keep_id, merge_id),
            )
            # Re-point claims
            await conn.execute(
                "UPDATE semantic_claim SET subject_entity_id = ? WHERE subject_entity_id = ?",
                (keep_id, merge_id),
            )
            await conn.execute(
                "UPDATE semantic_claim SET object_entity_id = ? WHERE object_entity_id = ?",
                (keep_id, merge_id),
            )
            # Re-point memory edges
            await conn.execute(
                "UPDATE memory_edge SET source_id = ? WHERE source_type = 'entity' AND source_id = ?",
                (keep_id, merge_id),
            )
            await conn.execute(
                "UPDATE memory_edge SET target_id = ? WHERE target_type = 'entity' AND target_id = ?",
                (keep_id, merge_id),
            )
            # Re-point peer records (FK: peer.entity_id → entity.id, no CASCADE)
            await conn.execute(
                "UPDATE peer SET entity_id = ? WHERE entity_id = ?",
                (keep_id, merge_id),
            )
            # Update mention_count and updated_at on keep
            await conn.execute(
                """
                UPDATE entity
                   SET mention_count = mention_count + (
                       SELECT mention_count FROM entity WHERE id = ?
                   ),
                       updated_at = ?
                 WHERE id = ?
                """,
                (merge_id, now, keep_id),
            )
            # Delete merged entity (cascades aliases)
            await conn.execute("DELETE FROM entity WHERE id = ?", (merge_id,))
            await conn.commit()

            kept = await self.get_entity(keep_id)
            assert kept is not None
            return kept


    async def list_entities(
        self,
        owner_id: str = "default",
        entity_type: Optional[str] = None,
        sort_by: str = "last_seen",
        limit: int = 50,
        offset: int = 0,
        search: Optional[str] = None,
    ) -> List["Entity"]:
        """Browse entities with filtering, sorting, and pagination."""
        conn = await self._get_conn()
        sort_col = "last_seen" if sort_by != "mention_count" else "mention_count"
        params: list = [owner_id]
        where_clauses = ["owner_id = ?"]
        if entity_type:
            where_clauses.append("entity_type = ?")
            params.append(entity_type)
        if search:
            # Escape LIKE wildcards in user input (#220)
            safe_search = search.replace("\\", "\\\\").replace("%", "\\%").replace("_", "\\_")
            where_clauses.append("canonical_name LIKE ? ESCAPE '\\'")
            params.append(f"%{safe_search}%")
        where_sql = " AND ".join(where_clauses)
        params += [limit, offset]
        async with conn.execute(
            f"""
            SELECT * FROM entity
            WHERE {where_sql}
            ORDER BY {sort_col} DESC
            LIMIT ? OFFSET ?
            """,
            params,
        ) as cur:
            rows = await cur.fetchall()
        return [_row_to_entity(r) for r in rows]


    async def count_entities(self, owner_id: str = "default") -> int:
        conn = await self._get_conn()
        async with conn.execute(
            "SELECT COUNT(*) FROM entity WHERE owner_id = ?", (owner_id,)
        ) as cur:
            row = await cur.fetchone()
        return row[0] if row else 0


    async def get_entity_aliases(self, entity_id: str) -> List["EntityAlias"]:
        conn = await self._get_conn()
        async with conn.execute(
            "SELECT * FROM entity_alias WHERE entity_id = ? ORDER BY created_at ASC",
            (entity_id,),
        ) as cur:
            rows = await cur.fetchall()
        return [_row_to_entity_alias(r) for r in rows]


    async def get_entity_claim_count(self, entity_id: str) -> int:
        conn = await self._get_conn()
        async with conn.execute(
            """
            SELECT COUNT(*) FROM semantic_claim
            WHERE subject_entity_id = ? OR object_entity_id = ?
            """,
            (entity_id, entity_id),
        ) as cur:
            row = await cur.fetchone()
        return row[0] if row else 0


    async def get_entity_relationships(
        self, entity_id: str, limit: int = 20
    ) -> List[Dict[str, Any]]:
        """Return top relationships where entity appears as subject or object."""
        conn = await self._get_conn()
        async with conn.execute(
            """
            SELECT c.id, c.predicate, c.object_text, c.subject_entity_id,
                   c.object_entity_id, c.confidence, c.created_at
            FROM semantic_claim c
            WHERE c.subject_entity_id = ? OR c.object_entity_id = ?
            ORDER BY c.confidence DESC
            LIMIT ?
            """,
            (entity_id, entity_id, limit),
        ) as cur:
            rows = await cur.fetchall()
        return [dict(r) for r in rows]


