from __future__ import annotations

import json
import re

from pydantic import BaseModel

from nexus.protocols.providers.base import AbstractProvider, LLMRequest, Message
from nexus.resilience.errors import OrchestratorError


class SketchNode(BaseModel):
    node_id: str
    label: str
    description: str
    depends_on: list[str]
    estimated_complexity: float
    suggested_tools: list[str]


class ExecutionSketch(BaseModel):
    nodes: list[SketchNode]
    edges: list[tuple[str, str]]
    terminal_node_id: str
    reasoning: str


class SketchPlanner:
    """Single-call sketch planner producing a minimal execution graph."""

    def __init__(self, provider: AbstractProvider) -> None:
        self.provider = provider

    async def plan(
        self,
        objective: str,
        available_tools: list[str] | None = None,
        available_agents: list[str] | None = None,
    ) -> ExecutionSketch:
        """Generate and validate an execution sketch."""
        prompt = self._build_prompt(objective, available_tools or [], available_agents or [])
        req = LLMRequest(messages=[Message(role="user", content=prompt)], temperature=0.2, max_tokens=2000)
        resp = await self.provider.complete(req)
        sketch = self._parse_response(resp.content)

        node_ids = {n.node_id for n in sketch.nodes}
        if sketch.terminal_node_id not in node_ids:
            raise OrchestratorError("terminal_node_id missing in nodes")
        for n in sketch.nodes:
            for dep in n.depends_on:
                if dep not in node_ids:
                    raise OrchestratorError(f"invalid dependency {dep} in node {n.node_id}")
        return sketch

    def _build_prompt(self, objective: str, tools: list[str], agents: list[str]) -> str:
        """Build strict JSON-only planning prompt."""
        return (
            "Return ONLY valid JSON matching this schema: "
            "{nodes:[{node_id,label,description,depends_on,estimated_complexity,suggested_tools}],"
            "edges:[[from_id,to_id]],terminal_node_id,reasoning}. "
            f"Objective: {objective}\nAvailable tools: {tools}\nAvailable agents: {agents}"
        )

    def _parse_response(self, response_text: str) -> ExecutionSketch:
        """Extract JSON payload and parse into model."""
        block = re.search(r"```json\s*(.*?)\s*```", response_text, flags=re.DOTALL)
        raw = block.group(1) if block else response_text.strip()
        data = json.loads(raw)
        return ExecutionSketch.model_validate(data)
