awoooi/apps/api/src/services/awooop_ansible_learning_writeback.py

"""Durable RAG and PlayBook trust writeback for the Ansible executor."""

from __future__ import annotations

import hashlib
import json
import os
import re
from typing import Any

import structlog
from sqlalchemy import select, text

from src.db.base import get_db_context
from src.db.models import PlaybookRecord
from src.services.awooop_ansible_audit_service import get_ansible_catalog_item
from src.services.knowledge_service import get_knowledge_service
from src.utils.timezone import now_taipei

logger = structlog.get_logger(__name__)

_CANONICAL_PLAYBOOK_PREFIX = "PB-ANSIBLE-"


def canonical_ansible_playbook_id(catalog_id: str) -> str:
    """Return a stable PlayBook primary key for one allowlisted catalog row."""

    normalized = re.sub(
        r"[^A-Z0-9]+",
        "-",
        str(catalog_id or "").removeprefix("ansible:").upper(),
    ).strip("-")
    candidate = f"{_CANONICAL_PLAYBOOK_PREFIX}{normalized or 'UNKNOWN'}"
    if len(candidate) <= 32:
        return candidate
    digest = hashlib.sha256(candidate.encode("utf-8")).hexdigest()[:8].upper()
    return f"{candidate[:23].rstrip('-')}-{digest}"


def _json_dict(value: Any) -> dict[str, Any]:
    if isinstance(value, dict):
        return value
    if isinstance(value, str):
        try:
            parsed = json.loads(value)
        except json.JSONDecodeError:
            return {}
        return parsed if isinstance(parsed, dict) else {}
    return {}


def resolve_ansible_playbook_identity(
    catalog_id: str,
) -> dict[str, Any] | None:
    """Resolve one allowlisted catalog action to its canonical PlayBook row."""

    catalog = get_ansible_catalog_item(catalog_id)
    if not catalog:
        return None
    playbook_path = str(catalog.get("playbook_path") or "")
    canonical_id = canonical_ansible_playbook_id(catalog_id)
    identity_payload = {
        "canonical_playbook_id": canonical_id,
        "catalog_id": catalog_id,
        "playbook_path": playbook_path,
    }
    return {
        "schema_version": "ansible_playbook_identity_v1",
        **identity_payload,
        "identity_fingerprint": hashlib.sha256(
            json.dumps(
                identity_payload,
                ensure_ascii=True,
                sort_keys=True,
                separators=(",", ":"),
            ).encode("utf-8")
        ).hexdigest(),
    }


def _playbook_row_fingerprint(playbook: PlaybookRecord) -> str:
    updated_at = playbook.updated_at.isoformat() if playbook.updated_at else None
    payload = {
        "playbook_id": playbook.playbook_id,
        "project_id": playbook.project_id,
        "version": int(playbook.version or 0),
        "success_count": int(playbook.success_count or 0),
        "failure_count": int(playbook.failure_count or 0),
        "trust_score": round(float(playbook.trust_score or 0.0), 6),
        "updated_at": updated_at,
    }
    return hashlib.sha256(
        json.dumps(
            payload,
            ensure_ascii=True,
            sort_keys=True,
            separators=(",", ":"),
        ).encode("utf-8")
    ).hexdigest()


def _playbook_trust_readback(
    *,
    playbook: PlaybookRecord,
    identity: dict[str, Any],
    verification_result: str,
) -> dict[str, Any]:
    return {
        "schema_version": "ansible_playbook_trust_writeback_v2",
        "identity_schema_version": identity["schema_version"],
        "identity_fingerprint": identity["identity_fingerprint"],
        "canonical_playbook_id": playbook.playbook_id,
        "catalog_id": identity["catalog_id"],
        "playbook_path": identity["playbook_path"],
        "playbook_row_version": int(playbook.version or 0),
        "playbook_row_updated_at": (
            playbook.updated_at.isoformat() if playbook.updated_at else None
        ),
        "playbook_row_fingerprint": _playbook_row_fingerprint(playbook),
        "success_count": int(playbook.success_count or 0),
        "failure_count": int(playbook.failure_count or 0),
        "trust_observation_count": (
            int(playbook.success_count or 0)
            + int(playbook.failure_count or 0)
        ),
        "trust_score": round(float(playbook.trust_score or 0.0), 6),
        "verification_result": verification_result,
        "learning_recorded": True,
        "trust_updated": True,
        "repository_write_acknowledged": True,
        "repository_readback_verified": True,
        "operation_receipt_readback_verified": True,
        "durable_write_acknowledged": True,
        "source_table": "playbooks",
        "writer_source_sha": os.getenv("AWOOOI_BUILD_COMMIT_SHA", "")
        .strip()
        .lower()
        or None,
        "raw_log_payload_stored": False,
        "secret_value_stored": False,
    }


