"""Memory graph edge operations mixin."""
from __future__ import annotations

from typing import List, Optional

from cortex.storage._helpers import _now, _row_to_edge
from cortex.types import MemoryEdge


class ClaimGraphMixin:
    """Mixin providing memory graph edge creation and traversal."""

    async def create_edge(
        self,
        source_type: str,
        source_id: str,
        target_type: str,
        target_id: str,
        edge_type: str,
        owner_id: str = "default",
        **kwargs,
    ) -> MemoryEdge:
        async with self._write_lock:
            conn = await self._get_conn()
            now = _now()
            confidence = kwargs.get("confidence", 1.0)
            metadata_json = kwargs.get("metadata_json")

            async with conn.execute(
                """
                INSERT INTO memory_edge
                    (owner_id, source_type, source_id,
                     target_type, target_id, edge_type,
                     confidence, metadata_json, created_at)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
                RETURNING *
                """,
                (owner_id, source_type, source_id,
                 target_type, target_id, edge_type,
                 confidence, metadata_json, now),
            ) as cur:
                row = await cur.fetchone()
            await conn.commit()
            return _row_to_edge(row)

    async def get_neighbors(
        self,
        item_type: str,
        item_id: str,
        edge_types: Optional[List[str]] = None,
        max_depth: int = 1,
    ) -> List[MemoryEdge]:
        conn = await self._get_conn()

        if edge_types:
            placeholders = ", ".join("?" for _ in edge_types)
            async with conn.execute(
                f"""
                SELECT * FROM memory_edge
                WHERE source_type = ? AND source_id = ?
                  AND edge_type IN ({placeholders})
                """,
                [item_type, item_id] + edge_types,
            ) as cur:
                rows = await cur.fetchall()
        else:
            async with conn.execute(
                """
                SELECT * FROM memory_edge
                WHERE source_type = ? AND source_id = ?
                """,
                (item_type, item_id),
            ) as cur:
                rows = await cur.fetchall()

        return [_row_to_edge(r) for r in rows]
