Project_Velocity/backend/services/colony_repository.py

from __future__ import annotations

import json
from typing import Any

import asyncpg


_JSON_KEYS = {
    "context_refs",
    "requested_outputs",
    "payload",
    "input",
    "output",
    "citations",
    "before_state",
    "after_state",
    "detail",
}


def _row_dict(row: asyncpg.Record | None) -> dict[str, Any] | None:
    if row is None:
        return None
    data = dict(row)
    for key in _JSON_KEYS:
        if isinstance(data.get(key), str):
            data[key] = json.loads(data[key])
    return data


class ColonyRepository:
    def __init__(self, pool: asyncpg.Pool) -> None:
        self.pool = pool

    async def create_mission(self, mission: dict[str, Any]) -> dict[str, Any]:
        async with self.pool.acquire() as conn:
            row = await conn.fetchrow(
                """
                INSERT INTO colony_missions (
                    mission_id, tenant_id, mission_type, origin_surface, actor_id,
                    actor_role, risk_level, sensitivity_class, status, review_status,
                    time_budget_ms, token_budget, user_goal, normalized_goal,
                    context_refs, requested_outputs, payload, created_at, updated_at
                ) VALUES (
                    $1::uuid, $2, $3, $4, $5,
                    $6, $7, $8, 'pending', NULL,
                    $9, $10, $11, $12,
                    $13::jsonb, $14::jsonb, $15::jsonb, NOW(), NOW()
                )
                RETURNING *
                """,
                mission["mission_id"],
                mission["tenant_id"],
                mission["mission_type"],
                mission["origin_surface"],
                mission["actor_id"],
                mission.get("actor_role"),
                mission["risk_level"],
                mission["sensitivity_class"],
                mission["time_budget_ms"],
                mission["token_budget"],
                mission["user_goal"],
                mission["normalized_goal"],
                json.dumps(mission["context_refs"]),
                json.dumps(mission["requested_outputs"]),
                json.dumps(mission["payload"]),
            )
        return _row_dict(row)

    async def get_mission(self, mission_id: str, tenant_id: str) -> dict[str, Any] | None:
        async with self.pool.acquire() as conn:
            row = await conn.fetchrow(
                """
                SELECT *
                FROM colony_missions
                WHERE mission_id = $1::uuid
                  AND tenant_id = $2
                """,
                mission_id,
                tenant_id,
            )
        return _row_dict(row)

    async def update_status(
        self,
        mission_id: str,
        tenant_id: str,
        status: str,
        *,
        review_status: str | None = None,
    ) -> dict[str, Any] | None:
        async with self.pool.acquire() as conn:
            row = await conn.fetchrow(
                """
                UPDATE colony_missions
                SET status = $3,
                    review_status = COALESCE($4, review_status),
                    updated_at = NOW(),
                    completed_at = CASE
                        WHEN $3 IN ('completed', 'failed', 'dispatch_failed') THEN NOW()
                        ELSE completed_at
                    END
                WHERE mission_id = $1::uuid
                  AND tenant_id = $2
                RETURNING *
                """,
                mission_id,
                tenant_id,
                status,
                review_status,
            )
        return _row_dict(row)

    async def list_missions(self, tenant_id: str, *, limit: int, offset: int) -> list[dict[str, Any]]:
        async with self.pool.acquire() as conn:
            rows = await conn.fetch(
                """
                SELECT *
                FROM colony_missions
                WHERE tenant_id = $1
                ORDER BY created_at DESC
                LIMIT $2 OFFSET $3
                """,
                tenant_id,
                limit,
                offset,
            )
        return [_row_dict(row) for row in rows]

    async def artifacts(self, mission_id: str, tenant_id: str) -> dict[str, list[dict[str, Any]]]:
        async with self.pool.acquire() as conn:
            mission = await conn.fetchrow(
                "SELECT mission_id FROM colony_missions WHERE mission_id = $1::uuid AND tenant_id = $2",
                mission_id,
                tenant_id,
            )
            if mission is None:
                return {}
            tasks = await conn.fetch(
                """
                SELECT *
                FROM colony_tasks
                WHERE mission_id = $1::uuid
                  AND tenant_id = $2
                ORDER BY created_at ASC
                """,
                mission_id,
                tenant_id,
            )
            results = await conn.fetch(
                """
                SELECT *
                FROM colony_worker_results
                WHERE mission_id = $1::uuid
                  AND tenant_id = $2
                ORDER BY created_at ASC
                """,
                mission_id,
                tenant_id,
            )
            proposals = await conn.fetch(
                """
                SELECT *
                FROM colony_writeback_proposals
                WHERE mission_id = $1::uuid
                  AND tenant_id = $2
                ORDER BY created_at DESC
                """,
                mission_id,
                tenant_id,
            )
        return {
            "tasks": [_row_dict(row) for row in tasks],
            "worker_results": [_row_dict(row) for row in results],
            "writeback_proposals": [_row_dict(row) for row in proposals],
        }

    async def pending_writeback_proposals(self, mission_id: str, tenant_id: str) -> list[dict[str, Any]]:
        async with self.pool.acquire() as conn:
            rows = await conn.fetch(
                """
                SELECT *
                FROM colony_writeback_proposals
                WHERE mission_id = $1::uuid
                  AND tenant_id = $2
                  AND approval_status = 'pending'
                ORDER BY created_at ASC
                """,
                mission_id,
                tenant_id,
            )
        return [_row_dict(row) for row in rows]

    async def approve_pending_writebacks(self, mission_id: str, tenant_id: str, actor_id: str) -> int:
        async with self.pool.acquire() as conn:
            result = await conn.execute(
                """
                UPDATE colony_writeback_proposals
                SET approval_status = 'approved',
                    approved_by = $3,
                    approved_at = NOW()
                WHERE mission_id = $1::uuid
                  AND tenant_id = $2
                  AND approval_status = 'pending'
                """,
                mission_id,
                tenant_id,
                actor_id,
            )
        return int(result.rsplit(" ", 1)[-1])

    async def reject_pending_writebacks(self, mission_id: str, tenant_id: str, actor_id: str, reason: str) -> int:
        async with self.pool.acquire() as conn:
            result = await conn.execute(
                """
                UPDATE colony_writeback_proposals
                SET approval_status = 'rejected',
                    rejected_by = $3,
                    rejected_at = NOW(),
                    rejection_reason = $4
                WHERE mission_id = $1::uuid
                  AND tenant_id = $2
                  AND approval_status = 'pending'
                """,
                mission_id,
                tenant_id,
                actor_id,
                reason,
            )
        return int(result.rsplit(" ", 1)[-1])

    async def log_event(
        self,
        *,
        mission_id: str | None,
        tenant_id: str,
        event_type: str,
        actor: str | None,
        detail: dict[str, Any] | None = None,
    ) -> None:
        async with self.pool.acquire() as conn:
            await conn.execute(
                """
                INSERT INTO colony_event_log (mission_id, tenant_id, event_type, actor, detail, created_at)
                VALUES ($1::uuid, $2, $3, $4, $5::jsonb, NOW())
                """,
                mission_id,
                tenant_id,
                event_type,
                actor,
                json.dumps(detail or {}),
            )