nas-tools/dashboard/backend/db/opc_db.py

"""
OPC Database Module - PostgreSQL connection and schema management
"""

import json
import os
from datetime import datetime
from typing import Any

import asyncpg

# Database connection settings
DB_HOST = os.getenv("OPC_DB_HOST", "gitea-db")
DB_PORT = int(os.getenv("OPC_DB_PORT", "5432"))
DB_NAME = os.getenv("OPC_DB_NAME", "opc")
DB_USER = os.getenv("OPC_DB_USER", "gitea")
DB_PASSWORD = os.getenv("GITEA_DB_PASSWORD", "")

# Connection pool
_pool: asyncpg.Pool | None = None


async def get_pool() -> asyncpg.Pool:
    """Get or create database connection pool with retry logic"""
    global _pool
    if _pool is None:
        import asyncio
        import logging

        logger = logging.getLogger(__name__)

        max_retries = 3
        retry_delay = 5

        for attempt in range(max_retries):
            try:
                logger.info(f"Attempting to connect to OPC database (attempt {attempt + 1}/{max_retries})")
                _pool = await asyncpg.create_pool(
                    host=DB_HOST,
                    port=DB_PORT,
                    database=DB_NAME,
                    user=DB_USER,
                    password=DB_PASSWORD,
                    min_size=2,
                    max_size=10,
                )
                logger.info("Successfully connected to OPC database")
                break
            except Exception as e:
                logger.error(f"Failed to connect to OPC database (attempt {attempt + 1}/{max_retries}): {e}")
                if attempt < max_retries - 1:
                    logger.info(f"Retrying in {retry_delay} seconds...")
                    await asyncio.sleep(retry_delay)
                else:
                    logger.error("All connection attempts failed")
                    raise Exception(f"Failed to connect to OPC database after {max_retries} attempts: {e}")
    return _pool


async def close_pool():
    """Close database connection pool"""
    global _pool
    if _pool:
        await _pool.close()
        _pool = None


