#!/usr/bin/env python3
"""
Claude Code Parallel Orchestrator v1.0
并行运行多个 CC 子 agent，每个在独立的 git worktree 中工作。

架构：
  Hermes (协调器)
    ├─ CC Worker 1 → worktree-1 (branch: parallel/task-1)
    ├─ CC Worker 2 → worktree-2 (branch: parallel/task-2)
    └─ CC Worker N → worktree-N (branch: parallel/task-N)

用法：
  # 双任务并行
  python3 cc_parallel.py \\
    --task "Implement CharacterModel.ts" \\
    --task "Implement CharacterView.tsx" \\
    --cwd ~/projects/oes-web2 --max-turns 8 --timeout 300

  # 三任务 + 自定义 worktree 前缀
  python3 cc_parallel.py \\
    --task "Task A" --task "Task B" --task "Task C" \\
    --cwd ~/project --prefix feat

  # 合并模式（默认 auto = 全部成功才合并）
  python3 cc_parallel.py \\
    --task "A" --task "B" \\
    --cwd ~/project --merge auto

隔离原理：
  git worktree add ../project-worker-1 -b parallel/task-1
  每个 CC 在自己的 worktree 中工作，互不干扰。
  完成后 git merge 将各分支合回主分支。
"""

import asyncio
import argparse
import json
import os
import sys
import time
import shutil
from pathlib import Path

# 复用 cc_sdk 的核心函数
sys.path.insert(0, str(Path(__file__).parent))
from cc_sdk import run_task, STATUS_DIR, read_status, find_status_file


def create_worktree(repo_path: str, branch: str, worktree_name: str) -> str | None:
    """创建 git worktree，返回路径"""
    wt_path = os.path.join(os.path.dirname(repo_path), worktree_name)
    
    # 清理旧的 worktree
    if os.path.exists(wt_path):
        shutil.rmtree(wt_path, ignore_errors=True)
    
    # 删除旧的 worktree 引用
    os.system(f'cd "{repo_path}" && git worktree prune 2>/dev/null')
    
    # 创建新 worktree
    ret = os.system(
        f'cd "{repo_path}" && git worktree add "{wt_path}" -b "{branch}" 2>&1'
    )
    if ret != 0:
        # 分支可能已存在，尝试 checkout
        ret = os.system(
            f'cd "{repo_path}" && git worktree add "{wt_path}" "{branch}" 2>&1'
        )
        if ret != 0:
            print(f"❌ Failed to create worktree: {wt_path}", file=sys.stderr)
            return None
    
    print(f"📁 Worktree: {wt_path} (branch: {branch})", file=sys.stderr)
    return wt_path


def cleanup_worktree(repo_path: str, worktree_path: str):
    """清理 worktree"""
    if os.path.exists(worktree_path):
        os.system(f'cd "{repo_path}" && git worktree remove "{worktree_path}" --force 2>/dev/null')


def merge_branch(repo_path: str, branch: str, strategy: str = "auto") -> bool:
    """合并分支到主分支"""
    # 先检查是否有冲突
    ret = os.system(
        f'cd "{repo_path}" && git merge "{branch}" --no-edit 2>&1'
    )
    if ret != 0:
        print(f"⚠️ Merge conflict on {branch}, aborting merge", file=sys.stderr)
        os.system(f'cd "{repo_path}" && git merge --abort 2>/dev/null')
        return False
    return True


async def run_worker(
    worker_id: int,
    task: str,
    cwd: str,
    max_turns: int,
    timeout: int,
    tools: str | None,
    effort: str,
    quiet: bool,
) -> dict:
    """运行单个 CC worker"""
    start = time.time()
    print(f"\n🚀 Worker {worker_id}: starting", file=sys.stderr)
    print(f"   Task: {task[:80]}...", file=sys.stderr)
    print(f"   CWD: {cwd}", file=sys.stderr)
    
    result = await run_task(
        prompt=task,
        cwd=cwd,
        max_turns=max_turns,
        allowed_tools=tools.split(",") if tools else None,
        effort=effort,
        timeout=timeout,
        output_json=True,
        quiet=quiet,
    )
    
    elapsed = time.time() - start
    success = result.get("success", False)
    tool_count = result.get("tool_count", 0)
    errors = result.get("errors", [])
    
    status = "✅" if success else "❌"
    print(f"\n{status} Worker {worker_id}: done in {elapsed:.0f}s | tools={tool_count} | errors={len(errors)}",
          file=sys.stderr)
    
    result["worker_id"] = worker_id
    result["elapsed"] = elapsed
    return result


