import asyncio
import random
import shutil
from pathlib import Path

from mixersystem.data.repository import context_scope, get_context, parse_result, update_session_run
from mixersystem.workflows.shared import artifact_builder as builder
from mixersystem.workflows.plan.agents import plan_merger as merger
from mixersystem.workflows.plan.agents import plan_reviewer as reviewer
from mixersystem.workflows.shared.implementation_checker import run as check_implementation

MAX_REVISIONS = 2
_CONCRETE_PROVIDERS = ["claude", "gemini", "codex"]
_SUPPORTED_PROVIDERS = {*_CONCRETE_PROVIDERS, "random"}


def _resolve_provider(provider: str) -> str:
    """Resolve provider name, picking randomly if 'random'."""
    if provider == "random":
        return random.choice(_CONCRETE_PROVIDERS)
    return provider


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


def _prepare_branches(session_folder: str, branch_count: int) -> list[str]:
    """Create branch session folders and copy root task.md into each."""
    root = Path(session_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_session_folder: str, provider: str,
                        instructions: str | None = None,
                        use_questions: bool = True) -> str:
    """Run plan artifact builder in a branch session folder and return branch plan path."""
    with context_scope(session_folder=branch_session_folder):
        await builder.run(stage="plan", provider=provider, instructions=instructions,
                          use_questions=use_questions)
        return str(Path(branch_session_folder) / "plan.md")


# --- Orchestration ---

async def run(session_folder: str, lite: bool = False,
              depth: int | None = None,
              plan_builder_provider: str = "claude",
              branch_count: int = 1,
              max_revisions: int = MAX_REVISIONS,
              instructions: str | None = None,
              use_questions: bool = True) -> 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}."
        )

    resolved_provider = _resolve_provider(plan_builder_provider)

    wf = Path(session_folder)
    wf.mkdir(parents=True, exist_ok=True)

    with context_scope(
        session_folder=str(wf.resolve()),
        lite=lite,
        depth=depth,
    ):
        update_session_run(
            str(wf.resolve()), "plan",
            provider=resolved_provider,
            branch_count=branch_count,
            max_revisions=max_revisions,
            lite=lite,
        )

        # Build (single branch) or build branches + merger feedback + artifact builder
        if branch_count == 1:
            session_id = await builder.run(stage="plan", provider=resolved_provider,
                                           instructions=instructions, use_questions=use_questions)
        else:
            branch_folders = _prepare_branches(get_context().session_folder, branch_count)
            if plan_builder_provider == "random":
                branch_providers = [_resolve_provider("random") for _ in branch_folders]
            else:
                branch_providers = [resolved_provider] * len(branch_folders)
            branch_plan_paths = await asyncio.gather(
                *(_build_branch(folder, bp, instructions=instructions,
                                use_questions=use_questions)
                  for folder, bp in zip(branch_folders, branch_providers))
            )
            merger_response = await merger.run(branch_plan_paths=branch_plan_paths)
            merger_feedback = parse_result(merger_response).notes or merger_response
            merged_instructions = f"{instructions}\n\nMerger feedback:\n{merger_feedback}" if instructions else f"Merger feedback:\n{merger_feedback}"
            session_id = await builder.run(
                stage="plan",
                provider=resolved_provider,
                instructions=merged_instructions,
                use_questions=use_questions,
            )

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

            if result.passed:
                break

            feedback = result.notes or review_result
            session_id = await builder.run_resume(
                stage="plan",
                feedback=feedback,
                session_id=session_id,
                provider=resolved_provider,
            )

        await check_implementation("plan")


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