src/hf_mcp_agent/database.py

from __future__ import annotations

import os
import sqlite3
from pathlib import Path
from typing import Any


WORKSPACE_ROOT = Path(
    os.getenv(
        "HF_MCP_AGENT_WORKSPACE",
        str(Path.home() / "huggingface-mcp-agent-workspace"),
    )
).resolve()
DB_PATH = WORKSPACE_ROOT / "workspace.sqlite3"


SCHEMA = """
CREATE TABLE IF NOT EXISTS projects (
  id INTEGER PRIMARY KEY AUTOINCREMENT,
  name TEXT NOT NULL UNIQUE,
  kind TEXT NOT NULL DEFAULT 'project',
  description TEXT NOT NULL DEFAULT '',
  created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP,
  updated_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP
);

CREATE TABLE IF NOT EXISTS files (
  id INTEGER PRIMARY KEY AUTOINCREMENT,
  project_name TEXT,
  path TEXT NOT NULL UNIQUE,
  kind TEXT NOT NULL,
  size INTEGER NOT NULL DEFAULT 0,
  created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP,
  updated_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP
);

CREATE TABLE IF NOT EXISTS exports (
  id INTEGER PRIMARY KEY AUTOINCREMENT,
  project_name TEXT,
  path TEXT NOT NULL,
  format TEXT NOT NULL,
  size INTEGER NOT NULL DEFAULT 0,
  created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP
);
"""


def connect() -> sqlite3.Connection:
    WORKSPACE_ROOT.mkdir(parents=True, exist_ok=True)
    conn = sqlite3.connect(DB_PATH)
    conn.row_factory = sqlite3.Row
    return conn


def init_database() -> dict[str, Any]:
    with connect() as conn:
        conn.executescript(SCHEMA)
    return {"initialized": True, "database": str(DB_PATH)}


def row_to_dict(row: sqlite3.Row) -> dict[str, Any]:
    return {key: row[key] for key in row.keys()}


def record_project(name: str, kind: str = "project", description: str = "") -> dict[str, Any]:
    init_database()
    with connect() as conn:
        conn.execute(
            """
            INSERT INTO projects (name, kind, description)
            VALUES (?, ?, ?)
            ON CONFLICT(name) DO UPDATE SET
              kind = excluded.kind,
              description = excluded.description,
              updated_at = CURRENT_TIMESTAMP
            """,
            (name, kind, description),
        )
        row = conn.execute("SELECT * FROM projects WHERE name = ?", (name,)).fetchone()
    return row_to_dict(row)


def record_file(path: str, kind: str, size: int = 0, project_name: str | None = None) -> dict[str, Any]:
    init_database()
    project = project_name or infer_project_name(path)
    with connect() as conn:
        conn.execute(
            """
            INSERT INTO files (project_name, path, kind, size)
            VALUES (?, ?, ?, ?)
            ON CONFLICT(path) DO UPDATE SET
              project_name = excluded.project_name,
              kind = excluded.kind,
              size = excluded.size,
              updated_at = CURRENT_TIMESTAMP
            """,
            (project, path, kind, size),
        )
        row = conn.execute("SELECT * FROM files WHERE path = ?", (path,)).fetchone()
    return row_to_dict(row)


def record_export(path: str, file_format: str, size: int = 0, project_name: str | None = None) -> dict[str, Any]:
    init_database()
    project = project_name or infer_project_name(path)
    with connect() as conn:
        conn.execute(
            "INSERT INTO exports (project_name, path, format, size) VALUES (?, ?, ?, ?)",
            (project, path, file_format, size),
        )
        row = conn.execute("SELECT * FROM exports ORDER BY id DESC LIMIT 1").fetchone()
    return row_to_dict(row)


def list_projects(limit: int = 100) -> list[dict[str, Any]]:
    init_database()
    with connect() as conn:
        rows = conn.execute(
            "SELECT * FROM projects ORDER BY updated_at DESC LIMIT ?",
            (max(1, min(int(limit), 500)),),
        ).fetchall()
    return [row_to_dict(row) for row in rows]


def list_file_records(project_name: str = "", limit: int = 200) -> list[dict[str, Any]]:
    init_database()
    with connect() as conn:
        if project_name:
            rows = conn.execute(
                "SELECT * FROM files WHERE project_name = ? ORDER BY updated_at DESC LIMIT ?",
                (project_name, max(1, min(int(limit), 1000))),
            ).fetchall()
        else:
            rows = conn.execute(
                "SELECT * FROM files ORDER BY updated_at DESC LIMIT ?",
                (max(1, min(int(limit), 1000)),),
            ).fetchall()
    return [row_to_dict(row) for row in rows]


def list_exports(project_name: str = "", limit: int = 100) -> list[dict[str, Any]]:
    init_database()
    with connect() as conn:
        if project_name:
            rows = conn.execute(
                "SELECT * FROM exports WHERE project_name = ? ORDER BY created_at DESC LIMIT ?",
                (project_name, max(1, min(int(limit), 500))),
            ).fetchall()
        else:
            rows = conn.execute(
                "SELECT * FROM exports ORDER BY created_at DESC LIMIT ?",
                (max(1, min(int(limit), 500)),),
            ).fetchall()
    return [row_to_dict(row) for row in rows]


def search_records(query: str, limit: int = 200) -> dict[str, list[dict[str, Any]]]:
    init_database()
    pattern = f"%{query}%"
    capped = max(1, min(int(limit), 500))
    with connect() as conn:
        projects = conn.execute(
            """
            SELECT * FROM projects
            WHERE name LIKE ? OR description LIKE ? OR kind LIKE ?
            ORDER BY updated_at DESC LIMIT ?
            """,
            (pattern, pattern, pattern, capped),
        ).fetchall()
        files = conn.execute(
            """
            SELECT * FROM files
            WHERE path LIKE ? OR kind LIKE ? OR project_name LIKE ?
            ORDER BY updated_at DESC LIMIT ?
            """,
            (pattern, pattern, pattern, capped),
        ).fetchall()
        exports = conn.execute(
            """
            SELECT * FROM exports
            WHERE path LIKE ? OR format LIKE ? OR project_name LIKE ?
            ORDER BY created_at DESC LIMIT ?
            """,
            (pattern, pattern, pattern, capped),
        ).fetchall()
    return {
        "projects": [row_to_dict(row) for row in projects],
        "files": [row_to_dict(row) for row in files],
        "exports": [row_to_dict(row) for row in exports],
    }


def database_summary() -> dict[str, Any]:
    init_database()
    with connect() as conn:
        project_count = conn.execute("SELECT COUNT(*) AS count FROM projects").fetchone()["count"]
        file_count = conn.execute("SELECT COUNT(*) AS count FROM files").fetchone()["count"]
        export_count = conn.execute("SELECT COUNT(*) AS count FROM exports").fetchone()["count"]
    return {
        "database": str(DB_PATH),
        "projects": project_count,
        "files": file_count,
        "exports": export_count,
    }


def infer_project_name(path: str) -> str | None:
    first = path.strip("/").split("/", 1)[0]
    return first or None