async def init_db():
    """Initialize database schema"""
    pool = await get_pool()
    async with pool.acquire() as conn:
        # Create tables
        await conn.execute(
            """
            CREATE TABLE IF NOT EXISTS projects (
                id SERIAL PRIMARY KEY,
                name TEXT NOT NULL,
                description TEXT,
                status TEXT DEFAULT 'active',
                client_id INTEGER,
                metadata JSONB,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )
        """
        )

        await conn.execute(
            """
            CREATE TABLE IF NOT EXISTS tasks (
                id SERIAL PRIMARY KEY,
                title TEXT NOT NULL,
                description TEXT,
                status TEXT NOT NULL DEFAULT 'parking_lot',
                priority TEXT DEFAULT 'medium',
                project_id INTEGER REFERENCES projects(id),
                assigned_to TEXT,
                assigned_type TEXT DEFAULT 'human',
                tags JSONB,
                metadata JSONB DEFAULT '{}',
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                completed_at TIMESTAMP,
                due_date TIMESTAMP
            )
        """
        )

        await conn.execute(
            """
            CREATE TABLE IF NOT EXISTS task_history (
                id SERIAL PRIMARY KEY,
                task_id INTEGER NOT NULL REFERENCES tasks(id) ON DELETE CASCADE,
                action TEXT NOT NULL,
                actor TEXT NOT NULL,
                actor_type TEXT DEFAULT 'human',
                old_value JSONB,
                new_value JSONB,
                timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )
        """
        )

        await conn.execute(
            """
            CREATE TABLE IF NOT EXISTS agents (
                id TEXT PRIMARY KEY,
                name TEXT NOT NULL,
                role TEXT NOT NULL,
                description TEXT,
                capabilities JSONB,
                system_prompt TEXT,
                config JSONB,
                enabled BOOLEAN DEFAULT TRUE,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )
        """
        )

        await conn.execute(
            """
            CREATE TABLE IF NOT EXISTS agent_executions (
                id SERIAL PRIMARY KEY,
                task_id INTEGER NOT NULL REFERENCES tasks(id) ON DELETE CASCADE,
                agent_id TEXT NOT NULL REFERENCES agents(id),
                status TEXT DEFAULT 'pending',
                input_context JSONB,
                output_result JSONB,
                actions_proposed JSONB,
                actions_executed JSONB,
                error_message TEXT,
                started_at TIMESTAMP,
                completed_at TIMESTAMP,
                requires_approval BOOLEAN DEFAULT FALSE,
                approved_by TEXT,
                approved_at TIMESTAMP
            )
        """
        )

        await conn.execute(
            """
            CREATE TABLE IF NOT EXISTS time_entries (
                id SERIAL PRIMARY KEY,
                task_id INTEGER REFERENCES tasks(id) ON DELETE CASCADE,
                project_id INTEGER REFERENCES projects(id),
                description TEXT,
                duration_minutes INTEGER NOT NULL,
                billable BOOLEAN DEFAULT TRUE,
                hourly_rate NUMERIC(10, 2),
                started_at TIMESTAMP,
                ended_at TIMESTAMP,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )
        """
        )

        await conn.execute(
            """
            CREATE TABLE IF NOT EXISTS clients (
                id SERIAL PRIMARY KEY,
                name TEXT NOT NULL,
                email TEXT,
                company TEXT,
                notes TEXT,
                metadata JSONB,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )
        """
        )

        # Create indexes
        await conn.execute("CREATE INDEX IF NOT EXISTS idx_tasks_status ON tasks(status)")
        await conn.execute("CREATE INDEX IF NOT EXISTS idx_tasks_assigned ON tasks(assigned_to, assigned_type)")
        await conn.execute("CREATE INDEX IF NOT EXISTS idx_tasks_project ON tasks(project_id)")
        await conn.execute("CREATE INDEX IF NOT EXISTS idx_task_history_task ON task_history(task_id)")
        await conn.execute("CREATE INDEX IF NOT EXISTS idx_agent_executions_task ON agent_executions(task_id)")
        await conn.execute("CREATE INDEX IF NOT EXISTS idx_agent_executions_agent ON agent_executions(agent_id)")
        await conn.execute("CREATE INDEX IF NOT EXISTS idx_time_entries_task ON time_entries(task_id)")
        await conn.execute("CREATE INDEX IF NOT EXISTS idx_time_entries_project ON time_entries(project_id)")

        # Insert default agents
        await seed_agents(conn)


