from __future__ import annotations

"""
Simple orchestration layer on top of BloomOSClient.

This is intentionally lightweight and rule-based. You can wire this into
your dot-command layer (e.g., a `bloomos.agent` dot-command that calls into
this module) or import it from experiments.
"""

from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional

from helix_cdc.bloomos.bloomos_client import BloomOSClient, TaskResult


@dataclass
class AgentStep:
    """One step in a multi-step BloomOS plan."""

    name: str
    task_type: str
    description: str
    anchors: Optional[List[str]] = None
    gate_policy: Optional[str] = None
    extra: Dict[str, Any] = field(default_factory=dict)


@dataclass
class AgentStepResult:
    step: AgentStep
    result: TaskResult


class BloomOSAgent:
    """
    A thin orchestration layer that:
    - Routes tasks to the right gate / anchors
    - Chains multiple BloomOS calls
    - Collects receipts and logits for analysis
    """

    def __init__(
        self,
        client: Optional[BloomOSClient] = None,
    ) -> None:
        self.client = client or BloomOSClient()

    # ------------------------------
    # High-level "auto routing"
    # ------------------------------

    def classify_goal(self, goal: str) -> AgentStep:
        """
        Very simple heuristic classifier that returns an AgentStep with the
        appropriate task_type, gate_policy, and anchors.

        You can replace this with a proper LLM-based classifier later.
        """
        g = goal.lower()

        # Heuristic: JSON-ish tasks
        if "json" in g or "schema" in g or "dict" in g:
            return AgentStep(
                name="json_structuring",
                task_type="json",
                description=goal,
                anchors=["json"],
                gate_policy="json_bias",
            )

        # Heuristic: mathy tasks (equations, integrals, etc.)
        math_tokens = ["integrate", "differentiate", "limit", "equation", "solve"]
        if any(tok in g for tok in math_tokens):
            return AgentStep(
                name="math_solve",
                task_type="math",
                description=goal,
                anchors=["math"],
                gate_policy="math_bias",
            )

        # Heuristic: code tasks
        code_tokens = [
            "python",
            "code",
            "function",
            "class",
            "refactor",
            "script",
            "bug",
        ]
        if any(tok in g for tok in code_tokens):
            return AgentStep(
                name="code_edit",
                task_type="code",
                description=goal,
                anchors=["code"],
                gate_policy="code_bias",
            )

        # Fallback: reasoning / balanced
        return AgentStep(
            name="reasoning",
            task_type="reason",
            description=goal,
            anchors=["math", "code", "json"],
            gate_policy="balanced",
        )

    def solve_goal(self, goal: str) -> AgentStepResult:
        """
        Single-step "agent" that auto-classifies the goal and dispatches
        a single BloomOS task.
        """
        step = self.classify_goal(goal)
        res = self._run_step(step)
        return AgentStepResult(step=step, result=res)

    # ------------------------------
    # Multi-step orchestrations
    # ------------------------------

    def run_plan(self, steps: List[AgentStep]) -> List[AgentStepResult]:
        """
        Execute a sequence of steps, returning the per-step results.

        For now, this is a straight-line execution. You can add branching,
        retries, or dynamic plan expansion later.
        """
        results: List[AgentStepResult] = []
        for step in steps:
            res = self._run_step(step)
            results.append(AgentStepResult(step=step, result=res))
        return results

    def _run_step(self, step: AgentStep) -> TaskResult:
        """
        Internal helper that maps an AgentStep to a BloomOSClient call.
        """
        return self.client.run_task(
            task_type=step.task_type,
            description=step.description,
            anchors=step.anchors,
            gate_policy=step.gate_policy,
            extra=step.extra,
        )

    # Example of a hand-crafted multi-step agent

    def analyze_code_then_write_tests(
        self,
        code_snippet: str,
        language: str = "python",
    ) -> List[AgentStepResult]:
        """
        Example composite agent:
            1. Ask Helix64/BloomOS to review the code for issues.
            2. Ask it to write tests for the (possibly fixed) code.
        """
        steps = [
            AgentStep(
                name="code_review",
                task_type="code",
                description=(
                    f"Review the following {language} code for bugs, "
                    f"edge cases, and improvements, then suggest fixes:\n\n"
                    f"{code_snippet}"
                ),
                anchors=["code"],
                gate_policy="code_bias",
            ),
            AgentStep(
                name="write_tests",
                task_type="code",
                description=(
                    f"Write unit tests for the following {language} code. "
                    f"Assume pytest is available:\n\n{code_snippet}"
                ),
                anchors=["code"],
                gate_policy="code_bias",
            ),
        ]
        return self.run_plan(steps)