async def run_parallel(
    tasks: list[str],
    cwd: str,
    max_turns: int,
    timeout: int,
    tools: str | None,
    effort: str,
    prefix: str,
    merge_mode: str,
    quiet: bool,
    use_worktrees: bool,
) -> dict:
    """并行运行多个 CC 任务"""
    repo_path = os.path.abspath(cwd)
    total_start = time.time()
    
    worktrees = []
    worktree_paths = {}
    
    if use_worktrees and os.path.exists(os.path.join(repo_path, ".git")):
        # Git worktree 模式：每个 worker 独立 worktree
        print(f"🌳 Creating {len(tasks)} git worktrees...", file=sys.stderr)
        for i, task in enumerate(tasks):
            branch = f"{prefix}/task-{i+1}"
            wt_name = f".{prefix}-worker-{i+1}"
            wt_path = create_worktree(repo_path, branch, wt_name)
            if wt_path:
                worktrees.append((branch, wt_path))
                worktree_paths[i] = wt_path
            else:
                # fallback: 用原始目录
                worktree_paths[i] = repo_path
        
        if not worktrees:
            print("⚠️ No worktrees created, falling back to shared directory", file=sys.stderr)
    else:
        # 非 git 项目或禁用 worktree：所有 worker 共享目录
        print("📂 Shared directory mode (no worktree isolation)", file=sys.stderr)
    
    # 并行启动所有 worker
    worker_tasks = []
    for i, task in enumerate(tasks):
        worker_cwd = worktree_paths.get(i, repo_path)
        worker_tasks.append(run_worker(
            worker_id=i + 1,
            task=task,
            cwd=worker_cwd,
            max_turns=max_turns,
            timeout=timeout,
            tools=tools,
            effort=effort,
            quiet=quiet,
        ))
    
    # 等待所有完成
    results = await asyncio.gather(*worker_tasks, return_exceptions=True)
    
    total_elapsed = time.time() - total_start
    
    # 统计
    successes = 0
    failures = 0
    for r in results:
        if isinstance(r, Exception):
            failures += 1
        elif r.get("success"):
            successes += 1
        else:
            failures += 1
    
    print(f"\n{'='*50}", file=sys.stderr)
    print(f"📊 并行执行完成: {successes}✅ {failures}❌ | 总耗时 {total_elapsed:.0f}s", file=sys.stderr)
    
    # 合并阶段
    merged_branches = []
    merge_errors = []
    
    if use_worktrees and worktrees and merge_mode != "none":
        should_merge = (
            merge_mode == "always" or
            (merge_mode == "auto" and failures == 0)
        )
        
        if should_merge:
            print(f"\n🔀 Merging {len(worktrees)} branches...", file=sys.stderr)
            os.system(f'cd "{repo_path}" && git checkout main 2>/dev/null || git checkout master 2>/dev/null')
            
            for branch, wt_path in worktrees:
                # 先提交 worktree 中的改动
                os.system(f'cd "{wt_path}" && git add -A && git diff --cached --quiet || git commit -m "CC worker: {branch}" 2>/dev/null')
                
                if merge_branch(repo_path, branch):
                    merged_branches.append(branch)
                    print(f"  ✅ Merged {branch}", file=sys.stderr)
                else:
                    merge_errors.append(branch)
                    print(f"  ❌ Conflict: {branch}", file=sys.stderr)
        else:
            print(f"\n⏭️ Skipping merge (mode={merge_mode}, failures={failures})", file=sys.stderr)
    
    # 清理 worktrees
    for branch, wt_path in worktrees:
        cleanup_worktree(repo_path, wt_path)
    
    # 删除已合并的分支
    for branch in merged_branches:
        os.system(f'cd "{repo_path}" && git branch -d "{branch}" 2>/dev/null')
    
    # 构建最终结果
    worker_results = []
    for i, r in enumerate(results):
        if isinstance(r, Exception):
            worker_results.append({
                "worker_id": i + 1,
                "success": False,
                "error": str(r),
                "text": "",
                "tool_count": 0,
            })
        else:
            worker_results.append({
                "worker_id": r.get("worker_id", i + 1),
                "success": r.get("success", False),
                "text": r.get("text", "")[:500],
                "tool_count": r.get("tool_count", 0),
                "tool_uses": r.get("tool_uses", []),
                "elapsed": r.get("elapsed", 0),
                "session_id": r.get("session_id"),
                "errors": r.get("errors", []),
            })
    
    return {
        "success": failures == 0,
        "total_elapsed": round(total_elapsed, 1),
        "workers": len(tasks),
        "successes": successes,
        "failures": failures,
        "merged_branches": merged_branches,
        "merge_errors": merge_errors,
        "results": worker_results,
    }