async def seed_agents(conn):
    """Seed default agents"""
    agents = [
        {
            "id": "pm",
            "name": "Project Manager",
            "role": "Project Manager",
            "description": "Breaks down tasks, plans timelines, monitors progress",
            "capabilities": ["task_breakdown", "timeline_planning", "risk_assessment", "status_reporting"],
            "system_prompt": """You are a Project Manager agent. Your role is to:
- Break down large tasks into manageable subtasks
- Plan realistic timelines and identify dependencies
- Monitor project progress and identify blockers
- Provide status reports and risk assessments

When assigned a task, analyze it and propose concrete actions.""",
            "config": {"model": "cc-kiro", "temperature": 0.7, "approval_level": "auto"},
        },
        {
            "id": "cto",
            "name": "CTO",
            "role": "Chief Technology Officer",
            "description": "Reviews code, designs architecture, manages technical debt",
            "capabilities": ["code_review", "architecture_design", "tech_debt_management", "security_audit"],
            "system_prompt": """You are a CTO agent. Your role is to:
- Review code quality and architecture decisions
- Design scalable and maintainable systems
- Identify and prioritize technical debt
- Ensure security best practices

When assigned a task, provide technical leadership and propose improvements.""",
            "config": {"model": "cc-kiro", "temperature": 0.5, "approval_level": "cxo"},
        },
        {
            "id": "coo",
            "name": "COO",
            "role": "Chief Operating Officer",
            "description": "Optimizes processes, automates workflows, improves efficiency",
            "capabilities": ["process_improvement", "workflow_automation", "efficiency_analysis"],
            "system_prompt": """You are a COO agent. Your role is to:
- Identify opportunities for process improvement
- Automate repetitive workflows using n8n
- Analyze operational efficiency
- Optimize resource allocation

When assigned a task, look for automation and optimization opportunities.""",
            "config": {"model": "cc-kiro", "temperature": 0.6, "approval_level": "cxo"},
        },
    ]

    for agent in agents:
        await conn.execute(
            """
            INSERT INTO agents (id, name, role, description, capabilities, system_prompt, config, enabled)
            VALUES ($1, $2, $3, $4, $5, $6, $7, $8)
            ON CONFLICT (id) DO UPDATE SET
                name = EXCLUDED.name,
                role = EXCLUDED.role,
                description = EXCLUDED.description,
                capabilities = EXCLUDED.capabilities,
                system_prompt = EXCLUDED.system_prompt,
                config = EXCLUDED.config
        """,
            agent["id"],
            agent["name"],
            agent["role"],
            agent["description"],
            json.dumps(agent["capabilities"]),
            agent["system_prompt"],
            json.dumps(agent["config"]),
            True,
        )


# Task operations
async def create_task(
    title: str,
    description: str | None = None,
    status: str = "parking_lot",
    priority: str = "medium",
    project_id: int | None = None,
    assigned_to: str | None = None,
    assigned_type: str = "human",
    tags: list[str] | None = None,
    due_date: datetime | None = None,
    actor: str = "system",
) -> dict[str, Any]:
    """Create a new task"""
    pool = await get_pool()
    async with pool.acquire() as conn:
        row = await conn.fetchrow(
            """
            INSERT INTO tasks (title, description, status, priority, project_id, assigned_to, assigned_type, tags, due_date)
            VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9)
            RETURNING id, title, description, status, priority, project_id, assigned_to, assigned_type, tags,
                      created_at, updated_at, completed_at, due_date
        """,
            title,
            description,
            status,
            priority,
            project_id,
            assigned_to,
            assigned_type,
            json.dumps(tags or []),
            due_date,
        )

        task = dict(row)
        task["tags"] = json.loads(task["tags"]) if task["tags"] else []

        # Convert datetime objects to ISO format strings for JSON serialization
        task_for_history = task.copy()
        for key in ["created_at", "updated_at", "completed_at", "due_date"]:
            if task_for_history.get(key):
                task_for_history[key] = task_for_history[key].isoformat()

        # Log history
        await conn.execute(
            """
            INSERT INTO task_history (task_id, action, actor, actor_type, new_value)
            VALUES ($1, $2, $3, $4, $5)
        """,
            task["id"],
            "created",
            actor,
            "human",
            json.dumps(task_for_history),
        )

        return task


async def get_tasks(
    status: str | None = None,
    assigned_to: str | None = None,
    project_id: int | None = None,
    priority: str | None = None,
    tags: list[str] | None = None,
    limit: int = 50,
    offset: int = 0,
) -> list[dict[str, Any]]:
    """Get tasks with filters"""
    pool = await get_pool()
    async with pool.acquire() as conn:
        query = "SELECT * FROM tasks WHERE 1=1"
        params = []
        param_count = 0

        if status:
            param_count += 1
            query += f" AND status = ${param_count}"
            params.append(status)

        if assigned_to:
            param_count += 1
            query += f" AND assigned_to = ${param_count}"
            params.append(assigned_to)

        if project_id:
            param_count += 1
            query += f" AND project_id = ${param_count}"
            params.append(project_id)

        if priority:
            param_count += 1
            query += f" AND priority = ${param_count}"
            params.append(priority)

        query += " ORDER BY created_at DESC"
        param_count += 1
        query += f" LIMIT ${param_count}"
        params.append(limit)
        param_count += 1
        query += f" OFFSET ${param_count}"
        params.append(offset)

        rows = await conn.fetch(query, *params)
        tasks = [dict(row) for row in rows]

        for task in tasks:
            task["tags"] = json.loads(task["tags"]) if task["tags"] else []

        return tasks


