graph_db.py

"""
graph_db.py
Database abstraction for the HF Space.
Uses SQLite by default (free, no external service needed).
Set SUPABASE_URL + SUPABASE_KEY env vars to switch to Supabase.
"""

import os
import json
import uuid
import logging
import aiosqlite
from datetime import datetime, timezone
from pathlib import Path

log = logging.getLogger("graph_db")

DB_PATH      = os.environ.get("DB_PATH", "/app/data/graph.db")
SUPABASE_URL = os.environ.get("SUPABASE_URL", "")
SUPABASE_KEY = os.environ.get("SUPABASE_KEY", "")

USE_SUPABASE = bool(SUPABASE_URL and SUPABASE_KEY)


def now_iso():
    return datetime.now(timezone.utc).isoformat()

def new_id():
    return str(uuid.uuid4())

def _row_to_dict(row, cursor):
    if row is None:
        return None
    cols = [d[0] for d in cursor.description]
    return dict(zip(cols, row))

def _rows_to_list(rows, cursor):
    cols = [d[0] for d in cursor.description]
    return [dict(zip(cols, r)) for r in rows]


def _is_supabase_error(payload) -> bool:
    """Supabase returns {'code': ..., 'message': ...} on error instead of a list.
    Treat anything that isn't a list as an error response."""
    if isinstance(payload, list):
        return False
    if isinstance(payload, dict) and ("code" in payload or "message" in payload or "error" in payload):
        return True
    return True