def main():
    parser = argparse.ArgumentParser(
        description="Claude Code Parallel Orchestrator — run multiple CC workers concurrently",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
示例:
  # 双任务并行（自动 worktree 隔离）
  python3 cc_parallel.py \\
    --task "Create src/models/Character.ts" \\
    --task "Create src/views/CharacterView.tsx" \\
    --cwd ~/project --max-turns 8

  # 三任务 + 不合并
  python3 cc_parallel.py \\
    --task "A" --task "B" --task "C" \\
    --cwd ~/project --merge none

  # 查询进度
  python3 cc_parallel.py --progress
"""
    )
    parser.add_argument("--task", action="append", required=True, help="Task for a CC worker (repeatable)")
    parser.add_argument("--cwd", default=".", help="Project root directory")
    parser.add_argument("--max-turns", type=int, default=8, help="Max turns per worker (default: 8)")
    parser.add_argument("--timeout", type=int, default=300, help="Timeout per worker in seconds (default: 300)")
    parser.add_argument("--tools", default=None, help="Allowed tools (comma-separated)")
    parser.add_argument("--effort", default="medium", choices=["low", "medium", "high", "max"])
    parser.add_argument("--prefix", default="parallel", help="Branch/worktree prefix (default: parallel)")
    parser.add_argument("--merge", default="auto", choices=["auto", "always", "none"],
                        help="Merge strategy: auto=only if all succeed, always=always, none=skip")
    parser.add_argument("--no-worktree", action="store_true", help="Disable git worktree isolation")
    parser.add_argument("--json", action="store_true", help="JSON output")
    parser.add_argument("--quiet", action="store_true", help="Suppress worker stderr output")
    
    args = parser.parse_args()
    
    if len(args.task) > 5:
        print("⚠️ Max 5 parallel workers supported", file=sys.stderr)
        sys.exit(1)
    
    result = asyncio.run(run_parallel(
        tasks=args.task,
        cwd=args.cwd,
        max_turns=args.max_turns,
        timeout=args.timeout,
        tools=args.tools,
        effort=args.effort,
        prefix=args.prefix,
        merge_mode=args.merge,
        quiet=args.quiet,
        use_worktrees=not args.no_worktree,
    ))
    
    if args.json:
        print(json.dumps(result, ensure_ascii=False, indent=2))
    else:
        # 人类友好输出
        print(f"\n{'='*60}")
        print(f"📊 结果: {result['successes']}✅ {result['failures']}❌ | {result['total_elapsed']:.0f}s")
        if result["merged_branches"]:
            print(f"🔀 已合并: {', '.join(result['merged_branches'])}")
        if result["merge_errors"]:
            print(f"⚠️ 合并失败: {', '.join(result['merge_errors'])}")
        print()
        for w in result["results"]:
            status = "✅" if w["success"] else "❌"
            text = w["text"][:100].replace("\n", " ") if w["text"] else "(no output)"
            print(f"  Worker {w['worker_id']}: {status} | tools={w['tool_count']} | {text}")
    
    sys.exit(0 if result["success"] else 1)


if __name__ == "__main__":
    main()