async def update_task(task_id: int, updates: dict[str, Any], actor: str = "system") -> dict[str, Any]:
    """Update a task"""
    pool = await get_pool()
    async with pool.acquire() as conn:
        # Get old values
        old_task = await conn.fetchrow("SELECT * FROM tasks WHERE id = $1", task_id)
        if not old_task:
            raise ValueError(f"Task {task_id} not found")

        old_values = dict(old_task)

        # Build update query
        set_clauses = ["updated_at = CURRENT_TIMESTAMP"]
        params = []
        param_count = 0

        for key, value in updates.items():
            if key in [
                "title",
                "description",
                "status",
                "priority",
                "project_id",
                "assigned_to",
                "assigned_type",
                "due_date",
                "completed_at",
            ]:
                param_count += 1
                set_clauses.append(f"{key} = ${param_count}")
                params.append(value)
            elif key == "tags":
                param_count += 1
                set_clauses.append(f"tags = ${param_count}")
                params.append(json.dumps(value))

        param_count += 1
        params.append(task_id)

        query = f"UPDATE tasks SET {', '.join(set_clauses)} WHERE id = ${param_count} RETURNING *"
        row = await conn.fetchrow(query, *params)

        task = dict(row)
        task["tags"] = json.loads(task["tags"]) if task["tags"] else []

        # Convert datetime objects to ISO format strings for JSON serialization
        old_values_for_history = old_values.copy()
        updates_for_history = updates.copy()
        for key in ["created_at", "updated_at", "completed_at", "due_date"]:
            if old_values_for_history.get(key):
                old_values_for_history[key] = old_values_for_history[key].isoformat()
            if updates_for_history.get(key):
                updates_for_history[key] = updates_for_history[key].isoformat()

        # Log history
        await conn.execute(
            """
            INSERT INTO task_history (task_id, action, actor, actor_type, old_value, new_value)
            VALUES ($1, $2, $3, $4, $5, $6)
        """,
            task_id,
            "updated",
            actor,
            "human",
            json.dumps(old_values_for_history),
            json.dumps(updates_for_history),
        )

        return task


async def delete_task(task_id: int, actor: str = "system"):
    """Delete a task"""
    pool = await get_pool()
    async with pool.acquire() as conn:
        await conn.execute(
            """
            INSERT INTO task_history (task_id, action, actor, actor_type)
            VALUES ($1, $2, $3, $4)
        """,
            task_id,
            "deleted",
            actor,
            "human",
        )

        await conn.execute("DELETE FROM tasks WHERE id = $1", task_id)


# Time tracking operations
async def start_time_entry(task_id: int, project_id: int | None = None, actor: str = "system") -> dict[str, Any]:
    """Start automatic time tracking for a task"""
    pool = await get_pool()
    async with pool.acquire() as conn:
        # Check if there's already an active entry
        active = await conn.fetchrow(
            """
            SELECT * FROM time_entries
            WHERE task_id = $1 AND ended_at IS NULL
            ORDER BY started_at DESC LIMIT 1
        """,
            task_id,
        )

        if active:
            return dict(active)

        # Create new time entry
        row = await conn.fetchrow(
            """
            INSERT INTO time_entries (task_id, project_id, started_at, duration_minutes)
            VALUES ($1, $2, CURRENT_TIMESTAMP, 0)
            RETURNING *
        """,
            task_id,
            project_id,
        )

        return dict(row)