def _existing_trust_readback_matches(
    existing_trust: dict[str, Any],
    *,
    playbook: PlaybookRecord,
    identity: dict[str, Any],
) -> bool:
    return bool(
        existing_trust.get("schema_version")
        == "ansible_playbook_trust_writeback_v2"
        and existing_trust.get("canonical_playbook_id")
        == identity["canonical_playbook_id"]
        and existing_trust.get("identity_fingerprint")
        == identity["identity_fingerprint"]
        and existing_trust.get("playbook_row_version")
        == int(playbook.version or 0)
        and existing_trust.get("playbook_row_fingerprint")
        == _playbook_row_fingerprint(playbook)
        and existing_trust.get("repository_readback_verified") is True
        and existing_trust.get("operation_receipt_readback_verified") is True
        and existing_trust.get("learning_recorded") is True
        and existing_trust.get("trust_updated") is True
    )


async def record_ansible_playbook_trust_writeback(
    *,
    project_id: str,
    automation_run_id: str,
    incident_id: str,
    incident_db_id: int | None,
    catalog_id: str,
    playbook_path: str,
    apply_op_id: str,
    verification_result: str,
) -> dict[str, Any] | None:
    """Atomically update canonical PlayBook trust and its learning receipt."""

    catalog = get_ansible_catalog_item(catalog_id)
    identity = resolve_ansible_playbook_identity(catalog_id)
    if not catalog or identity is None:
        return None
    canonical_id = str(identity["canonical_playbook_id"])
    canonical_playbook_path = str(identity["playbook_path"])
    if str(playbook_path) != canonical_playbook_path:
        logger.warning(
            "ansible_learning_playbook_identity_mismatch",
            incident_id=incident_id,
            catalog_id=catalog_id,
            apply_op_id=apply_op_id,
        )
        return None
    playbook_path = canonical_playbook_path
    success = verification_result == "success"
    now = now_taipei()

    try:
        async with get_db_context(project_id) as db:
            existing_receipt = await db.execute(
                text("""
                    SELECT op_id::text AS op_id, input, output
                    FROM automation_operation_log
                    WHERE operation_type = 'ansible_learning_writeback_recorded'
                      AND parent_op_id = CAST(:apply_op_id AS uuid)
                    ORDER BY created_at DESC
                    LIMIT 1
                    FOR UPDATE
                """),
                {"apply_op_id": apply_op_id},
            )
            existing_row = existing_receipt.mappings().first()
            existing_output = _json_dict(
                existing_row.get("output") if existing_row else None
            )
            existing_trust = _json_dict(
                existing_output.get("playbook_trust_writeback")
            )

            selected = await db.execute(
                select(PlaybookRecord)
                .where(PlaybookRecord.playbook_id == canonical_id)
                .with_for_update()
            )
            playbook = selected.scalar_one_or_none()
            if existing_row:
                if playbook is None:
                    logger.warning(
                        "ansible_learning_receipt_playbook_row_missing",
                        incident_id=incident_id,
                        catalog_id=catalog_id,
                        apply_op_id=apply_op_id,
                    )
                    return None
                if _existing_trust_readback_matches(
                    existing_trust,
                    playbook=playbook,
                    identity=identity,
                ):
                    return existing_trust
                logger.warning(
                    "ansible_learning_receipt_readback_not_verified",
                    incident_id=incident_id,
                    catalog_id=catalog_id,
                    apply_op_id=apply_op_id,
                    receipt_schema_version=existing_trust.get(
                        "schema_version"
                    ),
                )
                return None
            if playbook is None:
                playbook = PlaybookRecord(
                    playbook_id=canonical_id,
                    project_id=project_id,
                    name=f"Ansible catalog: {catalog_id}",
                    description=(
                        "Canonical trust record for the allowlisted Ansible "
                        f"catalog action {playbook_path}."
                    ),
                    status="approved",
                    source="yaml_rule",
                    symptom_pattern={
                        "alert_names": [],
                        "affected_services": list(catalog.get("domains") or []),
                        "severity_range": ["P0", "P1", "P2", "P3"],
                        "label_patterns": {},
                        "keywords": list(catalog.get("keywords") or [])[:50],
                    },
                    repair_steps=[
                        {
                            "step_number": 1,
                            "action_type": "script",
                            "command": playbook_path,
                            "expected_result": "independent post-verifier passes",
                            "rollback_command": None,
                            "requires_approval": bool(
                                catalog.get("approval_required")
                            ),
                            "risk_level": str(
                                catalog.get("risk_level") or "medium"
                            ).upper(),
                        }
                    ],
                    estimated_duration_minutes=5,
                    source_incident_ids=[incident_id],
                    version=1,
                    ai_confidence=1.0,
                    success_count=0,
                    failure_count=0,
                    trust_score=0.3,
                    approved_by="ansible_catalog_policy",
                    approved_at=now,
                    tags=["ansible_catalog", catalog_id],
                    notes="Generated from the controlled executor allowlist.",
                    requires_approval_level=(
                        "critical"
                        if catalog.get("break_glass_required") is True
                        else "standard"
                        if catalog.get("approval_required") is True
                        else "auto"
                    ),
                    stateful_targets=[],
                    requires_pre_backup=False,
                    review_required=False,
                    created_at=now,
                    updated_at=now,
                )
                db.add(playbook)
                await db.flush()

            source_incident_ids = list(playbook.source_incident_ids or [])
            if incident_id and incident_id not in source_incident_ids:
                source_incident_ids.append(incident_id)
                playbook.source_incident_ids = source_incident_ids

            if success:
                playbook.success_count = int(playbook.success_count or 0) + 1
                playbook.trust_score = 0.9 * float(playbook.trust_score) + 0.1
            else:
                playbook.failure_count = int(playbook.failure_count or 0) + 1
                playbook.trust_score = 0.8 * float(playbook.trust_score)
            playbook.trust_score = max(0.0, min(1.0, playbook.trust_score))
            playbook.last_used_at = now
            playbook.updated_at = now
            await db.flush()
            await db.refresh(playbook)
            trust_writeback = _playbook_trust_readback(
                playbook=playbook,
                identity=identity,
                verification_result=verification_result,
            )
            input_payload = {
                "schema_version": "ansible_learning_writeback_receipt_v3",
                "automation_run_id": automation_run_id,
                "incident_id": incident_id,
                "catalog_id": catalog_id,
                "canonical_playbook_id": canonical_id,
                "identity_schema_version": identity["schema_version"],
                "identity_fingerprint": identity["identity_fingerprint"],
                "playbook_path": playbook_path,
                "apply_op_id": apply_op_id,
                "verification_result": verification_result,
                "learning_repository": "playbooks",
                "writer_source_sha": trust_writeback["writer_source_sha"],
                "stores_raw_logs": False,
                "stores_secret_values": False,
            }
            output_payload = {
                "learning_recorded": True,
                "trust_updated": True,
                "success": success,
                "repository_readback_verified": True,
                "playbook_trust_writeback": trust_writeback,
            }
            await db.execute(
                text("""
                    INSERT INTO automation_operation_log (
                        operation_type, actor, status, incident_id,
                        input, output, dry_run_result, parent_op_id, tags
                    ) VALUES (
                        'ansible_learning_writeback_recorded',
                        'ansible_controlled_apply_worker',
                        'success',
                        :incident_db_id,
                        CAST(:input AS jsonb),
                        CAST(:output AS jsonb),
                        '{}'::jsonb,
                        CAST(:parent_op_id AS uuid),
                        :tags
                    )
                """),
                {
                    "incident_db_id": incident_db_id,
                    "input": json.dumps(input_payload, ensure_ascii=False),
                    "output": json.dumps(output_payload, ensure_ascii=False),
                    "parent_op_id": apply_op_id,
                    "tags": [
                        "ansible",
                        "controlled_apply",
                        "learning_writeback",
                        "playbook_trust",
                        "ai_agent_auto_execution",
                        f"automation_run_id:{automation_run_id}",
                    ],
                },
            )
            await db.flush()
            operation_readback = await db.execute(
                text("""
                    SELECT status, input, output
                    FROM automation_operation_log
                    WHERE operation_type
                            = 'ansible_learning_writeback_recorded'
                      AND parent_op_id = CAST(:apply_op_id AS uuid)
                    ORDER BY created_at DESC
                    LIMIT 1
                """),
                {"apply_op_id": apply_op_id},
            )
            operation_row = operation_readback.mappings().first()
            operation_output = _json_dict(
                operation_row.get("output") if operation_row else None
            )
            operation_trust = _json_dict(
                operation_output.get("playbook_trust_writeback")
            )
            if not (
                operation_row
                and operation_row.get("status") == "success"
                and operation_output.get("learning_recorded") is True
                and operation_output.get("trust_updated") is True
                and operation_output.get("repository_readback_verified")
                is True
                and operation_trust.get("playbook_row_fingerprint")
                == trust_writeback["playbook_row_fingerprint"]
            ):
                raise RuntimeError(
                    "ansible_learning_operation_receipt_readback_failed"
                )
            return trust_writeback
    except Exception as exc:
        logger.warning(
            "ansible_canonical_playbook_trust_writeback_failed",
            incident_id=incident_id,
            catalog_id=catalog_id,
            apply_op_id=apply_op_id,
            error=str(exc),
        )
        return None


