"""SQLite persistence layer"""

from __future__ import annotations

import json
import os
import sqlite3
import threading
import time
import uuid
import zlib
from abc import ABC, abstractmethod
from contextlib import contextmanager
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from hashlib import sha256
from typing import Any, Callable, Dict, List, Optional, Tuple

from backend.config import get_config
from backend.constants import DRAFT_ZONE, PUBLISHED_ZONE
from backend.models.agents import AgentCreate, AgentUpdate
from backend.utils.events_utils import get_artifacts_metadata
from backend.utils.logging_utils import get_logger
from backend.utils.project_utils import get_agent_details, get_ua_project
from backend.utils.user_utils import get_user_info
from backend.utils.utils import _has_unpublished_changes, invalidate_request_cache, request_cached

logger = get_logger(__name__)


@dataclass(frozen=True)
class Migration:
    to: int
    name: str
    script: Optional[str] = None
    func: Optional[Callable[[sqlite3.Connection], None]] = None

    def __post_init__(self):
        if (self.script is None) == (self.func is None):
            raise ValueError("Provide exactly one of: script or func")


MIGRATIONS: list[Migration] = []
assert all(MIGRATIONS[i - 1].to < MIGRATIONS[i].to for i in range(1, len(MIGRATIONS)))


class Store(ABC):
    """Abstract base class for storage implementations"""

    @abstractmethod
    def get_all_agents(self) -> List[Dict[str, Any]]: ...

    @abstractmethod
    def get_agents_by_owner(self, owner: str) -> List[Dict[str, Any]]: ...

    @abstractmethod
    def create_agent(self, agent_dict: Dict[str, Any]) -> Dict[str, Any]: ...

    @abstractmethod
    def update_agent(self, agent_id: str, update_data: Dict[str, Any]) -> Dict[str, Any]: ...

    @abstractmethod
    def delete_agent(self, agent_id: str) -> bool: ...

    @abstractmethod
    def get_agent(self, agent_id: str) -> Optional[Dict[str, Any]]: ...

    @abstractmethod
    def get_conversations(self) -> Dict[str, Dict[str, Any]]: ...

    @abstractmethod
    def get_conversations_by_user(self, user_id: str) -> Dict[str, Dict[str, Any]]: ...

    @abstractmethod
    def get_conversations_ids_by_user(self, user_id: str) -> List[str]: ...

    @abstractmethod
    def get_conversation(self, conversation_id: str, user_id: Optional[str] = None) -> Optional[Dict[str, Any]]: ...

    @abstractmethod
    def get_conversations_metadata(self, user_id: str) -> List[Dict[str, Any]]: ...

    @abstractmethod
    def update_conversation(self, conversation_id: str, conversation_obj: Dict[str, Any]) -> None: ...

    @abstractmethod
    def delete_conversation(self, conversation_id: str, permanent_delete: bool = False) -> bool: ...

    @abstractmethod
    def ensure_conversation_exists(self, conv_id: str, user_id: str, agent_ids: List[str]) -> None: ...

    @abstractmethod
    def append_message(self, conversation_id: str, message: Dict[str, Any]) -> None: ...

    @abstractmethod
    def update_message(self, message_id: str, updates: dict) -> None: ...

    @abstractmethod
    def append_messages(self, conversation_id: str, messages: List[Dict[str, Any]]) -> None: ...

    @abstractmethod
    def get_message(self, message_id: str, user_id: str) -> Dict[str, Any]: ...

    @abstractmethod
    def get_message_artifacts_meta(self, message_id: str, user_id: str) -> Dict[str, Any]: ...

    @abstractmethod
    def update_conversation_meta(
            self, conversation_id: str, *, title: Optional[str] = None, agent_ids: Optional[List[str]] = None
    ) -> None: ...

    @abstractmethod
    def get_message_events(self, message_id: str, user_id: Optional[str] = None) -> List[Dict[str, Any]]: ...

    @abstractmethod
    def get_preferences(self, user_id: str) -> Dict[str, Any]: ...

    @abstractmethod
    def update_preferences(self, user_id: str, prefs: Dict[str, Any]) -> Dict[str, Any]: ...

    @abstractmethod
    def get_draft_conversation(self, agent_id: str, user_id: str) -> Dict[str, Any]: ...

    @abstractmethod
    def upsert_draft_conversation(self, agent_id: str, convo: Dict[str, Any], user_id: str) -> None: ...

    @abstractmethod
    def delete_draft_conversation(self, agent_id: str, user_id: str) -> None: ...

    @abstractmethod
    def replace_agent_shares(self, agent_id: str, shares: List[Dict[str, str]]) -> None: ...

    @abstractmethod
    def get_agent_shares(self, agent_id: str) -> List[Dict[str, str]]: ...

    @abstractmethod
    def get_agents_shared_with(self, user_id: str, groups: List[str]) -> List[Dict[str, Any]]: ...

    @abstractmethod
    def get_share_counts(self, agent_ids: List[str]) -> Dict[str, int]: ...

    @abstractmethod
    def agent_exists(self, agent_id: str) -> bool: ...

    # --- Dashboard analytics ---
    @abstractmethod
    def analytics_counts(self, owner: str) -> Dict[str, int]: ...

    @abstractmethod
    def analytics_shared_users(self, owner: str) -> int: ...

    # Aross all user agents (all owners)
    @abstractmethod
    def analytics_shared_users_for_all_user_agents(self) -> int: ...

    # Number of *agents* that are shared (≥1 share), for an owner
    @abstractmethod
    def analytics_shared_agents(self, owner: str) -> int: ...

    # Number of *agents* that are shared (≥1 share), across all owners
    @abstractmethod
    def analytics_shared_agents_for_all_user_agents(self) -> int: ...

    # Counts across all user agents (not scoped to a single owner)
    @abstractmethod
    def analytics_counts_all_user_agents(self) -> Dict[str, int]: ...

    # Helper for /users/agent-owners
    @abstractmethod
    def list_user_agent_owners(self) -> List[Dict[str, Any]]: ...

    @abstractmethod
    def analytics_usage_buckets(
            self,
            owner: str,
            start: Optional[str],
            end: Optional[str],
            bucket: str,
            agent_id: Optional[str],
            *,
            is_project_admin: bool = False,
            agent_type: Optional[str] = None,  # "all" | "user" | "enterprise"
            owner_id: Optional[str] = None,
            enterprise_ids: Optional[List[str]] = None,
    ) -> List[Dict[str, Any]]: ...

    @abstractmethod
    def analytics_feedback_counts(
            self,
            owner: str,
            start: Optional[str],
            end: Optional[str],
            agent_id: Optional[str],
            *,
            is_project_admin: bool = False,
            agent_type: Optional[str] = None,
            owner_id: Optional[str] = None,
            enterprise_ids: Optional[List[str]] = None,
    ) -> Dict[str, int]: ...

    @abstractmethod
    def analytics_activity(
            self,
            owner: str,
            start: Optional[str],
            end: Optional[str],
            agent_id: Optional[str],
            limit: int,
            offset: int,
            *,
            is_project_admin: bool = False,
            agent_type: Optional[str] = None,
            owner_id: Optional[str] = None,
            enterprise_ids: Optional[List[str]] = None,
            q: Optional[str] = None,
            sort_by: Optional[str] = None,
            sort_dir: Optional[str] = None,
            group_by: Optional[str] = None,
    ) -> Tuple[List[Dict[str, Any]], int]: ...

    @abstractmethod
    def analytics_active_users_buckets(
            self,
            owner: str,
            start: Optional[str],
            end: Optional[str],
            bucket: str,
            agent_id: Optional[str],
            *,
            is_project_admin: bool = False,
            agent_type: Optional[str] = None,
            owner_id: Optional[str] = None,
            enterprise_ids: Optional[List[str]] = None,
    ) -> List[Dict[str, Any]]: ...

    # --- Message feedback (column-backed) ---
    @abstractmethod
    def update_message_feedback(
            self, message_id: str, *, rating: Optional[int], text: Optional[str], by: Optional[str]
    ) -> None: ...
    @abstractmethod
    def clear_message_feedback(self, message_id: str) -> None: ...

    @abstractmethod
    def delete_message(self, message_id: str) -> bool: ...

    @abstractmethod
    def get_message_trace(self, message_id: str, user_id: Optional[str] = None) -> str: ...


