import asyncio
import shutil
from pathlib import Path

from mixersystem.data.repository import context_scope, get_context, parse_result, update_frontmatter
from mixersystem.workflows.plan import builder
from mixersystem.workflows.plan.agents import merger
from mixersystem.workflows.plan.agents import formatter
from mixersystem.workflows.plan import review

MAX_REVISIONS = 2
_SUPPORTED_PROVIDERS = {"claude", "gemini", "codex"}


def _alpha_suffix(index: int) -> str:
    if 0 <= index < 26:
        return chr(ord("a") + index)
    return f"{index + 1:02d}"


def _prepare_branches(work_folder: str, branch_count: int) -> list[str]:
    """Create branch work folders and copy root task.md into each."""
    root = Path(work_folder)
    task_path = root / "task.md"
    if not task_path.is_file():
        raise FileNotFoundError(f"Missing task file: {task_path}")
    if branch_count < 2:
        raise ValueError("branch_count must be >= 2 when preparing branches")

    branches_root = root / "branches"
    branch_folders = []
    for i in range(branch_count):
        branch = branches_root / f"plan_{_alpha_suffix(i)}"
        branch.mkdir(parents=True, exist_ok=True)
        shutil.copy2(task_path, branch / "task.md")
        plan_path = branch / "plan.md"
        if plan_path.exists():
            plan_path.unlink()
        branch_folders.append(str(branch.resolve()))
    return branch_folders


async def _build_branch(branch_work_folder: str, provider: str) -> str:
    """Run plan builder in a branch work folder and return branch plan path."""
    with context_scope(work_folder=branch_work_folder):
        await builder.run(provider=provider)
        return str(Path(branch_work_folder) / "plan.md")


def _stamp_plan_frontmatter(*, provider: str, branch_count: int,
                             max_revisions: int) -> None:
    """Stamp YAML frontmatter on plan.md with workflow parameters."""
    ctx = get_context()
    update_frontmatter(
        ctx.paths["plan_path"],
        size=ctx.size,
        plan_builder_provider=provider,
        branch_count=branch_count,
        max_revisions=max_revisions,
    )


# --- Orchestration ---

async def run(work_folder: str, size: str = "M",
              depth: int | None = None,
              plan_builder_provider: str = "claude",
              branch_count: int = 1,
              max_revisions: int = MAX_REVISIONS) -> None:
    """Build a plan, run review/revision cycles, and format the final output."""
    if branch_count < 1:
        raise ValueError("branch_count must be >= 1")
    if max_revisions < 0:
        raise ValueError("max_revisions must be >= 0")
    if plan_builder_provider not in _SUPPORTED_PROVIDERS:
        allowed = ", ".join(sorted(_SUPPORTED_PROVIDERS))
        raise ValueError(
            f"Unsupported plan_builder_provider: {plan_builder_provider}. "
            f"Supported providers: {allowed}."
        )

    with context_scope(
        work_folder=str(Path(work_folder).resolve()),
        size=size,
        depth=depth,
    ):
        # Build (single branch) or build branches + merger feedback + builder
        if branch_count == 1:
            session_id = await builder.run(provider=plan_builder_provider)
        else:
            branch_folders = _prepare_branches(get_context().work_folder, branch_count)
            branch_plan_paths = await asyncio.gather(
                *(_build_branch(folder, plan_builder_provider) for folder in branch_folders)
            )
            merger_response = await merger.run(
                branch_plan_paths=branch_plan_paths,
            )
            merger_feedback = parse_result(merger_response).notes or merger_response
            session_id = await builder.run(
                provider=plan_builder_provider,
                merger_feedback=merger_feedback,
            )

        _stamp_plan_frontmatter(
            provider=plan_builder_provider,
            branch_count=branch_count,
            max_revisions=max_revisions,
        )

        # Review loop: each fail triggers a revision. After max cycles, accept final draft.
        for _ in range(max_revisions):
            review_result = await review.run()
            result = parse_result(review_result)

            if result.passed:
                break

            feedback = result.notes or review_result
            session_id = await builder.run_resume(
                feedback=feedback,
                session_id=session_id,
                provider=plan_builder_provider,
            )

            _stamp_plan_frontmatter(
                provider=plan_builder_provider,
                branch_count=branch_count,
                max_revisions=max_revisions,
            )

        await formatter.run()

        _stamp_plan_frontmatter(
            provider=plan_builder_provider,
            branch_count=branch_count,
            max_revisions=max_revisions,
        )


if __name__ == "__main__":
    from mixersystem.data.repository import run_cli
    run_cli(run)