async def stop_time_entry(task_id: int) -> dict[str, Any] | None:
    """Stop automatic time tracking for a task"""
    pool = await get_pool()
    async with pool.acquire() as conn:
        # Find active entry
        active = await conn.fetchrow(
            """
            SELECT * FROM time_entries
            WHERE task_id = $1 AND ended_at IS NULL
            ORDER BY started_at DESC LIMIT 1
        """,
            task_id,
        )

        if not active:
            return None

        # Calculate duration and update
        row = await conn.fetchrow(
            """
            UPDATE time_entries
            SET ended_at = CURRENT_TIMESTAMP,
                duration_minutes = EXTRACT(EPOCH FROM (CURRENT_TIMESTAMP - started_at)) / 60
            WHERE id = $1
            RETURNING *
        """,
            active["id"],
        )

        return dict(row)


async def get_time_entries(task_id: int | None = None, project_id: int | None = None) -> list[dict[str, Any]]:
    """Get time entries"""
    pool = await get_pool()
    async with pool.acquire() as conn:
        query = "SELECT * FROM time_entries WHERE 1=1"
        params = []

        if task_id:
            params.append(task_id)
            query += f" AND task_id = ${len(params)}"

        if project_id:
            params.append(project_id)
            query += f" AND project_id = ${len(params)}"

        query += " ORDER BY started_at DESC"

        rows = await conn.fetch(query, *params)
        return [dict(row) for row in rows]


async def get_task_history(task_id: int) -> list[dict[str, Any]]:
    """Get task history"""
    pool = await get_pool()
    async with pool.acquire() as conn:
        rows = await conn.fetch(
            """
            SELECT * FROM task_history
            WHERE task_id = $1
            ORDER BY timestamp DESC
        """,
            task_id,
        )
        return [dict(row) for row in rows]


# Agent operations
async def get_agents(enabled_only: bool = True) -> list[dict[str, Any]]:
    """Get all agents"""
    pool = await get_pool()
    async with pool.acquire() as conn:
        query = "SELECT * FROM agents"
        if enabled_only:
            query += " WHERE enabled = TRUE"
        query += " ORDER BY id"

        rows = await conn.fetch(query)
        agents = []
        for row in rows:
            agent = dict(row)
            agent["capabilities"] = json.loads(agent["capabilities"]) if agent["capabilities"] else []
            agent["config"] = json.loads(agent["config"]) if agent["config"] else {}
            agents.append(agent)
        return agents


async def get_agent(agent_id: str) -> dict[str, Any] | None:
    """Get agent by ID"""
    pool = await get_pool()
    async with pool.acquire() as conn:
        row = await conn.fetchrow("SELECT * FROM agents WHERE id = $1", agent_id)
        if not row:
            return None

        agent = dict(row)
        agent["capabilities"] = json.loads(agent["capabilities"]) if agent["capabilities"] else []
        agent["config"] = json.loads(agent["config"]) if agent["config"] else {}
        return agent


async def create_agent_execution(
    task_id: int, agent_id: str, actor: str = "system", requires_approval: bool = False
) -> dict[str, Any]:
    """Create agent execution record"""
    pool = await get_pool()
    async with pool.acquire() as conn:
        # Get task and agent context
        task = await conn.fetchrow("SELECT * FROM tasks WHERE id = $1", task_id)
        agent = await conn.fetchrow("SELECT * FROM agents WHERE id = $1", agent_id)

        if not task or not agent:
            raise ValueError("Task or agent not found")

        # Build input context
        # Prepare input context with serialized datetimes
        task_for_context = dict(task)
        for key in ["created_at", "updated_at", "completed_at", "due_date"]:
            if task_for_context.get(key):
                task_for_context[key] = (
                    task_for_context[key].isoformat()
                    if hasattr(task_for_context[key], "isoformat")
                    else task_for_context[key]
                )

        agent_for_context = dict(agent)
        for key in ["created_at"]:
            if agent_for_context.get(key):
                agent_for_context[key] = (
                    agent_for_context[key].isoformat()
                    if hasattr(agent_for_context[key], "isoformat")
                    else agent_for_context[key]
                )

        input_context = {
            "task": task_for_context,
            "agent": agent_for_context,
            "timestamp": datetime.utcnow().isoformat(),
        }

        # Check if agent requires approval (CxO level)
        agent_config = json.loads(agent["config"]) if agent["config"] else {}
        requires_approval = agent_config.get("approval_level") == "cxo"

        row = await conn.fetchrow(
            """
            INSERT INTO agent_executions (task_id, agent_id, status, input_context, requires_approval, started_at)
            VALUES ($1, $2, $3, $4, $5, CURRENT_TIMESTAMP)
            RETURNING *
        """,
            task_id,
            agent_id,
            "pending",
            json.dumps(input_context),
            requires_approval,
        )

        execution = dict(row)
        execution["input_context"] = json.loads(execution["input_context"]) if execution["input_context"] else {}
        return execution


