"""Wave scheduler for the Agent pipeline."""

from concurrent.futures import ThreadPoolExecutor, as_completed
from dataclasses import dataclass, field
from typing import Callable, Dict, Iterable, List, Optional

from backend.agents.base import AgentResult, AgentScope, ScopedAgent
from backend.agents.memory import ProjectMemory


@dataclass
class AgentStats:
    """Job-local accounting; never shared across concurrent jobs."""

    calls: int = 0
    scopes: int = 0
    merges: int = 0
    failed_scopes: List[str] = field(default_factory=list)

    def as_dict(self) -> Dict:
        return {
            "calls": self.calls,
            "scopes": self.scopes,
            "merges": self.merges,
            "failed_scopes": list(self.failed_scopes),
        }


class Orchestrator:
    """Coordinates bounded fan-out/fan-in waves against one ProjectMemory."""

    def __init__(
        self,
        memory: ProjectMemory,
        on_stage: Optional[Callable[[str], None]] = None,
    ) -> None:
        self.memory = memory
        self.on_stage = on_stage
        self.stats = AgentStats()

    def stage(self, name: str) -> None:
        print(f"\n=== {name} ===")
        if self.on_stage:
            self.on_stage(name)

    def initialize(self) -> Dict:
        """Initialize the job-local artifact store."""
        self.stage("Initialize agent pipeline")
        self.memory.dump()
        return self.stats.as_dict()

    def run_scopes(
        self,
        agent: ScopedAgent,
        scopes: Iterable[AgentScope],
        concurrency: int,
    ) -> List[AgentResult]:
        """Run independent scopes; one failure never aborts the wave."""
        scope_list = list(scopes)
        if not scope_list:
            return []
        results: List[AgentResult] = []
        with ThreadPoolExecutor(max_workers=max(1, concurrency)) as pool:
            futures = {pool.submit(agent.run, scope): scope for scope in scope_list}
            for future in as_completed(futures):
                scope = futures[future]
                self.stats.scopes += 1
                try:
                    result = future.result()
                except Exception as exc:
                    result = AgentResult(scope_id=scope.scope_id, error=str(exc))
                if result.error:
                    self.stats.failed_scopes.append(
                        f"{agent.name}:{scope.scope_id}: {result.error}"
                    )
                results.append(result)
        results.sort(key=lambda result: result.scope_id)
        return results