class GraphDB:
    def __init__(self):
        self._db = None

    async def init(self):
        if USE_SUPABASE:
            log.info("Using Supabase backend")
            return
        Path(DB_PATH).parent.mkdir(parents=True, exist_ok=True)
        self._db = await aiosqlite.connect(DB_PATH)
        await self._db.execute("PRAGMA journal_mode=WAL")
        await self._db.execute("PRAGMA foreign_keys=ON")
        await self._create_tables()
        log.info(f"SQLite ready: {DB_PATH}")

    async def _create_tables(self):
        await self._db.executescript("""
        CREATE TABLE IF NOT EXISTS agents (
            id TEXT PRIMARY KEY, name TEXT NOT NULL, agent_type TEXT,
            description TEXT, capabilities TEXT, wallet_address TEXT,
            nwo_agent_id TEXT UNIQUE, nwo_api_key TEXT,
            persona_color TEXT DEFAULT '#7F77DD',
            avatar_label TEXT, is_active INTEGER DEFAULT 1,
            created_at TEXT
        );
        CREATE TABLE IF NOT EXISTS graph_nodes (
            id TEXT PRIMARY KEY, user_id TEXT, agent_id TEXT,
            name TEXT NOT NULL, val REAL DEFAULT 1.0, color TEXT,
            depth_level INTEGER DEFAULT 0, description TEXT,
            category TEXT DEFAULT 'topic', actor_type TEXT DEFAULT 'human',
            x_position REAL, y_position REAL,
            expand_requested INTEGER DEFAULT 0, expand_done INTEGER DEFAULT 0,
            nwo_agent_id TEXT, nwo_task_id TEXT,
            robot_position TEXT, sensor_data TEXT, battery_level REAL,
            created_at TEXT, updated_at TEXT
        );
        CREATE INDEX IF NOT EXISTS idx_nodes_actor   ON graph_nodes(actor_type);
        CREATE INDEX IF NOT EXISTS idx_nodes_expand  ON graph_nodes(expand_requested);
        CREATE INDEX IF NOT EXISTS idx_nodes_created ON graph_nodes(created_at);
        CREATE TABLE IF NOT EXISTS graph_links (
            id TEXT PRIMARY KEY, source_id TEXT, target_id TEXT,
            similarity_score REAL DEFAULT 0.5, link_type TEXT DEFAULT 'semantic',
            connection_reason TEXT,
            created_by_actor_type TEXT, created_by_id TEXT, created_at TEXT
        );
        CREATE TABLE IF NOT EXISTS graph_posts (
            id TEXT PRIMARY KEY, node_id TEXT, content TEXT NOT NULL,
            actor_type TEXT DEFAULT 'human', actor_id TEXT, nwo_agent_id TEXT,
            parent_post_id TEXT, metadata TEXT, created_at TEXT
        );
        CREATE INDEX IF NOT EXISTS idx_posts_created ON graph_posts(created_at);
        CREATE TABLE IF NOT EXISTS robot_telemetry (
            id TEXT PRIMARY KEY, nwo_agent_id TEXT, agent_id TEXT,
            task_id TEXT, rl_session_id TEXT,
            joint_angles TEXT, gripper_state REAL, position TEXT,
            battery_level REAL, reward REAL,
            state_vector TEXT, action_vector TEXT, raw_telemetry TEXT,
            node_id TEXT, created_at TEXT
        );
        CREATE TABLE IF NOT EXISTS swarm_events (
            id TEXT PRIMARY KEY, swarm_id TEXT, event_type TEXT,
            payload TEXT, node_id TEXT, created_at TEXT
        );
        CREATE TABLE IF NOT EXISTS iot_snapshots (
            id TEXT PRIMARY KEY, device_id TEXT, status TEXT,
            node_id TEXT, created_at TEXT
        );
        CREATE TABLE IF NOT EXISTS expansion_queue (
            id TEXT PRIMARY KEY, node_id TEXT, model_used TEXT,
            prompt_used TEXT, result_raw TEXT, status TEXT DEFAULT 'pending',
            error_msg TEXT, started_at TEXT, completed_at TEXT, created_at TEXT
        );
        """)
        await self._db.commit()

    # ---------------------------------------------------------------
    # Agents
    # ---------------------------------------------------------------
    async def create_agent(self, data: dict) -> dict:
        if USE_SUPABASE: return await self._sb_post("agents", data)
        data.setdefault("id", new_id())
        data.setdefault("created_at", now_iso())
        if "capabilities" in data and isinstance(data["capabilities"], list):
            data["capabilities"] = json.dumps(data["capabilities"])
        cols = list(data.keys())
        vals = list(data.values())
        q = f"INSERT INTO agents ({','.join(cols)}) VALUES ({','.join('?' * len(vals))})"
        async with self._db.execute(q, vals) as cur:
            pass
        await self._db.commit()
        return data

    async def get_agent_by_key(self, key: str):
        if USE_SUPABASE: return await self._sb_get_one("agents", f"nwo_api_key=eq.{key}&is_active=eq.true")
        async with self._db.execute("SELECT * FROM agents WHERE nwo_api_key=? AND is_active=1", [key]) as cur:
            row = await cur.fetchone()
            return _row_to_dict(row, cur) if row else None

    async def get_agent_by_nwo_id(self, nwo_id: str):
        if USE_SUPABASE: return await self._sb_get_one("agents", f"nwo_agent_id=eq.{nwo_id}")
        async with self._db.execute("SELECT * FROM agents WHERE nwo_agent_id=?", [nwo_id]) as cur:
            row = await cur.fetchone()
            return _row_to_dict(row, cur) if row else None

    async def get_agent_by_name(self, name: str):
        if USE_SUPABASE: return await self._sb_get_one("agents", f"name=eq.{name}")
        async with self._db.execute("SELECT * FROM agents WHERE name=?", [name]) as cur:
            row = await cur.fetchone()
            return _row_to_dict(row, cur) if row else None

    async def get_robots(self, active_only=True, has_nwo_id=False):
        if USE_SUPABASE:
            qs = "agent_type=neq.ai_agent"
            if active_only: qs += "&is_active=eq.true"
            rows = await self._sb_get_many("agents", qs)
            return rows if isinstance(rows, list) else []
        q = "SELECT * FROM agents WHERE agent_type != 'ai_agent'"
        if active_only: q += " AND is_active=1"
        if has_nwo_id: q += " AND nwo_agent_id IS NOT NULL"
        async with self._db.execute(q) as cur:
            rows = await cur.fetchall()
            return _rows_to_list(rows, cur)

    # ---------------------------------------------------------------
    # Nodes
    # ---------------------------------------------------------------
    async def create_node(self, data: dict) -> dict:
        if USE_SUPABASE: return await self._sb_post("graph_nodes", data)
        data.setdefault("id", new_id())
        data.setdefault("created_at", now_iso())
        data.setdefault("updated_at", now_iso())
        for k in ("robot_position", "sensor_data"):
            if k in data and not isinstance(data[k], str):
                data[k] = json.dumps(data[k])
        data["expand_requested"] = 1 if data.get("expand_requested") else 0
        data["expand_done"] = 1 if data.get("expand_done") else 0
        cols = list(data.keys())
        vals = list(data.values())
        q = f"INSERT INTO graph_nodes ({','.join(cols)}) VALUES ({','.join('?' * len(vals))})"
        await self._db.execute(q, vals)
        await self._db.commit()
        return data

    async def get_nodes(self, limit=500) -> list:
        if USE_SUPABASE:
            rows = await self._sb_get_many("graph_nodes", f"order=created_at.desc&limit={limit}")
            return rows if isinstance(rows, list) else []
        async with self._db.execute(
            "SELECT * FROM graph_nodes ORDER BY created_at DESC LIMIT ?", [limit]
        ) as cur:
            rows = await cur.fetchall()
            result = _rows_to_list(rows, cur)
        for n in result:
            for k in ("robot_position", "sensor_data"):
                if isinstance(n.get(k), str):
                    try: n[k] = json.loads(n[k])
                    except: pass
        return result

    async def get_pending_expansions(self, limit=5) -> list:
        if USE_SUPABASE:
            rows = await self._sb_get_many("graph_nodes",
                f"expand_requested=eq.true&expand_done=eq.false&depth_level=lt.6&limit={limit}&order=created_at.asc")
            return rows if isinstance(rows, list) else []
        async with self._db.execute(
            "SELECT * FROM graph_nodes WHERE expand_requested=1 AND expand_done=0 AND depth_level < 6 ORDER BY created_at ASC LIMIT ?",
            [limit]
        ) as cur:
            rows = await cur.fetchall()
            return _rows_to_list(rows, cur)

    async def queue_expansion(self, node_id: str) -> bool:
        if USE_SUPABASE:
            await self._sb_patch("graph_nodes", node_id, {
                "expand_requested": True,
                "expand_done": False
            })
            return True
        await self._db.execute(
            "UPDATE graph_nodes SET expand_requested=1, expand_done=0 WHERE id=?",
            [node_id]
        )
        await self._db.commit()
        return True

    async def get_node_children(self, node_id: str) -> list:
        if USE_SUPABASE:
            links = await self._sb_get_many("graph_links",
                f"source_id=eq.{node_id}&select=target_id")
            if not isinstance(links, list) or not links:
                return []
            target_ids = [l["target_id"] for l in links if isinstance(l, dict) and "target_id" in l]
            if not target_ids:
                return []
            id_filter = ",".join(target_ids)
            rows = await self._sb_get_many("graph_nodes",
                f"id=in.({id_filter})&select=id,name")
            return rows if isinstance(rows, list) else []
        async with self._db.execute(
            "SELECT n.id, n.name FROM graph_nodes n "
            "JOIN graph_links l ON l.target_id = n.id "
            "WHERE l.source_id = ?",
            [node_id]
        ) as cur:
            rows = await cur.fetchall()
            return [{"id": r[0], "name": r[1]} for r in rows]

    async def mark_node_expanded(self, node_id: str):
        if USE_SUPABASE: return await self._sb_patch("graph_nodes", node_id, {"expand_done": True})
        await self._db.execute("UPDATE graph_nodes SET expand_done=1 WHERE id=?", [node_id])
        await self._db.commit()

    # ---------------------------------------------------------------
    # Links
    # ---------------------------------------------------------------
    async def create_link(self, data: dict) -> dict:
        if USE_SUPABASE: return await self._sb_post("graph_links", data)
        data.setdefault("id", new_id())
        data.setdefault("created_at", now_iso())
        cols, vals = list(data.keys()), list(data.values())
        await self._db.execute(f"INSERT INTO graph_links ({','.join(cols)}) VALUES ({','.join('?' * len(vals))})", vals)
        await self._db.commit()
        return data

    async def get_links(self, limit=1000) -> list:
        if USE_SUPABASE:
            rows = await self._sb_get_many("graph_links", f"limit={limit}")
            return rows if isinstance(rows, list) else []
        async with self._db.execute("SELECT * FROM graph_links LIMIT ?", [limit]) as cur:
            rows = await cur.fetchall()
            return _rows_to_list(rows, cur)

    # ---------------------------------------------------------------
    # Posts
    # ---------------------------------------------------------------
    async def create_post(self, data: dict) -> dict:
        if USE_SUPABASE: return await self._sb_post("graph_posts", data)
        data.setdefault("id", new_id())
        data.setdefault("created_at", now_iso())
        if "metadata" in data and not isinstance(data["metadata"], str):
            data["metadata"] = json.dumps(data["metadata"])
        cols, vals = list(data.keys()), list(data.values())
        await self._db.execute(f"INSERT INTO graph_posts ({','.join(cols)}) VALUES ({','.join('?' * len(vals))})", vals)
        await self._db.commit()
        return data

    async def get_posts(self, since=None, limit=50) -> list:
        if USE_SUPABASE:
            qs = f"order=created_at.desc&limit={limit}"
            if since: qs += f"&created_at=gt.{since}"
            rows = await self._sb_get_many("graph_posts", qs)
            return rows if isinstance(rows, list) else []
        q = "SELECT * FROM graph_posts"
        params = []
        if since:
            q += " WHERE created_at > ?"
            params.append(since)
        q += " ORDER BY created_at DESC LIMIT ?"
        params.append(limit)
        async with self._db.execute(q, params) as cur:
            rows = await cur.fetchall()
            result = _rows_to_list(rows, cur)
        for p in result:
            if isinstance(p.get("metadata"), str):
                try: p["metadata"] = json.loads(p["metadata"])
                except: pass
        return result

    # ---------------------------------------------------------------
    # Telemetry
    # ---------------------------------------------------------------
    async def create_telemetry(self, data: dict) -> dict:
        if USE_SUPABASE: return await self._sb_post("robot_telemetry", data)
        data.setdefault("id", new_id())
        data.setdefault("created_at", now_iso())
        for k in ("joint_angles", "state_vector", "action_vector", "position", "raw_telemetry"):
            if k in data and not isinstance(data[k], str):
                data[k] = json.dumps(data[k])
        cols, vals = list(data.keys()), list(data.values())
        await self._db.execute(f"INSERT INTO robot_telemetry ({','.join(cols)}) VALUES ({','.join('?' * len(vals))})", vals)
        await self._db.commit()
        return data

    async def get_telemetry(self, nwo_agent_id: str, limit=20) -> list:
        if USE_SUPABASE:
            rows = await self._sb_get_many("robot_telemetry", f"nwo_agent_id=eq.{nwo_agent_id}&order=created_at.desc&limit={limit}")
            return rows if isinstance(rows, list) else []
        async with self._db.execute(
            "SELECT * FROM robot_telemetry WHERE nwo_agent_id=? ORDER BY created_at DESC LIMIT ?",
            [nwo_agent_id, limit]
        ) as cur:
            rows = await cur.fetchall()
            return _rows_to_list(rows, cur)

    # ---------------------------------------------------------------
    # Swarm + IoT
    # ---------------------------------------------------------------
    async def create_swarm_event(self, data: dict) -> dict:
        if USE_SUPABASE: return await self._sb_post("swarm_events", data)
        data.setdefault("id", new_id())
        data.setdefault("created_at", now_iso())
        if "payload" in data and not isinstance(data["payload"], str):
            data["payload"] = json.dumps(data["payload"])
        cols, vals = list(data.keys()), list(data.values())
        await self._db.execute(f"INSERT INTO swarm_events ({','.join(cols)}) VALUES ({','.join('?' * len(vals))})", vals)
        await self._db.commit()
        return data

    async def create_iot_snapshot(self, data: dict) -> dict:
        if USE_SUPABASE: return await self._sb_post("iot_snapshots", data)
        data.setdefault("id", new_id())
        data.setdefault("created_at", now_iso())
        if "status" in data and not isinstance(data["status"], str):
            data["status"] = json.dumps(data["status"])
        cols, vals = list(data.keys()), list(data.values())
        await self._db.execute(f"INSERT INTO iot_snapshots ({','.join(cols)}) VALUES ({','.join('?' * len(vals))})", vals)
        await self._db.commit()
        return data

    # ---------------------------------------------------------------
    # Expansion queue
    # ---------------------------------------------------------------
    async def create_expansion_queue_entry(self, data: dict) -> dict:
        if USE_SUPABASE: return await self._sb_post("expansion_queue", data)
        data.setdefault("id", new_id())
        data.setdefault("created_at", now_iso())
        cols, vals = list(data.keys()), list(data.values())
        await self._db.execute(f"INSERT INTO expansion_queue ({','.join(cols)}) VALUES ({','.join('?' * len(vals))})", vals)
        await self._db.commit()
        return data

    async def update_expansion_queue(self, eq_id: str, data: dict):
        if USE_SUPABASE: return await self._sb_patch("expansion_queue", eq_id, data)
        sets = ", ".join(f"{k}=?" for k in data)
        vals = list(data.values()) + [eq_id]
        await self._db.execute(f"UPDATE expansion_queue SET {sets} WHERE id=?", vals)
        await self._db.commit()

    # ---------------------------------------------------------------
    # Supabase REST helpers (hardened — detect errors and log them)
    # ---------------------------------------------------------------
    async def _sb_post(self, table, data):
        import aiohttp
        url = f"{SUPABASE_URL}/rest/v1/{table}"
        headers = {
            "apikey": SUPABASE_KEY, "Authorization": f"Bearer {SUPABASE_KEY}",
            "Content-Type": "application/json", "Prefer": "return=representation"
        }
        try:
            async with aiohttp.ClientSession() as sess:
                async with sess.post(url, json=data, headers=headers) as resp:
                    result = await resp.json()
                    if resp.status >= 400:
                        log.error(f"Supabase POST {table} failed: HTTP {resp.status} body={result}")
                        return {}
                    return result[0] if isinstance(result, list) and result else (result if isinstance(result, dict) else {})
        except Exception as e:
            log.error(f"Supabase POST {table} exception: {e}")
            return {}

    async def _sb_get_many(self, table, qs=""):
        """Always returns a list. Logs and returns [] on any error."""
        import aiohttp
        url = f"{SUPABASE_URL}/rest/v1/{table}?{qs}"
        headers = {"apikey": SUPABASE_KEY, "Authorization": f"Bearer {SUPABASE_KEY}"}
        try:
            async with aiohttp.ClientSession() as sess:
                async with sess.get(url, headers=headers) as resp:
                    result = await resp.json()
                    if resp.status >= 400:
                        log.error(f"Supabase GET {table} failed: HTTP {resp.status} body={result} url={url}")
                        return []
                    if not isinstance(result, list):
                        log.error(f"Supabase GET {table} returned non-list: {result} url={url}")
                        return []
                    return result
        except Exception as e:
            log.error(f"Supabase GET {table} exception: {e}")
            return []

    async def _sb_get_one(self, table, qs):
        rows = await self._sb_get_many(table, qs + "&limit=1")
        return rows[0] if rows else None

    async def _sb_patch(self, table, row_id, data):
        import aiohttp
        url = f"{SUPABASE_URL}/rest/v1/{table}?id=eq.{row_id}"
        headers = {
            "apikey": SUPABASE_KEY, "Authorization": f"Bearer {SUPABASE_KEY}",
            "Content-Type": "application/json"
        }
        try:
            async with aiohttp.ClientSession() as sess:
                async with sess.patch(url, json=data, headers=headers) as resp:
                    if resp.status >= 400:
                        body = await resp.text()
                        log.error(f"Supabase PATCH {table} failed: HTTP {resp.status} body={body[:400]}")
        except Exception as e:
            log.error(f"Supabase PATCH {table} exception: {e}")


db = GraphDB()