async def get_agent_executions(
    task_id: int | None = None, agent_id: str | None = None, status: str | None = None, limit: int = 50
) -> list[dict[str, Any]]:
    """Get agent executions"""
    pool = await get_pool()
    async with pool.acquire() as conn:
        query = "SELECT * FROM agent_executions WHERE 1=1"
        params = []

        if task_id:
            params.append(task_id)
            query += f" AND task_id = ${len(params)}"

        if agent_id:
            params.append(agent_id)
            query += f" AND agent_id = ${len(params)}"

        if status:
            params.append(status)
            query += f" AND status = ${len(params)}"

        query += " ORDER BY started_at DESC"
        params.append(limit)
        query += f" LIMIT ${len(params)}"

        rows = await conn.fetch(query, *params)
        executions = []
        for row in rows:
            execution = dict(row)
            for json_field in ["input_context", "output_result", "actions_proposed", "actions_executed"]:
                if execution[json_field]:
                    execution[json_field] = json.loads(execution[json_field])
            executions.append(execution)
        return executions


async def get_task_by_id(task_id: int) -> dict[str, Any] | None:
    """Get a single task by ID"""
    pool = await get_pool()
    async with pool.acquire() as conn:
        row = await conn.fetchrow("SELECT * FROM tasks WHERE id = $1", task_id)
        if not row:
            return None

        task = dict(row)
        task["tags"] = json.loads(task["tags"]) if task["tags"] else []
        return task


async def get_agent_execution(execution_id: int) -> dict[str, Any] | None:
    """Get a single agent execution by ID"""
    pool = await get_pool()
    async with pool.acquire() as conn:
        row = await conn.fetchrow("SELECT * FROM agent_executions WHERE id = $1", execution_id)
        if not row:
            return None

        execution = dict(row)
        for json_field in ["input_context", "output_result", "actions_proposed", "actions_executed"]:
            if execution[json_field]:
                execution[json_field] = json.loads(execution[json_field])
        return execution


async def approve_agent_execution(
    execution_id: int, approved: bool, approved_by: str, feedback: str | None = None
) -> dict[str, Any]:
    """Approve or reject an agent execution"""
    pool = await get_pool()
    async with pool.acquire() as conn:
        # Update execution with approval status
        row = await conn.fetchrow(
            """
            UPDATE agent_executions
            SET approved_by = $1,
                approved_at = CURRENT_TIMESTAMP,
                status = CASE WHEN $2 THEN 'approved' ELSE 'rejected' END
            WHERE id = $3
            RETURNING *
        """,
            approved_by,
            approved,
            execution_id,
        )

        if not row:
            raise ValueError(f"Execution {execution_id} not found")

        execution = dict(row)
        for json_field in ["input_context", "output_result", "actions_proposed", "actions_executed"]:
            if execution[json_field]:
                execution[json_field] = json.loads(execution[json_field])

        # If feedback provided, add to output_result
        if feedback:
            if not execution["output_result"]:
                execution["output_result"] = {}
            execution["output_result"]["approval_feedback"] = feedback

            await conn.execute(
                """
                UPDATE agent_executions
                SET output_result = $1
                WHERE id = $2
            """,
                json.dumps(execution["output_result"]),
                execution_id,
            )

        return execution


