"""
Phase 2 #2C: Loop Primitive + Verifier
Adds loop edges and verifier nodes to the orchestrator.
"""
import asyncio
from typing import Dict, List, Any, Callable


class LoopConfig:
    """Configuration for a loop in the agent DAG."""
    def __init__(self, max_iterations: int = 5, condition: str = "not_passed",
                 threshold: float = 0.8):
        self.max_iterations = max_iterations
        self.condition = condition
        self.threshold = threshold
        self.current_iteration = 0

    def should_continue(self, verifier_result: dict) -> bool:
        self.current_iteration += 1
        if self.current_iteration > self.max_iterations:
            return False
        if self.condition == "not_passed":
            return not verifier_result.get("passed", False)
        elif self.condition == "score_below":
            return verifier_result.get("score", 0) < self.threshold
        elif self.condition == "always":
            return self.current_iteration <= self.max_iterations
        return False


class VerifierNode:
    """Node that checks whether a goal has been achieved."""

    def __init__(self, goal: str, check_type: str = "quality"):
        self.goal = goal
        self.check_type = check_type

    async def verify(self, output: str, context: dict) -> dict:
        """Run verification via LLM to check if goal is met."""
        try:
            import httpx, json
            prompt = (
                f"Goal: {self.goal}\n"
                f"Output to verify:\n{output[:2000]}\n\n"
                f"Check if the output meets the goal. Return JSON:\n"
                f'{{"passed": true/false, "score": 0.0-1.0, "reason": "brief explanation"}}'
            )
            async with httpx.AsyncClient(timeout=60) as client:
                resp = await client.post(
                    "http://qihu-aiflow-litellm:4000/v1/chat/completions",
                    headers={
                        "Authorization": "Bearer sk-litellm-master-key-aiflow-2026",
                        "Content-Type": "application/json"
                    },
                    json={
                        "model": "deepseek-v4-flash",
                        "messages": [{"role": "user", "content": prompt}],
                        "temperature": 0.1,
                        "max_tokens": 200,
                    }
                )
                if resp.status_code == 200:
                    text = resp.json()["choices"][0]["message"]["content"].strip()
                    if text.startswith("```"):
                        text = text.split("\n", 1)[1][:-3]
                    return json.loads(text)
        except Exception:
            pass
        return {"passed": True, "score": 1.0, "reason": "verification skipped (offline fallback)"}


class LoopExecutor:
    """Executes a node graph with loop support."""

    def __init__(self, nodes: List[Dict], edges: List[Dict]):
        self.nodes = {n["id"]: n for n in nodes}
        self.edges = edges
        self.loop_configs: Dict[str, LoopConfig] = {}
        self.verifiers: Dict[str, VerifierNode] = {}

    def add_loop(self, from_node: str, to_node: str, config: LoopConfig):
        self.loop_configs[from_node] = config

    def add_verifier(self, node_id: str, verifier: VerifierNode):
        self.verifiers[node_id] = verifier

    async def execute(self, node_executor: Callable, initial_input: dict) -> dict:
        """Execute DAG with loop support."""
        results = {}
        to_process = ["input-1"]  # Start from input node

        while to_process:
            node_id = to_process.pop(0)
            node = self.nodes.get(node_id)
            if not node:
                continue

            # Check loop limit
            if node_id in self.loop_configs:
                lc = self.loop_configs[node_id]
                if lc.current_iteration >= lc.max_iterations:
                    continue

            output = await node_executor(node, initial_input, results)
            results[node_id] = output

            # Find outgoing edges
            next_nodes = [e["target"] for e in self.edges if e["source"] == node_id]
            for nn in next_nodes:
                if nn in self.loop_configs:
                    # Check if loop should continue
                    lc = self.loop_configs[nn]
                    if lc.should_continue(output):
                        target = [e["target"] for e in self.edges if e["source"] == nn]
                        if target:
                            to_process.append(target[0])  # Loop back
                else:
                    to_process.append(nn)

        return results