class SQLiteStore(Store):
    """SQLite implementation of Store interface"""

    def __init__(self, db_path: str):
        self.db_path = db_path
        self.timeout = int(os.getenv("SQLITE_TIMEOUT", "30"))
        self._local = threading.local()
        self._connection_refs: set[sqlite3.Connection] = set()
        self._connection_threads: dict[sqlite3.Connection, str] = {}

        self._initialize_db()

    def debug_connection_count(self) -> dict:
        """Debug method to check connection status"""
        import gc
        import threading

        active_connections = len(self._connection_refs)
        sqlite_objects = len([obj for obj in gc.get_objects() if isinstance(obj, sqlite3.Connection)])

        # Get thread information for each connection
        active_threads = [self._connection_threads.get(conn, "unknown") for conn in self._connection_refs]

        return {
            "tracked_connections": active_connections,
            "total_sqlite_objects_in_memory": sqlite_objects,
            "current_thread_has_connection": hasattr(self._local, "conn"),
            "current_thread_id": threading.current_thread().ident,
            "thread_names": active_threads,
            "db_path": self.db_path,
        }

    def start_connection_monitor(self, interval_seconds=60):
        """Start a background thread to monitor connections"""
        import time

        def monitor():
            while True:
                try:
                    info = self.debug_connection_count()
                    logger.info(
                        f"SQLite Monitor: {info['tracked_connections']} tracked, "
                        f"{info['total_sqlite_objects_in_memory']} total objects"
                    )

                    # Log threads holding connections
                    if info["thread_names"]:
                        thread_info = ", ".join(info["thread_names"])
                        logger.info(f"SQLite Connections held by threads: {thread_info}")

                    time.sleep(interval_seconds)
                except Exception as e:
                    logger.error(f"Connection monitor error: {e}")
                    break

        monitor_thread = threading.Thread(target=monitor, daemon=True)
        monitor_thread.start()
        logger.info("Started SQLite connection monitor")

    @property
    def _conn(self) -> sqlite3.Connection:
        """Get thread-local database connection"""
        if not hasattr(self._local, "conn"):
            self._local.conn = self._create_connection()
        return self._local.conn

    def _create_connection(self) -> sqlite3.Connection:
        """Create a new database connection with optimized settings"""
        logger.info("Opening SQLite connection to %s (timeout=%s)", self.db_path, self.timeout)
        conn = sqlite3.connect(
            self.db_path,
            timeout=self.timeout,
            check_same_thread=False,
            isolation_level=None,  # autocommit mode
        )
        # Enable WAL mode for better concurrency
        conn.execute("PRAGMA journal_mode=WAL")
        conn.execute("PRAGMA synchronous=NORMAL")
        conn.execute("PRAGMA foreign_keys=ON")
        conn.row_factory = sqlite3.Row
        # --- Track connections for debugging ---
        # self._connection_refs.add(conn)  # Track the connection
        # thread_name = threading.current_thread().name
        # self._connection_threads[conn] = thread_name

        # # Extract number from thread name like "Thread-18" or "MainThread"
        # import re
        # thread_match = re.search(r'Thread-(\d+)', thread_name)
        # if thread_match:
        #     thread_number = thread_match.group(1)
        #     logger.debug(f"Created new SQLite connection for thread Thread-{thread_number}")
        # else:
        #     # Fallback for MainThread or other named threads
        #     logger.debug(f"Created new SQLite connection for thread {thread_name}")
        # ---------------------------------------
        return conn

    def close_thread_connection(self):
        conn = getattr(self._local, "conn", None)
        if conn is None:
            return
        try:
            conn.close()
            logger.info("Closed SQLite connection for thread %s", threading.current_thread().name)
        finally:
            self._connection_refs.discard(conn)
            self._connection_threads.pop(conn, None)
            del self._local.conn

    @contextmanager
    def _transaction(self):
        """Execute operations within a database transaction"""
        conn = self._conn
        try:
            conn.execute("BEGIN IMMEDIATE")
            yield conn
            conn.commit()
        except Exception:
            conn.rollback()
            raise

    # system_prompt, kb_description, llmid, embedding_llm_id, short_example_queries, tools
    def _initialize_db(self):
        """Initialize database schema"""
        with sqlite3.connect(self.db_path, timeout=self.timeout) as conn:
            logger.info("Initializing SQLite DB at %s", self.db_path)
            conn.execute("PRAGMA journal_mode=WAL")
            conn.execute("PRAGMA synchronous=NORMAL")
            conn.execute("PRAGMA foreign_keys=ON")
            conn.executescript("""
                               CREATE TABLE IF NOT EXISTS agents (
                                                                     id                          TEXT PRIMARY KEY,
                                                                     owner                       TEXT NOT NULL,
                                                                     description                 TEXT,
                                                                     sample_questions            TEXT DEFAULT '[]',
                                                                     documents                   TEXT DEFAULT '[]',
                                                                     indexing                    TEXT,
                                                                     published_version           TEXT,
                                                                     published_at                TEXT,
                                                                     publishing_status           TEXT,
                                                                     publishing_job_id           TEXT,
                                                                     created_at                  TEXT DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ','now')),
                                   last_modified               TEXT DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ','now'))
                                   );

                               CREATE TABLE IF NOT EXISTS conversations (
                                                                            conversation_id TEXT PRIMARY KEY,
                                                                            user_id         TEXT NOT NULL,
                                                                            title           TEXT DEFAULT '',
                                                                            agent_ids       TEXT DEFAULT '[]',
                                                                            last_updated    TEXT DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ','now')),
                                   created_at      TEXT DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ','now')),
                                   last_modified   TEXT DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ','now')),
                                   status          TEXT DEFAULT 'active' CHECK (status IN ('active','deleted'))
                                   );

                               CREATE TABLE IF NOT EXISTS messages (
                                                                       id                   TEXT PRIMARY KEY,
                                                                       conversation_id      TEXT NOT NULL,
                                                                       role                 TEXT NOT NULL,
                                                                       content              TEXT NOT NULL,
                                                                       event_log            TEXT DEFAULT '[]',
                                                                       created_at           TEXT DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ','now')),
                                   actions              TEXT DEFAULT '{}',
                                   artifacts            BLOB,                          -- compressed JSON (zlib)
                               -- selections/usage captured at write-time to avoid recomputing
                                   selected_agent_ids   TEXT DEFAULT '[]',             -- JSON array of agent ids (for UI convenience)
                                   used_agent_ids       TEXT DEFAULT '[]',             -- JSON array of agent ids actually invoked
                               -- feedback columns for efficient analytics
                                   feedback_rating      INTEGER NULL CHECK (feedback_rating IN (0,1)),
                                   feedback_text        TEXT NULL,
                                   feedback_by          TEXT NULL,
                                   feedback_updated_at  TEXT NULL,
                                   trace               TEXT DEFAULT '{}',
                                   status              TEXT DEFAULT 'active' CHECK (status IN ('active','deleted')),
                                   FOREIGN KEY(conversation_id) REFERENCES conversations(conversation_id) ON DELETE CASCADE
                                   );

                               -- Narrow relational table for analytics
                               CREATE TABLE IF NOT EXISTS message_agents (
                                                                             message_id  TEXT NOT NULL,
                                                                             agent_id    TEXT NOT NULL,
                                                                             selected    INTEGER NOT NULL CHECK (selected IN (0,1)),
                                   used        INTEGER NOT NULL CHECK (used IN (0,1)),
                                   PRIMARY KEY (message_id, agent_id),
                                   FOREIGN KEY(message_id) REFERENCES messages(id) ON DELETE CASCADE
                                   );

                               CREATE TABLE IF NOT EXISTS preferences (
                                                                          user_id       TEXT PRIMARY KEY,
                                                                          prefs         TEXT DEFAULT '{}',
                                                                          created_at    TEXT DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ','now')),
                                   last_modified TEXT DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ','now'))
                                   );

                               CREATE TABLE IF NOT EXISTS draft_conversations (
                                                                                  agent_id      TEXT NOT NULL,
                                                                                  user_id       TEXT NOT NULL,
                                                                                  convo         TEXT DEFAULT '{}',
                                                                                  created_at    TEXT DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ','now')),
                                   last_modified TEXT DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ','now')),
                                   PRIMARY KEY(agent_id, user_id)
                                   );

                               CREATE TABLE IF NOT EXISTS agent_shares (
                                                                           agent_id       TEXT NOT NULL,
                                                                           principal      TEXT NOT NULL,
                                                                           principal_type TEXT NOT NULL CHECK (principal_type IN ('user','group')),
                                   PRIMARY KEY (agent_id, principal),
                                   FOREIGN KEY(agent_id) REFERENCES agents(id) ON DELETE CASCADE
                                   );
                               CREATE TABLE IF NOT EXISTS schema_migrations (
                                                                                version      INTEGER PRIMARY KEY,
                                                                                name         TEXT NOT NULL,
                                                                                applied_at   TEXT DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ','now')),
                                   checksum     TEXT,
                                   script       TEXT,
                                   runtime_ms   INTEGER
                                   );

                               CREATE INDEX IF NOT EXISTS idx_agents_owner ON agents(owner);
                               CREATE INDEX IF NOT EXISTS idx_conversations_user ON conversations(user_id);
                               CREATE INDEX IF NOT EXISTS idx_messages_conversation ON messages(conversation_id);
                               CREATE INDEX IF NOT EXISTS idx_messages_created ON messages(created_at);
                               CREATE INDEX IF NOT EXISTS idx_messages_role ON messages(role);
                               CREATE INDEX IF NOT EXISTS idx_shares_agent ON agent_shares(agent_id);
                               CREATE INDEX IF NOT EXISTS idx_shares_principal ON agent_shares(principal);
                               CREATE INDEX IF NOT EXISTS idx_ma_agent ON message_agents(agent_id);
                               CREATE INDEX IF NOT EXISTS idx_ma_used ON message_agents(used);
                               CREATE INDEX IF NOT EXISTS idx_ma_selected ON message_agents(selected);
                               CREATE INDEX IF NOT EXISTS idx_messages_conv_created ON messages(conversation_id, created_at);
                               CREATE INDEX IF NOT EXISTS idx_conversations_user_status ON conversations(user_id, status);
                               CREATE INDEX IF NOT EXISTS idx_messages_role_created ON messages(role, created_at);
                               CREATE INDEX IF NOT EXISTS idx_ma_message ON message_agents(message_id);
                               CREATE INDEX IF NOT EXISTS idx_messages_conv_created_active ON messages(conversation_id, created_at) WHERE status='active';
                               """)

            cur_version = self._get_user_version(conn)
            logger.info("Current user_version=%d", cur_version)

            if cur_version == 0:
                logger.info("Baselining fresh DB to user_version=1")
                conn.execute(
                    "INSERT OR IGNORE INTO schema_migrations(version, name) VALUES (?, ?)",
                    (1, "baseline (initial schema)"),
                )
                self._set_user_version(conn, 1)

            self._run_migrations(conn)
            new_version = self._get_user_version(conn)
            if new_version != cur_version:
                logger.info("DB schema upgraded: %d → %d", cur_version, new_version)
            else:
                logger.info("DB schema up-to-date at user_version=%d", new_version)

    def _table_exists(self, conn, name: str) -> bool:
        row = conn.execute("SELECT 1 FROM sqlite_master WHERE type='table' AND name=?", (name,)).fetchone()
        return bool(row)

    def _column_exists(self, conn, table: str, col: str) -> bool:
        return any(c[1] == col for c in conn.execute(f"PRAGMA table_info({table})"))

    def _get_user_version(self, conn) -> int:
        row = conn.execute("PRAGMA user_version").fetchone()
        return int(row[0] or 0)

    def _set_user_version(self, conn, v: int) -> None:
        conn.execute(f"PRAGMA user_version = {v}")

    def _run_migrations(self, conn):
        current = self._get_user_version(conn)

        # High-level context
        pending = [m for m in MIGRATIONS if m.to > current]
        if pending:
            logger.info("DB '%s' user_version=%d → %d pending migration(s)", self.db_path, current, len(pending))
        else:
            logger.info("DB '%s' user_version=%d — no migrations to apply", self.db_path, current)

        for m in MIGRATIONS:
            target = int(m.to)
            if target <= current:
                continue

            logger.info("Applying migration → v%d: %s", target, m.name)
            start = time.perf_counter()
            try:
                conn.execute("BEGIN IMMEDIATE")

                if m.script is not None:
                    script = m.script.strip()
                    conn.executescript(script)
                    checksum = sha256(script.encode("utf-8")).hexdigest()
                    script_to_store = script
                else:
                    # callable path
                    m.func(conn)
                    checksum = None
                    script_to_store = None

                # Bump user_version first (source of truth)
                self._set_user_version(conn, target)
                runtime_ms = int((time.perf_counter() - start) * 1000)

                # Audit trail
                conn.execute(
                    """
                    INSERT INTO schema_migrations(version, name, checksum, script, runtime_ms)
                    VALUES (?, ?, ?, ?, ?)
                    """,
                    (target, m.name, checksum, script_to_store, runtime_ms),
                )
                conn.commit()
                logger.info("Migration v%d OK in %d ms", target, runtime_ms)
                current = target
            except Exception:
                conn.rollback()
                logger.exception("Migration v%d FAILED, rolled back", target)
                raise

    def _get_utc_now(self) -> str:
        now = datetime.utcnow()
        return now.strftime("%Y-%m-%dT%H:%M:%S") + f".{now.microsecond // 1000:03d}Z"

    def _parse_json(self, value: Optional[str], default: Any = None) -> Any:
        """Parse JSON string with fallback to default value"""
        if not value:
            return default
        try:
            return json.loads(value)
        except (json.JSONDecodeError, TypeError):
            return default

    # Agent operations

    def _agent_row_to_dict(self, row: sqlite3.Row) -> Dict[str, Any]:
        """Transform database row to agent dictionary matching API contract"""
        agent = dict(row)
        project = get_ua_project(agent)
        draft_agent_details = get_agent_details(project, DRAFT_ZONE)
        # Parse JSON fields
        agent["sample_questions"] = self._parse_json(agent["sample_questions"], [])
        agent["documents"] = self._parse_json(agent["documents"], [])
        agent["tools"] = draft_agent_details["tools"]
        agent["kb_description"] = draft_agent_details["kb_description"]
        agent["system_prompt"] = draft_agent_details["system_prompt"]
        agent["llmid"] = draft_agent_details["llm_id"]
        agent["name"] = draft_agent_details["name"]
        logger.debug(f"Agent[{agent.get('name')}] details: {draft_agent_details}")
        # Transform indexing field
        indexing_data = self._parse_json(agent.pop("indexing", None))
        if indexing_data:
            agent["indexing"] = indexing_data
        pub_details = get_agent_details(project, PUBLISHED_ZONE)
        # Transform published_version field - ensure it's a dict or None
        published_version = self._parse_json(agent.pop("published_version", None))
        if published_version and isinstance(published_version, dict):
            published_version["tools"] = pub_details["tools"]
            published_version["kb_description"] = pub_details["kb_description"]
            published_version["system_prompt"] = pub_details["system_prompt"]
            published_version["llmid"] = pub_details["llm_id"]
            published_version["name"] = pub_details["name"]
            agent["published_version"] = published_version
            logger.debug(f"Published version details: {pub_details}")
        else:
            agent["published_version"] = None

        # Rename timestamp fields to match API contract
        agent["createdAt"] = agent.pop("created_at")
        # TODO should take this from visual agent, keeping in mind doc changes wont be reflected in the vis agent last modif time ?
        agent["lastModified"] = agent.pop("last_modified")
        return agent

    @request_cached
    def get_all_agents(self):
        c = self._conn.cursor()
        result = []
        for row in c.execute("SELECT * FROM agents"):
            try:
                agent_dict = self._agent_row_to_dict(row)
                result.append(agent_dict)
            except Exception as e:
                # Silently skip this row if agent retrieval fails
                logger.error(f"error getting agent metadata: {dict(row)}\n{e}")
                pass
        return result

    @request_cached
    def get_agents_by_owner(self, owner: str):
        c = self._conn.cursor()
        return [self._agent_row_to_dict(row) for row in c.execute("SELECT * FROM agents WHERE owner = ?", (owner,))]

    @request_cached
    def get_agent(self, agent_id: str):
        logger.info(f"Fetching agent {agent_id} from database")
        c = self._conn.cursor()
        row = c.execute("SELECT * FROM agents WHERE id = ?", (agent_id,)).fetchone()
        return self._agent_row_to_dict(row) if row else None

    @invalidate_request_cache("get_agent", "get_agents_by_owner", "get_all_agents")
    def create_agent(self, agent_dict):
        if "id" not in agent_dict:
            agent_dict["id"] = str(uuid.uuid4())

        agent_dict.setdefault("sample_questions", [])

        # Validate with model
        agent_create = AgentCreate(**agent_dict)

        # Check existence
        if self.get_agent(agent_dict["id"]):
            raise ValueError(f"Agent '{agent_dict['id']}' already exists")
        with self._transaction() as conn:
            try:
                now = self._get_utc_now()
                cur = self._conn.execute(
                    """
                    INSERT OR IGNORE INTO agents (
                        id, owner, description,
                        sample_questions,
                        documents, created_at, last_modified
                    ) VALUES (?, ?, ?, ?, ?, ?, ?)
                    """,
                    (
                        agent_dict["id"],
                        agent_dict.get("owner", "unknown"),
                        agent_create.description,
                        json.dumps(agent_create.sample_questions),
                        json.dumps([]),
                        now,
                        now,
                    ),
                )
                if cur.rowcount == 0:
                    raise ValueError(f"Agent '{agent_dict['id']}' already exists")
                else:
                    logger.info(f"Created agent: {agent_dict['id']}")
                    return agent_dict
            except Exception as e:
                logger.error(f"Error creating agent: {e}")
                raise

    @invalidate_request_cache("get_agent", "get_agents_by_owner", "get_all_agents")
    def update_agent(self, agent_id, update_data):
        update_model = AgentUpdate(**update_data)

        fields = []
        values = []

        # Update simple fields
        for field in [
            "description",
        ]:
            value = getattr(update_model, field, None)
            if value is not None:
                fields.append(f"{field} = ?")
                values.append(value)

        # Update JSON fields
        if update_model.sample_questions is not None:
            fields.append("sample_questions = ?")
            values.append(json.dumps(update_model.sample_questions))

        if update_model.documents is not None:
            fields.append("documents = ?")
            values.append(json.dumps(update_model.documents))

        if update_model.indexing is not None:
            fields.append("indexing = ?")
            values.append(json.dumps(update_model.indexing))

        # Update publishing fields
        if update_model.published_at is not None:
            fields.append("published_at = ?")
            values.append(update_model.published_at)

        if update_model.publishing_status is not None:
            fields.append("publishing_status = ?")
            values.append(update_model.publishing_status)

        if update_model.publishing_job_id is not None:
            fields.append("publishing_job_id = ?")
            values.append(update_model.publishing_job_id)

        # Handle published version - ensure it's a dict
        if "published_version" in update_data:
            pv = update_data["published_version"]
            if pv and isinstance(pv, dict):
                fields.append("published_version = ?")
                values.append(json.dumps(pv))
            elif pv is None:
                fields.append("published_version = ?")
                values.append(None)

        if fields:
            fields.append("last_modified = ?")
            values.extend([self._get_utc_now(), agent_id])

            self._conn.execute(f"UPDATE agents SET {', '.join(fields)} WHERE id = ?", values)

        return self.get_agent(agent_id)

    @invalidate_request_cache("get_agent", "get_agents_by_owner", "get_all_agents", "get_share_counts")
    def delete_agent(self, agent_id):
        c = self._conn.cursor()
        c.execute("DELETE FROM agents WHERE id = ?", (agent_id,))
        return c.rowcount > 0

    def agent_exists(self, agent_id: str) -> bool:
        """Check if agent exists in database (no permission check)"""
        c = self._conn.cursor()
        row = c.execute("SELECT 1 FROM agents WHERE id = ?", (agent_id,)).fetchone()
        return row is not None

    # Conversation operations

    def _conversation_row_to_dict(self, row: sqlite3.Row, include_messages: bool = True) -> Dict[str, Any]:
        """Transform database row to conversation dictionary"""
        conv = {
            "userId": row["user_id"],
            "title": row["title"],
            "lastUpdated": row["last_updated"],
            "agentIds": self._parse_json(row["agent_ids"], []),
        }

        if include_messages:
            c = self._conn.cursor()
            messages = []
            for msg_row in c.execute(
                    "SELECT * FROM messages WHERE conversation_id = ? AND status = 'active' ORDER BY created_at",
                    (row["conversation_id"],),
            ):
                messages.append(self._message_row_to_dict(msg_row))
            conv["messages"] = messages

        return conv

    def _message_row_to_dict(self, row: sqlite3.Row) -> Dict[str, Any]:
        """Transform message row to dictionary"""
        # TODO do not read all event log
        event_log = self._parse_json(row["event_log"], [])
        # artifacts are stored compressed; tolerate empty/None
        artifacts = {}
        try:
            compressed_data = row["artifacts"]
            if compressed_data:
                json_str = zlib.decompress(compressed_data).decode("utf-8")
                artifacts = json.loads(json_str)
        except zlib.error as e:
            logger.error("Artifacts decompression failed for message %s: %s", row["id"], e)
            artifacts = {}
        except json.JSONDecodeError as e:
            logger.error("Artifacts JSON invalid for message %s: %s", row["id"], e)
            artifacts = {}
        except Exception:
            logger.exception("Unexpected error reading artifacts for message %s", row["id"])
            artifacts = {}

        max_artifacts_size_mb = get_config().get("max_artifacts_size_mb", 2)
        # max_artifacts_size_mb = 0.0001  # For testing purpose only, set to 0.0001MB = 0.1KB
        arts_meta = get_artifacts_metadata(artifacts, max_artifacts_size_mb)

        # feedback → JSON shape for API convenience
        fb = {}
        if row["feedback_rating"] is not None:
            fb = {
                "rating": int(row["feedback_rating"]),
                "by": row["feedback_by"],
                "updatedAt": row["feedback_updated_at"],
            }
            if row["feedback_text"]:
                fb["text"] = row["feedback_text"]

        return {
            "id": row["id"],
            "role": row["role"],
            "content": row["content"],
            "createdAt": row["created_at"],
            "hasEventLog": bool(event_log),
            "actions": self._parse_json(row["actions"], {}),
            "artifactsMetadata": arts_meta,
            "feedback": fb,
            "selectedAgentIds": self._parse_json(row["selected_agent_ids"], []),
            "usedAgentIds": self._parse_json(row["used_agent_ids"], []),
        }

    @request_cached
    def get_conversations(self):
        c = self._conn.cursor()
        return {
            row["conversation_id"]: self._conversation_row_to_dict(row)
            for row in c.execute("SELECT * FROM conversations WHERE status = 'active'")
        }

    @request_cached
    def get_conversations_by_user(self, user_id):
        c = self._conn.cursor()
        return {
            row["conversation_id"]: self._conversation_row_to_dict(row)
            for row in c.execute("SELECT * FROM conversations WHERE user_id = ? AND status = 'active' ", (user_id,))
        }

    @request_cached
    def get_conversations_ids_by_user(self, user_id):
        c = self._conn.cursor()
        return [
            row["conversation_id"]
            for row in c.execute(
                "SELECT conversation_id FROM conversations WHERE user_id = ?  AND status = 'active'", (user_id,)
            )
        ]

    @request_cached
    def get_conversation(self, conversation_id: str, user_id: Optional[str] = None):
        c = self._conn.cursor()

        if user_id:
            row = c.execute(
                "SELECT * FROM conversations WHERE conversation_id = ? AND user_id = ?  AND status = 'active'",
                (conversation_id, user_id),
            ).fetchone()
        else:
            row = c.execute(
                "SELECT * FROM conversations WHERE conversation_id = ?  AND status = 'active'", (conversation_id,)
            ).fetchone()

        return self._conversation_row_to_dict(row) if row else None

    @request_cached
    def get_conversations_metadata(self, user_id: str) -> list[dict]:
        c = self._conn.cursor()
        return [
            {
                "id": row["conversation_id"],
                "title": row["title"],
                "lastUpdated": row["last_updated"],
                "agentIds": self._parse_json(row["agent_ids"], []),
            }
            for row in c.execute(
                "SELECT conversation_id, title, last_updated, agent_ids FROM conversations WHERE user_id = ?  AND status = 'active'",
                (user_id,),
            )
        ]

    @invalidate_request_cache(
        "get_conversations",
        "get_conversations_by_user",
        "get_conversation",
        "get_conversations_metadata",
        "get_conversations_ids_by_user",
    )
    def update_conversation(self, conversation_id, conversation_obj):
        with self._transaction() as conn:
            c = conn.cursor()

            now = self._get_utc_now()
            c.execute(
                """
                INSERT OR REPLACE INTO conversations (
                    conversation_id, user_id, title, last_updated,
                    agent_ids, created_at, last_modified, status
                ) VALUES (
                    ?, ?, ?, ?, ?,
                    COALESCE((SELECT created_at FROM conversations WHERE conversation_id = ?), ?),
                    ?,
                    COALESCE((SELECT status FROM conversations WHERE conversation_id = ?), 'active')
                )
            """,
                (
                    conversation_id,
                    conversation_obj.get("userId"),
                    conversation_obj.get("title", ""),
                    conversation_obj.get("lastUpdated", now),
                    json.dumps(conversation_obj.get("agentIds", [])),
                    conversation_id,
                    now,
                    now,
                    conversation_id,
                ),
            )

            # Replace all messages
            c.execute("DELETE FROM messages WHERE conversation_id = ?", (conversation_id,))

            for msg in conversation_obj.get("messages", []):
                self._insert_message(c, conversation_id, msg)

    @invalidate_request_cache(
        "get_conversations",
        "get_conversations_by_user",
        "get_conversation",
        "get_conversations_metadata",
        "get_conversations_ids_by_user",
    )
    def delete_conversation(self, conversation_id, permanent_delete: bool = False):
        c = self._conn.cursor()
        if permanent_delete:
            # rely on FK ON DELETE CASCADE (messages → conversations)
            c.execute("DELETE FROM conversations WHERE conversation_id = ?", (conversation_id,))
            logger.info(
                "Permanently deleted conversation %s (cascade removed its messages), rows=%d",
                conversation_id,
                c.rowcount,
            )
        else:
            c.execute("UPDATE messages SET status = 'deleted' WHERE conversation_id = ?", (conversation_id,))
            logger.info(f"Soft-deleted messages for conversation {conversation_id} {c.rowcount} rows affected")
            c.execute("UPDATE conversations SET status = 'deleted' WHERE conversation_id = ?", (conversation_id,))
        return c.rowcount > 0

    def ensure_conversation_exists(self, conv_id: str, user_id: str, agent_ids: list[str]):
        now = self._get_utc_now()
        with self._transaction() as conn:
            conn.execute(
                """
                INSERT OR IGNORE INTO conversations (
                    conversation_id, user_id, title, last_updated,
                    agent_ids, created_at, last_modified
                    ) VALUES (?, ?, '', ?, ?, ?, ?)
                """,
                (conv_id, user_id, now, json.dumps(agent_ids), now, now),
            )

    def _insert_message(self, cursor: sqlite3.Cursor, conversation_id: str, message: dict):
        """Insert a message with its precomputed fields (storage stays dumb)."""
        events = message.get("eventLog", [])
        actions = message.get("actions") or {}
        artifacts = message.get("artifactsMetadata") or {}
        selected_ids = message.get("selectedAgentIds") or []
        used_ids = message.get("usedAgentIds") or []

        # Compress artifacts JSON (to bytes)
        artifacts_str = json.dumps(artifacts)
        compressed_artifacts = zlib.compress(artifacts_str.encode("utf-8"))

        cursor.execute(
            """
            INSERT INTO messages (
                id, conversation_id, role, content, event_log, created_at,
                actions, artifacts, selected_agent_ids, used_agent_ids
            )
            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
            """,
            (
                message["id"],
                conversation_id,
                message["role"],
                message.get("content", ""),
                json.dumps(events),
                message.get("createdAt", self._get_utc_now()),
                json.dumps(actions),
                compressed_artifacts,
                json.dumps(selected_ids),
                json.dumps(used_ids),
            ),
        )

        # Write message→agent join rows (selected/used)
        if selected_ids or used_ids:
            seen = set(selected_ids) | set(used_ids)
            cursor.executemany(
                "INSERT OR REPLACE INTO message_agents (message_id, agent_id, selected, used) VALUES (?, ?, ?, ?)",
                [
                    (
                        message["id"],
                        aid,
                        1 if aid in set(selected_ids) else 0,
                        1 if aid in set(used_ids) else 0,
                    )
                    for aid in seen
                ],
            )

    @invalidate_request_cache(
        "get_conversations", "get_conversations_by_user", "get_conversation", "get_conversations_metadata"
    )
    def append_message(self, conversation_id: str, message: dict):
        with self._transaction() as conn:
            c = conn.cursor()
            self._insert_message(c, conversation_id, message)

            now = self._get_utc_now()
            c.execute(
                "UPDATE conversations SET last_updated = ?, last_modified = ? WHERE conversation_id = ?",
                (now, now, conversation_id),
            )

    def update_message(self, message_id: str, updates: dict):
        if not updates:
            return
        ALLOWED = {
            "content",
            "event_log",
            "actions",
            "status",
            "trace",
            "artifacts",
        }
        unknown = set(updates) - ALLOWED
        if unknown:
            raise ValueError(f"Unsupported message fields: {', '.join(sorted(unknown))}")

        with self._transaction() as conn:
            c = conn.cursor()
            if "artifacts" in updates:
                arts = updates["artifacts"] or {}
                # always store compressed version (zlib) as in insert
                updates["artifacts"] = zlib.compress(json.dumps(arts).encode("utf-8"))
            # serialize json columns
            for k in ("event_log", "actions", "trace"):
                if k in updates and not isinstance(updates[k], (str, bytes)):
                    updates[k] = json.dumps(updates[k])
            if updates:
                cols = [f"{col} = ?" for col in updates]
                vals = list(updates.values()) + [message_id]
                c.execute(f"UPDATE messages SET {', '.join(cols)} WHERE id = ?", vals)

    def update_message_feedback(
            self, message_id: str, *, rating: Optional[int], text: Optional[str], by: Optional[str]
    ) -> None:
        with self._transaction() as conn:
            c = conn.cursor()
            now = self._get_utc_now()
            c.execute(
                """
                UPDATE messages
                SET feedback_rating = ?, feedback_text = ?, feedback_by = ?, feedback_updated_at = ?
                WHERE id = ?
                """,
                (rating if rating in (0, 1) else None, (text or None), by, now, message_id),
            )

    def clear_message_feedback(self, message_id: str) -> None:
        with self._transaction() as conn:
            conn.execute(
                """
                UPDATE messages
                SET feedback_rating = NULL, feedback_text = NULL, feedback_by = NULL, feedback_updated_at = NULL
                WHERE id = ?
                """,
                (message_id,),
            )

    @invalidate_request_cache(
        "get_conversations", "get_conversations_by_user", "get_conversation", "get_conversations_metadata"
    )
    def append_messages(self, conversation_id: str, messages: list[dict]) -> None:
        with self._transaction() as conn:
            c = conn.cursor()

            for msg in messages:
                self._insert_message(c, conversation_id, msg)

            now = self._get_utc_now()
            c.execute(
                "UPDATE conversations SET last_updated = ?, last_modified = ? WHERE conversation_id = ?",
                (now, now, conversation_id),
            )

    @invalidate_request_cache(
        "get_conversations", "get_conversations_by_user", "get_conversation", "get_conversations_metadata"
    )
    def update_conversation_meta(self, conversation_id: str, *, title=None, agent_ids=None):
        fields = []
        values = []

        if title is not None:
            fields.append("title = ?")
            values.append(title)

        if agent_ids is not None:
            fields.append("agent_ids = ?")
            values.append(json.dumps(agent_ids))

        if fields:
            now = self._get_utc_now()
            fields.extend(["last_updated = ?", "last_modified = ?"])
            values.extend([now, now, conversation_id])

            self._conn.execute(f"UPDATE conversations SET {', '.join(fields)} WHERE conversation_id = ?", values)

    def get_message(self, message_id: str, user_id: Optional[str] = None) -> Dict[str, Any]:
        c = self._conn.cursor()
        msg_row = None
        if user_id:
            msg_row = c.execute(
                """
                SELECT * FROM messages m
                                  JOIN conversations c ON m.conversation_id = c.conversation_id
                WHERE m.id = ? AND c.user_id = ? AND m.status = 'active'
                """,
                (message_id, user_id),
            ).fetchone()

            if not msg_row:
                return {}
            msg = self._message_row_to_dict(msg_row)
            return msg
        return {}

    def get_message_artifacts_meta(self, message_id: str, user_id: Optional[str] = None) -> Dict[str, Any]:
        c = self._conn.cursor()
        row = None
        if user_id:
            row = c.execute(
                """
                SELECT m.artifacts
                FROM messages m
                         JOIN conversations c ON m.conversation_id = c.conversation_id
                WHERE m.id = ? AND c.user_id = ?
                """,
                (message_id, user_id),
            ).fetchone()
        artifacts = {}
        try:
            compressed_data = row[0] if row else None
            if compressed_data:
                json_str = zlib.decompress(compressed_data).decode("utf-8")
                artifacts = json.loads(json_str)
        except Exception:
            artifacts = {}
        return artifacts

    @request_cached
    def get_message_events(self, message_id: str, user_id: Optional[str] = None) -> list:
        c = self._conn.cursor()

        # Optional ACL check
        if user_id:
            authorized = c.execute(
                """
                SELECT 1 FROM messages m
                                  JOIN conversations c ON m.conversation_id = c.conversation_id
                WHERE m.id = ? AND c.user_id = ?
                """,
                (message_id, user_id),
            ).fetchone()

            if not authorized:
                return []

        # Get event log from message
        row = c.execute("SELECT event_log FROM messages WHERE id = ?", (message_id,)).fetchone()
        return self._parse_json(row["event_log"], []) if row else []

    @invalidate_request_cache(
        "get_conversations",
        "get_conversations_by_user",
        "get_conversation",
        "get_conversations_metadata",
    )
    def delete_message(self, message_id: str) -> bool:
        """Delete a specific message by ID"""
        with self._transaction() as conn:
            c = conn.cursor()

            # Get conversation_id before deleting the message
            conv_row = c.execute("SELECT conversation_id FROM messages WHERE id = ?", (message_id,)).fetchone()

            if not conv_row:
                return False

            conversation_id = conv_row["conversation_id"]

            # Delete the message
            c.execute("DELETE FROM messages WHERE id = ?", (message_id,))
            deleted = c.rowcount > 0

            # Update conversation timestamp if message was deleted
            if deleted:
                now = self._get_utc_now()
                c.execute(
                    "UPDATE conversations SET last_updated = ?, last_modified = ? WHERE conversation_id = ?",
                    (now, now, conversation_id),
                )

            return deleted

    # Preferences operations

    @request_cached
    def get_preferences(self, user_id: str) -> dict:
        c = self._conn.cursor()
        row = c.execute("SELECT prefs FROM preferences WHERE user_id = ?", (user_id,)).fetchone()
        return self._parse_json(row["prefs"], {}) if row else {}

    @invalidate_request_cache("get_preferences")
    def update_preferences(self, user_id: str, prefs: dict) -> dict:
        """Merge new preferences with existing ones"""
        current = self.get_preferences(user_id)
        current.update(prefs or {})

        now = self._get_utc_now()
        self._conn.execute(
            """
            INSERT OR REPLACE INTO preferences (user_id, prefs, created_at, last_modified)
            VALUES (?, ?, 
                    COALESCE((SELECT created_at FROM preferences WHERE user_id = ?), ?),
                    ?)
        """,
            (user_id, json.dumps(current), user_id, now, now),
        )

        return current

    # Draft conversation operations

    @request_cached
    def get_draft_conversation(self, agent_id: str, user_id: str) -> dict:
        c = self._conn.cursor()
        row = c.execute(
            "SELECT convo FROM draft_conversations WHERE agent_id = ? AND user_id = ?", (agent_id, user_id)
        ).fetchone()
        return self._parse_json(row["convo"], {}) if row else {}

    @invalidate_request_cache("get_draft_conversation")
    def upsert_draft_conversation(self, agent_id: str, convo: dict, user_id: str) -> None:
        now = self._get_utc_now()
        self._conn.execute(
            """
            INSERT OR REPLACE INTO draft_conversations (agent_id, user_id, convo, created_at, last_modified)
            VALUES (?, ?, ?, 
                    COALESCE((SELECT created_at FROM draft_conversations WHERE agent_id = ? AND user_id = ?), ?),
                    ?)
        """,
            (agent_id, user_id, json.dumps(convo), agent_id, user_id, now, now),
        )

    @invalidate_request_cache("get_draft_conversation")
    def delete_draft_conversation(self, agent_id: str, user_id: str) -> None:
        self._conn.execute("DELETE FROM draft_conversations WHERE agent_id = ? AND user_id = ?", (agent_id, user_id))

    # Agent sharing operations

    @invalidate_request_cache("get_share_counts")
    def replace_agent_shares(self, agent_id: str, shares: list[dict]):
        """Replace all shares for an agent atomically"""
        with self._transaction() as conn:
            c = conn.cursor()
            c.execute("DELETE FROM agent_shares WHERE agent_id = ?", (agent_id,))

            if shares:
                c.executemany(
                    "INSERT INTO agent_shares (agent_id, principal, principal_type) VALUES (?, ?, ?)",
                    [(agent_id, s["principal"], s["type"]) for s in shares],
                )

    @request_cached
    def get_agent_shares(self, agent_id: str) -> list[dict]:
        c = self._conn.cursor()
        return [
            {"principal": r["principal"], "type": r["principal_type"]}
            for r in c.execute(
                "SELECT principal, principal_type FROM agent_shares WHERE agent_id = ? ORDER BY principal", (agent_id,)
            )
        ]

    @request_cached
    def get_agents_shared_with(self, user_id: str, groups: list[str]) -> list[dict]:
        """Get agents shared with user or their groups, excluding agents they own."""
        if not user_id and not groups:
            return []

        c = self._conn.cursor()

        conditions = []
        params = []

        if user_id:
            conditions.append("(s.principal_type = 'user' AND s.principal = ?)")
            params.append(user_id)

        if groups:
            placeholders = ", ".join("?" for _ in groups)
            conditions.append(f"(s.principal_type = 'group' AND s.principal IN ({placeholders}))")
            params.extend(groups)

        # The `where_components` list will be joined by AND.
        where_components = [f"({' OR '.join(conditions)})"]

        if user_id:
            where_components.append("a.owner != ?")
            params.append(user_id)

        query = f"""
            SELECT DISTINCT a.*
            FROM agents a
            JOIN agent_shares s ON a.id = s.agent_id
            WHERE {" AND ".join(where_components)}
        """
        result = []
        for row in c.execute(query, params):
            try:
                result.append(self._agent_row_to_dict(row))
            except Exception as e:
                logger.error(f"error getting agent metadata: {dict(row)}\n{e}")
                pass
        return result

    @request_cached
    def get_share_counts(self, agent_ids: list[str]) -> dict[str, int]:
        """Get number of shares for each agent"""
        if not agent_ids:
            return {}

        c = self._conn.cursor()
        placeholders = ", ".join("?" for _ in agent_ids)

        return {
            r["agent_id"]: r["cnt"]
            for r in c.execute(
                f"""
                SELECT agent_id, COUNT(*) as cnt
                FROM agent_shares
                WHERE agent_id IN ({placeholders})
                GROUP BY agent_id
            """,
                agent_ids,
            )
        }

    # ---------------- Dashboard helpers ----------------

    def _owner_agent_ids(self, owner: str) -> set[str]:
        c = self._conn.cursor()
        return {r["id"] for r in c.execute("SELECT id FROM agents WHERE owner = ?", (owner,))}

    # Build the agent-scope WHERE clause (and params) for analytics
    def _agent_scope_clause(
            self,
            *,
            owner: str,
            agent_id: Optional[str],
            is_project_admin: bool,
            agent_type: Optional[str],
            owner_id: Optional[str],
            enterprise_ids: Optional[List[str]],
    ) -> tuple[str, list]:
        # AgentId always wins
        if agent_id:
            return ("AND ma.agent_id = ?", [agent_id])

        # Non project owners: restrict to "my agents"
        if not is_project_admin:
            # Use EXISTS to avoid constructing a large list
            return ("AND EXISTS (SELECT 1 FROM agents a WHERE a.id = ma.agent_id AND a.owner = ?)", [owner])

        # Project owner paths
        agent_type = (agent_type or "all").lower()
        enterprise_ids = list(enterprise_ids or [])

        if agent_type == "enterprise":
            if not enterprise_ids:
                return ("AND 1=0", [])  # nothing to include
            placeholders = ", ".join("?" for _ in enterprise_ids)
            return (f"AND ma.agent_id IN ({placeholders})", enterprise_ids)

        if agent_type == "user":
            if owner_id:
                return (
                    "AND EXISTS (SELECT 1 FROM agents a WHERE a.id = ma.agent_id AND a.owner = ?)",
                    [owner_id],
                )
            else:
                return ("AND EXISTS (SELECT 1 FROM agents a WHERE a.id = ma.agent_id)", [])

        # agent_type == "all"
        # user agents (optionally owned by owner_id) OR enterprise ids
        if owner_id:
            user_clause = "EXISTS (SELECT 1 FROM agents a WHERE a.id = ma.agent_id AND a.owner = ?)"
            user_params: list = [owner_id]
        else:
            user_clause = "EXISTS (SELECT 1 FROM agents a WHERE a.id = ma.agent_id)"
            user_params = []

        if enterprise_ids:
            placeholders = ", ".join("?" for _ in enterprise_ids)
            ent_clause = f"ma.agent_id IN ({placeholders})"
            ent_params = enterprise_ids
            return (f"AND ({user_clause} OR {ent_clause})", user_params + ent_params)
        else:
            return (f"AND ({user_clause})", user_params)

    # Normalize a date range (in case this is accessed by API)
    def _normalize_range(self, start: Optional[str], end: Optional[str]) -> tuple[datetime, datetime]:
        def _parse(s: Optional[str]) -> Optional[datetime]:
            if not s:
                return None
            try:
                return datetime.fromisoformat(s.replace("Z", "+00:00"))
            except Exception:
                try:
                    return datetime.strptime(s[:19], "%Y-%m-%dT%H:%M:%S")
                except Exception:
                    try:
                        return datetime.strptime(s[:10], "%Y-%m-%d")
                    except Exception:
                        return None

        end_dt = _parse(end) or datetime.utcnow()
        start_dt = _parse(start) or (end_dt - timedelta(days=56))  # last 8 weeks by default
        return start_dt, end_dt

    # --- helpers for analytics date handling  ---
    @staticmethod
    def _parse_created_ts(ts: str) -> datetime:
        """
        Parse ISO timestamps coming from SQLite (sometimes with 'Z') or Python (no 'Z')
        and normalise to a timezone-naive UTC datetime so comparisons never error.
        """
        try:
            dt = datetime.fromisoformat(ts.replace("Z", "+00:00"))
        except Exception:
            dt = datetime.strptime(ts[:19], "%Y-%m-%dT%H:%M:%S")
        if getattr(dt, "tzinfo", None) is not None:
            dt = dt.astimezone(timezone.utc).replace(tzinfo=None)
        return dt

    @staticmethod
    def _day_bounds(start_dt: datetime, end_dt: datetime) -> tuple[str, str, datetime]:
        """
        Return (lower_bound, exclusive_upper_bound, end_exclusive_dt).
        Bounds are YYYY-MM-DD strings so SQLite can use the index on messages.created_at.
        """
        d0 = start_dt.strftime("%Y-%m-%d")
        end_excl_dt = end_dt + timedelta(days=1)
        d1 = end_excl_dt.strftime("%Y-%m-%d")
        return d0, d1, end_excl_dt

    def analytics_counts(self, owner: str) -> dict[str, int]:
        agents = self.get_agents_by_owner(owner)
        live = sum(1 for a in agents if bool(a.get("published_version")))
        in_progress = sum(1 for a in agents if _has_unpublished_changes(a))
        return {"myAgents": len(agents), "live": live, "inProgress": in_progress}

    def analytics_shared_users(self, owner: str) -> int:
        c = self._conn.cursor()
        row = c.execute(
            """
            SELECT COUNT(DISTINCT s.principal) AS cnt
            FROM agent_shares s
                     JOIN agents a ON a.id = s.agent_id
            WHERE a.owner = ? AND s.principal_type = 'user'
            """,
            (owner,),
        ).fetchone()
        return int(row["cnt"] if row else 0)

    def analytics_shared_users_for_all_user_agents(self) -> int:
        """
        Count DISTINCT principals across shares on ANY user agent (i.e., those that exist in agents table).
        """
        c = self._conn.cursor()
        row = c.execute(
            """
            SELECT COUNT(DISTINCT s.principal) AS cnt
            FROM agent_shares s
                     JOIN agents a ON a.id = s.agent_id
            WHERE s.principal_type = 'user'
            """
        ).fetchone()
        return int(row["cnt"] if row else 0)

    def analytics_shared_agents(self, owner: str) -> int:
        """
        Number of user agents (owned by `owner`) that are shared (to a user or a group).
        """
        c = self._conn.cursor()
        row = c.execute(
            """
            SELECT COUNT(DISTINCT s.agent_id) AS cnt
            FROM agent_shares s
                     JOIN agents a ON a.id = s.agent_id
            WHERE a.owner = ?
            """,
            (owner,),
        ).fetchone()
        return int(row["cnt"] if row else 0)

    def analytics_shared_agents_for_all_user_agents(self) -> int:
        """Number of user agents (any owner) that are shared (to a user or a group)."""
        c = self._conn.cursor()
        row = c.execute(
            """
            SELECT COUNT(DISTINCT s.agent_id) AS cnt
            FROM agent_shares s
                     JOIN agents a ON a.id = s.agent_id
            """
        ).fetchone()
        return int(row["cnt"] if row else 0)

    def analytics_counts_all_user_agents(self) -> dict[str, int]:
        """
        Same semantics as analytics_counts(owner) but across *all* user agents.
        Counts pure drafts correctly: inProgress = _has_unpublished_changes(agent) regardless of publish.
        """
        agents = self.get_all_agents()
        live = sum(1 for a in agents if bool(a.get("published_version")))
        in_progress = sum(1 for a in agents if _has_unpublished_changes(a))
        return {"myAgents": len(agents), "live": live, "inProgress": in_progress}

    def list_user_agent_owners(self) -> List[Dict[str, Any]]:
        c = self._conn.cursor()
        owners = [
            {"id": r["owner"], "agentCount": int(r["cnt"])}
            for r in c.execute("SELECT owner, COUNT(*) AS cnt FROM agents GROUP BY owner ORDER BY owner")
        ]

        out: List[Dict[str, Any]] = []
        for o in owners:
            login = o["id"]
            display_name = ""
            email = ""
            try:
                ui = get_user_info(login)
                display_name = ui.display_name or ""
                email = ui.email or ""
            except Exception:
                logger.exception(f"Failed to get user info for {login}")
                pass
            out.append(
                {
                    "id": login,
                    "displayName": display_name,
                    "email": email,
                    "agentCount": o["agentCount"],
                }
            )
        return out

    # TODO: replace selected with used when find a way to compute them reliably
    def analytics_usage_buckets(
            self,
            owner: str,
            start: Optional[str],
            end: Optional[str],
            bucket: str,
            agent_id: Optional[str],
            *,
            is_project_admin: bool = False,
            agent_type: Optional[str] = None,
            owner_id: Optional[str] = None,
            enterprise_ids: Optional[List[str]] = None,
    ) -> list[dict]:
        start_dt, end_dt = self._normalize_range(start, end)
        # Build agent scope
        clause, scope_params = self._agent_scope_clause(
            owner=owner,
            agent_id=agent_id,
            is_project_admin=is_project_admin,
            agent_type=agent_type,
            owner_id=owner_id,
            enterprise_ids=enterprise_ids,
        )

        # group key builder
        def key_for(ts: str) -> str:
            if bucket == "day":
                return ts[:10]
            if bucket == "month":
                return ts[:7]
            # week: ISO year-week
            try:
                dt = datetime.fromisoformat(ts.replace("Z", "+00:00"))
            except Exception:
                dt = datetime.strptime(ts[:19], "%Y-%m-%dT%H:%M:%S")
            iso_y, iso_w, _ = dt.isocalendar()
            return f"{iso_y}-W{iso_w:02d}"

        counts: dict[str, int] = {}
        c = self._conn.cursor()
        d0, d1, end_excl_dt = self._day_bounds(start_dt, end_dt)
        params = [d0, d1, *scope_params]
        sql = f"""
            SELECT DISTINCT m.id, m.created_at
            FROM messages m
            JOIN conversations c ON m.conversation_id = c.conversation_id
            JOIN message_agents ma ON ma.message_id = m.id
            WHERE m.role = 'assistant'
              AND ma.selected = 1
              AND m.created_at >= ?
              AND m.created_at <  ?
             {clause}

        """
        for row in c.execute(sql, params):
            ts = row["created_at"]
            dt = self._parse_created_ts(ts)
            if not (start_dt <= dt < end_excl_dt):
                continue
            k = key_for(ts)
            counts[k] = counts.get(k, 0) + 1  # one per assistant message

        # produce sorted series
        series = [{"periodStart": k, "count": counts[k]} for k in sorted(counts.keys())]
        return series

    # TODO: replace selected with used when find a way to compute them reliably
    def analytics_feedback_counts(
            self,
            owner: str,
            start: Optional[str],
            end: Optional[str],
            agent_id: Optional[str],
            *,
            is_project_admin: bool = False,
            agent_type: Optional[str] = None,
            owner_id: Optional[str] = None,
            enterprise_ids: Optional[List[str]] = None,
    ) -> dict[str, int]:
        start_dt, end_dt = self._normalize_range(start, end)
        clause, scope_params = self._agent_scope_clause(
            owner=owner,
            agent_id=agent_id,
            is_project_admin=is_project_admin,
            agent_type=agent_type,
            owner_id=owner_id,
            enterprise_ids=enterprise_ids,
        )

        out = {"positive": 0, "negative": 0, "none": 0}
        c = self._conn.cursor()
        d0, d1, end_excl_dt = self._day_bounds(start_dt, end_dt)
        params = [d0, d1, *scope_params]
        sql = f"""
            SELECT DISTINCT m.id, m.created_at, m.feedback_rating
            FROM messages m
            JOIN conversations c ON m.conversation_id = c.conversation_id
            JOIN message_agents ma ON ma.message_id = m.id
            WHERE m.role = 'assistant'
              AND ma.selected = 1
              AND m.created_at >= ?
              AND m.created_at <  ?
              {clause}
        """
        for row in c.execute(sql, params):
            ts = row["created_at"]
            dt = self._parse_created_ts(ts)
            if not (start_dt <= dt < end_excl_dt):
                continue
            rating = row["feedback_rating"]
            if rating is None:
                out["none"] += 1
            else:
                out["positive" if int(rating) == 1 else "negative"] += 1
        return out

    # TODO: replace selected with used when find a way to compute them reliably

    def analytics_activity(
            self,
            owner: str,
            start: Optional[str],
            end: Optional[str],
            agent_id: Optional[str],
            limit: int,
            offset: int,
            *,
            is_project_admin: bool = False,
            agent_type: Optional[str] = None,
            owner_id: Optional[str] = None,
            enterprise_ids: Optional[List[str]] = None,
            q: Optional[str] = None,
            sort_by: Optional[str] = None,
            sort_dir: Optional[str] = None,
            group_by: Optional[str] = None,
    ) -> tuple[list[dict], int]:
        start_dt, end_dt = self._normalize_range(start, end)
        clause, scope_params = self._agent_scope_clause(
            owner=owner,
            agent_id=agent_id,
            is_project_admin=is_project_admin,
            agent_type=agent_type,
            owner_id=owner_id,
            enterprise_ids=enterprise_ids,
        )

        name_map = {a["id"]: a.get("name", a["id"]) for a in self.get_agents_by_owner(owner)}
        agg: dict[tuple[str, str], dict] = {}
        c = self._conn.cursor()
        d0, d1, end_excl_dt = self._day_bounds(start_dt, end_dt)
        params = [d0, d1, *scope_params]
        sql = f"""
            SELECT DISTINCT m.id, m.created_at, c.user_id, ma.agent_id, m.feedback_rating
            FROM messages m
            JOIN conversations c ON m.conversation_id = c.conversation_id
            JOIN message_agents ma ON ma.message_id = m.id
            WHERE m.role = 'assistant'
              AND ma.selected = 1
              AND m.created_at >= ?
              AND m.created_at <  ?
               {clause}
        """
        for row in c.execute(sql, params):
            ts = row["created_at"]
            dt = self._parse_created_ts(ts)
            if not (start_dt <= dt < end_excl_dt):
                continue
            aid = row["agent_id"]
            uid = row["user_id"]
            key = (aid, uid)
            entry = agg.setdefault(
                key,
                {
                    "agentId": aid,
                    "agentName": name_map.get(aid, aid),
                    "user": uid,
                    "questions": 0,
                    "rated": 0,
                    "sum": 0.0,
                },
            )
            entry["questions"] += 1
            rating = row["feedback_rating"]
            if rating is not None:
                entry["rated"] += 1
                entry["sum"] += 1.0 if int(rating) == 1 else 0.0

        # Build base rows (keep 'rated' and 'sum' for later aggregation)
        base_rows: list[dict] = list(agg.values())
        # --- apply q BEFORE grouping on detailed (agentId,user) rows ---
        q_norm = (q or "").strip().lower()
        if q_norm:

            def _contains_detail(r: dict) -> bool:
                return q_norm in (r.get("agentName") or "").lower() or q_norm in (r.get("user") or "").lower()

            base_rows = [r for r in base_rows if _contains_detail(r)]

        # --- optional grouping ---
        group_by = (group_by or "").lower()
        rows_all: list[dict] = []
        if group_by == "agent":
            by_agent: dict[str, dict] = {}
            for r in base_rows:
                aid = r["agentId"]
                e = by_agent.setdefault(
                    aid,
                    {
                        "agentId": aid,
                        "agentName": r["agentName"],
                        "user": None,  # keep shape consistent
                        "questions": 0,
                        "rated": 0,
                        "sum": 0.0,
                    },
                )
                e["questions"] += r.get("questions", 0)
                e["rated"] += r.get("rated", 0)
                e["sum"] += r.get("sum", 0.0)
            for e in by_agent.values():
                rated = e.pop("rated")
                s = e.pop("sum", 0.0)
                e["avgFeedback"] = (s / rated) if rated else None
                rows_all.append(e)
        elif group_by == "user":
            by_user: dict[str, dict] = {}
            for r in base_rows:
                uid = r["user"]
                e = by_user.setdefault(
                    uid,
                    {
                        "agentId": None,
                        "agentName": None,
                        "user": uid,
                        "questions": 0,
                        "rated": 0,
                        "sum": 0.0,
                    },
                )
                e["questions"] += r.get("questions", 0)
                e["rated"] += r.get("rated", 0)
                e["sum"] += r.get("sum", 0.0)
            for e in by_user.values():
                rated = e.pop("rated")
                s = e.pop("sum", 0.0)
                e["avgFeedback"] = (s / rated) if rated else None
                rows_all.append(e)
        else:
            # original (agent,user) pairs
            for r in base_rows:
                rated = r.pop("rated")
                s = r.pop("sum", 0.0)
                r["avgFeedback"] = (s / rated) if rated else None
                rows_all.append(r)

        # --- sorting ---
        sort_by = sort_by or "questions"
        sort_dir = (sort_dir or "desc").lower()
        if sort_by not in {"agentName", "user", "questions", "avgFeedback"}:
            sort_by = "questions"
        reverse = sort_dir == "desc"

        if sort_by == "avgFeedback":
            # Keep rows with no feedback at the end, regardless of direction
            if sort_dir == "asc":
                rows_all.sort(key=lambda r: (r.get("avgFeedback") is None, r.get("avgFeedback") or 0.0))
            else:
                rows_all.sort(key=lambda r: (r.get("avgFeedback") is None, -(r.get("avgFeedback") or 0.0)))
        elif sort_by == "questions":
            rows_all.sort(key=lambda r: int(r.get("questions", 0)), reverse=reverse)
        else:
            # agentName / user (case-insensitive)
            rows_all.sort(key=lambda r: (r.get(sort_by) or "").casefold(), reverse=reverse)

        # paginate
        total = len(rows_all)
        return rows_all[offset : offset + limit], total

    def analytics_active_users_buckets(
            self,
            owner: str,
            start: Optional[str],
            end: Optional[str],
            bucket: str,
            agent_id: Optional[str],
            *,
            is_project_admin: bool = False,
            agent_type: Optional[str] = None,
            owner_id: Optional[str] = None,
            enterprise_ids: Optional[List[str]] = None,
    ) -> list[dict]:
        """
        COUNT(DISTINCT user_id) per bucket for assistant messages that invoked an owned agent.
        """
        start_dt, end_dt = self._normalize_range(start, end)
        clause, scope_params = self._agent_scope_clause(
            owner=owner,
            agent_id=agent_id,
            is_project_admin=is_project_admin,
            agent_type=agent_type,
            owner_id=owner_id,
            enterprise_ids=enterprise_ids,
        )

        def key_for(ts: str) -> str:
            if bucket == "day":
                return ts[:10]
            if bucket == "month":
                return ts[:7]
            try:
                dt = datetime.fromisoformat(ts.replace("Z", "+00:00"))
            except Exception:
                dt = datetime.strptime(ts[:19], "%Y-%m-%dT%H:%M:%S")
            iso_y, iso_w, _ = dt.isocalendar()
            return f"{iso_y}-W{iso_w:02d}"

        c = self._conn.cursor()
        d0, d1, end_excl_dt = self._day_bounds(start_dt, end_dt)
        params = [d0, d1, *scope_params]
        sql = f"""
            SELECT DISTINCT m.id, m.created_at, c.user_id
            FROM messages m
            JOIN conversations c ON m.conversation_id = c.conversation_id
            JOIN message_agents ma ON ma.message_id = m.id
            WHERE m.role = 'assistant'
              AND ma.selected = 1
              AND m.created_at >= ?
              AND m.created_at <  ?
              {clause}
        """
        buckets: dict[str, set[str]] = {}
        for row in c.execute(sql, params):
            ts = row["created_at"]
            dt = self._parse_created_ts(ts)
            if not (start_dt <= dt < end_excl_dt):
                continue
            k = key_for(ts)
            s = buckets.setdefault(k, set())
            s.add(row["user_id"])

        return [{"periodStart": k, "activeUsers": len(uids)} for k, uids in sorted(buckets.items())]

    @request_cached
    def get_message_trace(self, message_id: str, user_id: Optional[str] = None) -> str:
        c = self._conn.cursor()
        if user_id:
            msg_row = c.execute(
                """
                SELECT * FROM messages m
                                  JOIN conversations c ON m.conversation_id = c.conversation_id
                WHERE m.id = ? AND c.user_id = ?
                """,
                (message_id, user_id),
            ).fetchone()
            if not msg_row or not msg_row["trace"]:
                return ""
            return msg_row["trace"]
        return ""