async def ensure_ansible_rag_writeback(
    *,
    project_id: str,
    incident_id: str,
    apply_op_id: str,
) -> dict[str, Any] | None:
    """Persist and read back the KM embedding used by semantic retrieval."""

    path_type = f"ansible_apply_receipt:{apply_op_id[:8]}"
    try:
        async with get_db_context(project_id) as db:
            selected = await db.execute(
                text("""
                    SELECT
                        id,
                        title,
                        content,
                        embedding IS NOT NULL AS embedding_persisted
                    FROM knowledge_entries
                    WHERE related_incident_id = :incident_id
                      AND path_type = :path_type
                    LIMIT 1
                """),
                {"incident_id": incident_id, "path_type": path_type},
            )
            row = selected.mappings().first()
        if not row:
            return None
        entry_id = str(row["id"])
        if row.get("embedding_persisted") is not True:
            saved = await get_knowledge_service().ensure_entry_embedding(
                entry_id,
                str(row.get("title") or ""),
                str(row.get("content") or ""),
                project_id=project_id,
            )
            if not saved:
                fallback = await _ensure_rag_chunk_fallback(
                    project_id=project_id,
                    incident_id=incident_id,
                    apply_op_id=apply_op_id,
                    entry_id=entry_id,
                    title=str(row.get("title") or ""),
                    content=str(row.get("content") or ""),
                )
                if fallback is not None:
                    return fallback
                await _record_rag_attempt(
                    project_id=project_id,
                    apply_op_id=apply_op_id,
                    status="primary_and_chunk_fallback_not_persisted",
                )
                return None
        async with get_db_context(project_id) as db:
            verified = await db.execute(
                text("""
                    SELECT embedding IS NOT NULL
                    FROM knowledge_entries
                    WHERE id = :entry_id
                """),
                {"entry_id": entry_id},
            )
            if verified.scalar() is not True:
                return None
        return {
            "schema_version": "ansible_rag_writeback_v1",
            "knowledge_entry_id": entry_id,
            "embedding_persisted": True,
            "source_table": "knowledge_entries",
            "retrieval_surface": "knowledge_semantic_search",
            "durable_write_acknowledged": True,
        }
    except Exception as exc:
        logger.warning(
            "ansible_rag_writeback_failed",
            incident_id=incident_id,
            apply_op_id=apply_op_id,
            error=str(exc),
        )
        await _record_rag_attempt(
            project_id=project_id,
            apply_op_id=apply_op_id,
            status=type(exc).__name__,
        )
        return None


