"""Planning phase for deliberate."""

import asyncio
import re
from dataclasses import dataclass

from deliberate.adapters.base import ModelAdapter
from deliberate.budget.tracker import BudgetTracker
from deliberate.prompts.planning import DEBATE_PROMPT, JUDGE_PROMPT, PLANNING_PROMPT
from deliberate.types import DebateMessage, Plan


@dataclass
class PlanningPhase:
    """Orchestrates the planning phase of the jury workflow.

    Multiple agents propose plans, optionally debate them,
    and a judge (or voting) selects the best plan.
    """

    agents: list[str]
    adapters: dict[str, ModelAdapter]
    budget: BudgetTracker
    debate_enabled: bool = False
    debate_rounds: int = 1
    selection_method: str = "first"  # first | llm_judge | borda
    judge_agent: str | None = None

    async def run(self, task: str) -> Plan | None:
        """Run the planning phase.

        Args:
            task: The task description.

        Returns:
            The selected plan, or None if no plans were generated.
        """
        plans = await self._collect_plans(task)

        if not plans:
            return None

        if self.debate_enabled and len(plans) > 1:
            await self._run_debate(task, plans)

        return await self._select_plan(task, plans)

    async def _collect_plans(self, task: str) -> list[Plan]:
        """Collect plans from all planning agents."""

        async def get_plan(agent_name: str) -> Plan | None:
            adapter = self.adapters.get(agent_name)
            if not adapter:
                return None

            try:
                prompt = PLANNING_PROMPT.format(task=task)
                response = await adapter.call(prompt)

                self.budget.record_usage(
                    agent_name,
                    response.token_usage,
                    adapter.estimate_cost(response.token_usage),
                )

                return Plan(
                    id=f"plan-{agent_name}",
                    agent=agent_name,
                    content=response.content,
                    token_usage=response.token_usage,
                )
            except Exception as e:
                print(f"Warning: {agent_name} planning failed: {e}")
                return None

        results = await asyncio.gather(*[get_plan(a) for a in self.agents])
        return [p for p in results if p is not None]

    async def _run_debate(
        self,
        task: str,
        plans: list[Plan],
    ) -> list[DebateMessage]:
        """Run debate rounds where agents critique each other's plans."""
        messages: list[DebateMessage] = []

        for round_num in range(self.debate_rounds):
            for i, agent_name in enumerate(self.agents):
                adapter = self.adapters.get(agent_name)
                if not adapter:
                    continue

                # Each agent reviews the next agent's plan (circular)
                other_plan = plans[(i + 1) % len(plans)]
                prompt = DEBATE_PROMPT.format(
                    task=task,
                    plan=other_plan.content,
                )

                try:
                    response = await adapter.call(prompt)
                    self.budget.record_usage(agent_name, response.token_usage)

                    messages.append(
                        DebateMessage(
                            agent=agent_name,
                            content=response.content,
                            round=round_num,
                            reply_to=other_plan.agent,
                        )
                    )
                except Exception as e:
                    print(f"Warning: {agent_name} debate failed: {e}")

        return messages

    async def _select_plan(self, task: str, plans: list[Plan]) -> Plan:
        """Select the best plan using the configured method."""
        if len(plans) == 1 or self.selection_method == "first":
            return plans[0]

        if self.selection_method == "llm_judge" and self.judge_agent:
            return await self._judge_select(task, plans)

        # Default to first plan
        return plans[0]

    async def _judge_select(self, task: str, plans: list[Plan]) -> Plan:
        """Use an LLM judge to select the best plan."""
        adapter = self.adapters.get(self.judge_agent)
        if not adapter:
            return plans[0]

        # Format plans for comparison
        plans_text = "\n\n---\n\n".join(
            f"## Plan {i + 1} (by {p.agent})\n\n{p.content}" for i, p in enumerate(plans)
        )

        prompt = JUDGE_PROMPT.format(
            task=task,
            plans=plans_text,
            num_plans=len(plans),
        )

        try:
            response = await adapter.call(prompt)
            self.budget.record_usage(self.judge_agent, response.token_usage)

            # Parse "Plan N" from response
            match = re.search(r"[Pp]lan\s*(\d+)", response.content)
            if match:
                idx = int(match.group(1)) - 1
                if 0 <= idx < len(plans):
                    return plans[idx]

        except Exception as e:
            print(f"Warning: Judge selection failed: {e}")

        return plans[0]