async def update_agent_config(
    agent_id: str,
    name: str | None = None,
    description: str | None = None,
    system_prompt: str | None = None,
    config: dict[str, Any] | None = None,
    enabled: bool | None = None,
) -> dict[str, Any]:
    """Update agent configuration"""
    pool = await get_pool()
    async with pool.acquire() as conn:
        # Build update query dynamically
        updates = []
        params = []
        param_count = 0

        if name is not None:
            param_count += 1
            updates.append(f"name = ${param_count}")
            params.append(name)

        if description is not None:
            param_count += 1
            updates.append(f"description = ${param_count}")
            params.append(description)

        if system_prompt is not None:
            param_count += 1
            updates.append(f"system_prompt = ${param_count}")
            params.append(system_prompt)

        if config is not None:
            param_count += 1
            updates.append(f"config = ${param_count}")
            params.append(json.dumps(config))

        if enabled is not None:
            param_count += 1
            updates.append(f"enabled = ${param_count}")
            params.append(enabled)

        if not updates:
            # No updates provided, just return current agent
            return await get_agent(agent_id)

        param_count += 1
        params.append(agent_id)

        query = f"UPDATE agents SET {', '.join(updates)} WHERE id = ${param_count} RETURNING *"
        row = await conn.fetchrow(query, *params)

        if not row:
            raise ValueError(f"Agent {agent_id} not found")

        agent = dict(row)
        agent["capabilities"] = json.loads(agent["capabilities"]) if agent["capabilities"] else []
        agent["config"] = json.loads(agent["config"]) if agent["config"] else {}
        return agent


async def update_execution_status(
    execution_id: int,
    status: str,
    output_result: dict[str, Any] | None = None,
    actions_proposed: list[dict[str, Any]] | None = None,
    actions_executed: list[dict[str, Any]] | None = None,
    error_message: str | None = None,
) -> dict[str, Any]:
    """Update agent execution status and results"""
    pool = await get_pool()
    async with pool.acquire() as conn:
        # Check if started_at is already set
        current = await conn.fetchrow("SELECT started_at FROM agent_executions WHERE id = $1", execution_id)

        updates = ["status = $1"]
        params = [status]
        param_count = 1

        # Only set started_at if transitioning to running and not already set
        if status == "running" and (not current or not current["started_at"]):
            updates.append("started_at = CURRENT_TIMESTAMP")

        if status in ["completed", "failed", "rejected"]:
            updates.append("completed_at = CURRENT_TIMESTAMP")

        if output_result is not None:
            param_count += 1
            updates.append(f"output_result = ${param_count}")
            params.append(json.dumps(output_result))

        if actions_proposed is not None:
            param_count += 1
            updates.append(f"actions_proposed = ${param_count}")
            params.append(json.dumps(actions_proposed))

        if actions_executed is not None:
            param_count += 1
            updates.append(f"actions_executed = ${param_count}")
            params.append(json.dumps(actions_executed))

        if error_message is not None:
            param_count += 1
            updates.append(f"error_message = ${param_count}")
            params.append(error_message)

        param_count += 1
        params.append(execution_id)

        query = f"UPDATE agent_executions SET {', '.join(updates)} WHERE id = ${param_count} RETURNING *"
        row = await conn.fetchrow(query, *params)

        if not row:
            raise ValueError(f"Execution {execution_id} not found")

        execution = dict(row)
        for json_field in ["input_context", "output_result", "actions_proposed", "actions_executed"]:
            if execution[json_field]:
                execution[json_field] = json.loads(execution[json_field])
        return execution