async def _ensure_rag_chunk_fallback(
    *,
    project_id: str,
    incident_id: str,
    apply_op_id: str,
    entry_id: str,
    title: str,
    content: str,
) -> dict[str, Any] | None:
    """Index through the longer-timeout RAG lane and verify durable chunks."""

    source_id = f"ansible-learning:{apply_op_id}"
    try:
        from src.repositories import rag_chunk_repository
        from src.services.knowledge_rag_service import get_knowledge_rag_service

        chunk_count = await rag_chunk_repository.count_chunks_by_source_id(
            source_id,
            project_id=project_id,
        )
        if chunk_count <= 0:
            indexed = await get_knowledge_rag_service().index_document(
                source="ansible_learning",
                source_id=source_id,
                title=title,
                text=content,
                metadata={
                    "incident_id": incident_id,
                    "apply_op_id": apply_op_id,
                    "knowledge_entry_id": entry_id,
                    "raw_log_payload_stored": False,
                    "secret_value_stored": False,
                },
                project_id=project_id,
            )
            chunk_count = await rag_chunk_repository.count_chunks_by_source_id(
                source_id,
                project_id=project_id,
            )
            if not indexed or chunk_count <= 0:
                return None

        embedding_copied_to_km = False
        try:
            async with get_db_context(project_id) as db:
                copied = await db.execute(
                    text("""
                        WITH source_embedding AS (
                            SELECT embedding
                            FROM rag_chunks
                            WHERE source_id = :source_id
                              AND embedding IS NOT NULL
                            ORDER BY id
                            LIMIT 1
                        )
                        UPDATE knowledge_entries AS km
                        SET embedding = source_embedding.embedding,
                            updated_at = NOW()
                        FROM source_embedding
                        WHERE km.id = :entry_id
                        RETURNING km.embedding IS NOT NULL
                    """),
                    {"source_id": source_id, "entry_id": entry_id},
                )
                embedding_copied_to_km = copied.scalar_one_or_none() is True
        except Exception as exc:
            logger.warning(
                "ansible_rag_km_embedding_copy_failed",
                incident_id=incident_id,
                apply_op_id=apply_op_id,
                source_id=source_id,
                error=str(exc),
            )

        return {
            "schema_version": "ansible_rag_writeback_v2",
            "knowledge_entry_id": entry_id,
            "embedding_persisted": embedding_copied_to_km,
            "source_table": (
                "knowledge_entries" if embedding_copied_to_km else "rag_chunks"
            ),
            "source_id": source_id,
            "rag_chunk_count": chunk_count,
            "chunk_index_verified": True,
            "retrieval_surface": (
                "knowledge_semantic_search"
                if embedding_copied_to_km
                else "knowledge_rag_query"
            ),
            "fallback_used": True,
            "durable_write_acknowledged": True,
        }
    except Exception as exc:
        logger.warning(
            "ansible_rag_chunk_fallback_failed",
            incident_id=incident_id,
            apply_op_id=apply_op_id,
            error=str(exc),
        )
        return None


async def _record_rag_attempt(
    *,
    project_id: str,
    apply_op_id: str,
    status: str,
) -> None:
    try:
        async with get_db_context(project_id) as db:
            await db.execute(
                text("""
                    UPDATE automation_operation_log
                    SET input = coalesce(input, '{}'::jsonb) || jsonb_build_object(
                        'rag_writeback_attempt',
                        jsonb_build_object(
                            'attempted_at', NOW(),
                            'status', CAST(:status AS text)
                        )
                    )
                    WHERE op_id = CAST(:apply_op_id AS uuid)
                """),
                {"apply_op_id": apply_op_id, "status": status[:80]},
            )
    except Exception as exc:
        logger.warning(
            "ansible_rag_attempt_receipt_failed",
            apply_op_id=apply_op_id,
            error=str(exc),
        )