src/engine/quadruflow.py

from __future__ import annotations

import asyncio
import hashlib
import os
import time
from datetime import datetime
from typing import Dict, Any

from src.cognition.curriculum import CurriculumManager, CurriculumTask
from src.memory.ledger import QuantumResistantLedger
from src.engine.validator import ArtifactValidator
from src.cognition._constants import logger
from src.core.lfm_controller import LFMController


class QuadruflowOrchestrator:
    """The Neurosynthetic Quadruflow Engine implementing four discrete cognitive flows."""

    def __init__(self, random_seed: int | None = None):
        logger.info("Initializing Quadruflow Engine Core...")
        self.curriculum = CurriculumManager(random_seed=random_seed)
        self.ledger = QuantumResistantLedger()
        self.validator = ArtifactValidator()

        candidates = [
            os.path.expanduser("~/.vitalis/models/Llama-3.2-3B-Instruct-Q4_K_M.gguf"),
            os.path.expanduser("~/.vitalis/models/LFM2.5-1.2B-Instruct-Q4_K_M.gguf"),
            "Llama-3.2-3B-Instruct-Q4_K_M.gguf",
            "LFM2.5-1.2B-Instruct-Q4_K_M.gguf",
        ]
        model_path = ""
        for c in candidates:
            if os.path.exists(c):
                model_path = c
                break
        if not model_path:
            model_path = candidates[0]
        logger.info("Wiring LFMController -> %s", model_path)
        self.controller = LFMController(
            model_path=model_path,
            n_ctx=8192,
            n_threads=6,
            n_gpu_layers=-1,
        )

        if not self.ledger.verify_integrity():
            raise SecurityError("Critical Error: Cryptographic ledger validation failed during boot sequence.")
        logger.info("[Evaluation Ready] Core engine verified and integrated with evaluation hooks.")

    def _normalize_val_report(self, v_out) -> dict:
        if isinstance(v_out, dict):
            return v_out
        if isinstance(v_out, tuple) and len(v_out) == 2:
            valid, reason = v_out
            valid = bool(valid)
            return {
                "valid": valid,
                "errors": [] if valid else [str(reason)],
                "score_components": {
                    "schema": 1.0 if valid else 0.0,
                    "semantics": 1.0 if valid else 0.0,
                    "length": 1.0,
                },
            }
        return {
            "valid": getattr(v_out, "valid", False),
            "errors": getattr(v_out, "errors", []),
            "score_components": getattr(v_out, "score_components", {"schema": 0.0, "semantics": 0.0, "length": 0.0}),
        }

    def flow1_ingest(self, state: dict) -> dict:
        if "task_id" not in state or "prompt" not in state:
            raise ValueError("Schema validation failed: missing task_id or prompt")
        return {
            "task_id": str(state["task_id"]),
            "prompt": str(state["prompt"]),
            "expected_type": state.get("expected_type", "code"),
            "provenance": "vitalis_devcore_ingest",
        }

    def flow2_simulate(self, state: dict, seed: int) -> dict:
        formatted_prompt = f"<|im_start|>user\n{state['prompt']}<|im_end|>\n<|im_start|>assistant\n"
        text = self.controller.execute_raw(formatted_prompt, max_tokens=256, temperature=0.0)
        return {"text": text, "task_id": state["task_id"]}

    def flow3_debug(self, candidate: dict, errors: list) -> tuple[dict, dict]:
        repair_prompt = (
            f"Fix the following errors in your previous output:\n"
            f"Errors: {', '.join(errors)}\n"
            f"Previous Output:\n{candidate['text']}"
        )
        return (
            {"task_id": candidate["task_id"], "prompt": repair_prompt},
            {"repair_prompt": repair_prompt, "previous_errors": errors},
        )

    def flow4_attest(self, candidate: dict, corrections: dict | None, val_report: dict, retries: int, runtime_ms: float, attested: bool, rejection_reason: str | None = None) -> dict:
        det_hash = hashlib.sha256(candidate["text"].encode("utf-8")).hexdigest()
        return {
            "artifact_id": f"art_{candidate['task_id']}_{int(time.time())}",
            "task_id": candidate["task_id"],
            "result": {"output": candidate["text"]},
            "validator": {
                "valid": val_report.get("valid", False),
                "errors": val_report.get("errors", []),
                "score_components": val_report.get("score_components", {"schema": 0.0, "semantics": 0.0, "length": 0.0}),
            },
            "attestation": {
                "hash": det_hash,
                "signature": "<placeholder>",
                "timestamp": datetime.utcnow().isoformat() + "Z",
                "attested": attested,
            },
            "meta": {
                "retries": retries,
                "runtime_ms": int(runtime_ms),
                "rejection_reason": rejection_reason,
            },
        }

    def run_cognitive_cycle(self, task: CurriculumTask, seed: int) -> dict:
        start_time = time.time()
        initial_state = {
            "task_id": task.task_id,
            "prompt": getattr(task, "prompt", f"Complete task {task.task_id}"),
            "expected_type": task.expected_type,
        }
        state = self.flow1_ingest(initial_state)
        retries = 0
        rejection_reason = None
        candidate = {"text": "", "task_id": task.task_id}
        corrections = None
        val_report: dict = {
            "valid": False,
            "errors": [],
            "score_components": {"schema": 0.0, "semantics": 0.0, "length": 0.0},
        }

        while retries <= 2:
            candidate = self.flow2_simulate(state, seed)
            v_out = self.validator.validate(candidate["text"])
            val_report = self._normalize_val_report(v_out)
            if val_report.get("valid"):
                break
            retries += 1
            if retries <= 2:
                state, corrections = self.flow3_debug(candidate, val_report.get("errors", []))
            else:
                rejection_reason = "Max retries exceeded without valid artifact"

        runtime_ms = (time.time() - start_time) * 1000
        attested = val_report.get("valid", False) and (rejection_reason is None)
        artifact = self.flow4_attest(candidate, corrections, val_report, min(retries, 2), runtime_ms, attested, rejection_reason)

        sc = val_report.get("score_components", {"schema": 0.0, "semantics": 0.0, "length": 0.0})
        risk = (
            0.5 * (1.0 - sc.get("schema", 0.0))
            + 0.3 * (1.0 - sc.get("semantics", 0.0))
            + 0.2 * (min(retries, 2) / 2.0)
        )

        metrics_payload = (
            f"status={'SUCCESS' if attested else 'REJECTED'}"
            f"|risk={round(risk, 4)}"
            f"|duration={int(runtime_ms)}"
            f"|attested={str(attested).lower()}"
            f"|rejection_reason={rejection_reason}"
        )
        block = self.ledger.append_record(task_id=task.task_id, outcome_metrics=metrics_payload)
        artifact["attestation"]["signature"] = getattr(block, "signature", "<placeholder>")
        self.curriculum.record_result(success=attested, risk_encountered=risk)
        return artifact

    async def execute_closed_loop_cycle(self) -> Dict[str, Any]:
        task: CurriculumTask = self.curriculum.generate_next_task()
        logger.info("Executing Task Channel: %s [Tier %d]", task.task_id, task.tier)
        loop = asyncio.get_running_loop()
        try:
            artifact = await loop.run_in_executor(None, lambda: self.run_cognitive_cycle(task, 42))
            success = artifact["attestation"]["attested"]
            sig = artifact["attestation"]["signature"]
            block_idx = artifact.get("ledger_block_index", 1)
        except Exception as err:
            logger.error("Execution exception encountered on %s: %s", task.task_id, str(err))
            block = self.ledger.append_record(
                task_id=task.task_id,
                outcome_metrics="status=CRASHED|risk=1.0|duration=0|attested=false|rejection_reason=catastrophic_failure",
            )
            success = False
            sig = getattr(block, "signature", "<placeholder>")
            block_idx = getattr(block, "index", 1)

        return {
            "task_id": task.task_id,
            "success": success,
            "ledger_block_index": block_idx,
            "ledger_signature": sig,
            "curriculum_state": self.curriculum.export_state(),
        }

    def verify_system_health(self) -> bool:
        return self.ledger.verify_integrity()

    def shutdown(self) -> None:
        self.controller.shutdown()


class SecurityError(Exception):
    """Raised when cryptographic data structures show validation failures."""