#!/usr/bin/env python3
"""开发分支工作区初始化脚本

用于创建基于远端分支或新建的开发分支工作区。
"""

import argparse
import subprocess
import sys
from pathlib import Path


def run_git(args: list[str]) -> subprocess.CompletedProcess:
    """执行 git 命令
    
    Args:
        args: git 命令参数列表
        
    Returns:
        subprocess.CompletedProcess 对象
    """
    return subprocess.run(
        ["git"] + args,
        capture_output=True,
        text=True,
        check=False
    )


def get_root_dir() -> Path:
    """获取 git 仓库根目录
    
    Raises:
        RuntimeError: 不在 git 仓库中
    """
    result = run_git(["rev-parse", "--show-toplevel"])
    if result.returncode != 0:
        raise RuntimeError("不在 git 仓库中")
    return Path(result.stdout.strip())


def fetch_remote() -> bool:
    """从远端获取最新分支信息
    
    Returns:
        是否成功获取
    """
    result = run_git(["fetch", "--quiet"])
    if result.returncode != 0:
        print("警告: 无法获取远端信息，继续使用本地数据")
        return False
    return True


def get_all_remote_branches() -> list[str]:
    """获取所有远端分支列表
    
    Returns:
        远端分支列表
    """
    result = run_git(["branch", "-r"])
    if result.returncode != 0:
        return []
    
    return [line.strip() for line in result.stdout.strip().split("\n") if line.strip()]


def get_remote_branches(branch_name: str) -> list[str]:
    """获取远端匹配的分支列表
    
    Args:
        branch_name: 分支名称
        
    Returns:
        匹配的远端分支列表
    """
    result = run_git(["branch", "-r"])
    if result.returncode != 0:
        return []
    
    return [
        line.strip()
        for line in result.stdout.strip().split("\n")
        if line.strip() and line.strip().endswith(f"/{branch_name}")
    ]


def branch_exists_local(branch_name: str) -> bool:
    """检查本地分支是否存在
    
    Args:
        branch_name: 分支名称
        
    Returns:
        是否存在
    """
    result = run_git(["show-ref", "--verify", "--quiet", f"refs/heads/{branch_name}"])
    return result.returncode == 0


def worktree_exists(worktree_path: Path) -> bool:
    """检查工作区是否存在
    
    Args:
        worktree_path: 工作区路径
        
    Returns:
        是否存在
    """
    result = run_git(["worktree", "list"])
    if result.returncode != 0:
        return False
    return str(worktree_path) in result.stdout


def extract_branch_name(remote_branch: str) -> str:
    """从远端分支名提取分支名
    
    Args:
        remote_branch: 远端分支名（如 origin/feature-123）
        
    Returns:
        分支名（如 feature-123）
    """
    return remote_branch.split("/", 1)[1] if "/" in remote_branch else remote_branch


def validate_branch_creation(branch_name: str, worktree_path: Path) -> None:
    """验证是否可以创建分支和工作区
    
    Args:
        branch_name: 分支名称
        worktree_path: 工作区路径
        
    Raises:
        RuntimeError: 验证失败
    """
    if worktree_path.exists():
        raise RuntimeError(f"工作区已存在于 {worktree_path}")
    
    if worktree_exists(worktree_path):
        raise RuntimeError(f"工作区 '{branch_name}' 已存在")
    
    if branch_exists_local(branch_name):
        raise RuntimeError(f"本地分支 '{branch_name}' 已存在")


def select_from_list(items: list[str], prompt: str, allow_create: bool = False) -> str | None:
    """让用户从列表中选择
    
    Args:
        items: 选项列表
        prompt: 提示信息
        allow_create: 是否允许创建新分支
        
    Returns:
        用户选择的项，或 None 表示创建新分支
    """
    if not items:
        return None
    
    print(prompt)
    for i, item in enumerate(items, 1):
        print(f"     {i}\t{item}")
    
    if allow_create:
        print(f"     {len(items) + 1}\t创建新分支")
    
    print()
    
    max_choice = len(items) + 1 if allow_create else len(items)
    
    while True:
        try:
            selection = int(input(f"请选择 (1-{max_choice}): "))
            if 1 <= selection <= len(items):
                selected = items[selection - 1]
                print(f"已选择: {selected}")
                return selected
            if allow_create and selection == len(items) + 1:
                return None
            print(f"错误: 请输入 1-{max_choice} 之间的数字")
        except ValueError:
            print("错误: 请输入有效的数字")
        except KeyboardInterrupt:
            print("\n操作已取消")
            sys.exit(1)


def input_branch_name() -> str:
    """让用户输入分支名称
    
    Returns:
        分支名称
    """
    while True:
        try:
            name = input("请输入新分支名称: ").strip()
            if name:
                return name
            print("错误: 分支名称不能为空")
        except KeyboardInterrupt:
            print("\n操作已取消")
            sys.exit(1)


def create_worktree(branch_name: str, worktree_path: Path, base_branch: str | None = None) -> None:
    """创建工作区
    
    Args:
        branch_name: 新分支名称
        worktree_path: 工作区路径
        base_branch: 基础分支（可选）
        
    Raises:
        RuntimeError: 创建失败
    """
    args = ["worktree", "add", "-b", branch_name, str(worktree_path)]
    if base_branch:
        args.append(base_branch)
    
    result = run_git(args)
    if result.returncode != 0:
        raise RuntimeError(f"创建工作区失败: {result.stderr.strip()}")


def create_worktree_with_validation(
    branch_name: str,
    worktree_path: Path,
    base_branch: str | None = None
) -> None:
    """验证并创建工作区
    
    Args:
        branch_name: 新分支名称
        worktree_path: 工作区路径
        base_branch: 基础分支（可选）
    """
    validate_branch_creation(branch_name, worktree_path)
    create_worktree(branch_name, worktree_path, base_branch)
    print(f"工作区已创建于 {worktree_path}")


def handle_with_branch_name(branch_name: str, worktrees_dir: Path) -> None:
    """处理有分支名参数的情况
    
    Args:
        branch_name: 分支名称
        worktrees_dir: 工作区目录
    """
    worktree_path = worktrees_dir / branch_name
    
    try:
        validate_branch_creation(branch_name, worktree_path)
    except RuntimeError as e:
        print(f"错误: {e}")
        sys.exit(1)
    
    remote_branches = get_remote_branches(branch_name)
    
    try:
        if remote_branches:
            selected_branch = select_from_list(
                remote_branches,
                "找到远端分支:",
                allow_create=True
            )
            if selected_branch:
                create_worktree(branch_name, worktree_path, selected_branch)
            else:
                create_worktree(branch_name, worktree_path)
        else:
            print("未找到远端分支，创建新分支")
            create_worktree(branch_name, worktree_path)
        
        print(f"工作区已创建于 {worktree_path}")
    except RuntimeError as e:
        print(f"错误: {e}")
        sys.exit(1)


def handle_without_branch_name(worktrees_dir: Path) -> None:
    """处理无分支名参数的情况
    
    Args:
        worktrees_dir: 工作区目录
    """
    remote_branches = get_all_remote_branches()
    
    if not remote_branches:
        print("未找到远端分支")
        sys.exit(1)
    
    selected_branch = select_from_list(
        remote_branches,
        "远端分支列表:",
        allow_create=True
    )
    
    if selected_branch:
        branch_name = extract_branch_name(selected_branch)
        worktree_path = worktrees_dir / branch_name
        
        try:
            create_worktree_with_validation(branch_name, worktree_path, selected_branch)
        except RuntimeError as e:
            print(f"错误: {e}")
            sys.exit(1)
    else:
        branch_name = input_branch_name()
        worktree_path = worktrees_dir / branch_name
        
        try:
            create_worktree_with_validation(branch_name, worktree_path)
        except RuntimeError as e:
            print(f"错误: {e}")
            sys.exit(1)


def main() -> None:
    parser = argparse.ArgumentParser(
        description="创建开发分支工作区",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
示例:
  %(prog)s                  # 列出远端分支供选择
  %(prog)s feature-123      # 创建 feature-123 工作区
  %(prog)s bugfix-456       # 创建 bugfix-456 工作区
        """
    )
    parser.add_argument("branch_name", nargs="?", help="分支名称（可选）")
    
    args = parser.parse_args()
    
    try:
        root_dir = get_root_dir()
    except RuntimeError as e:
        print(f"错误: {e}")
        sys.exit(1)
    
    worktrees_dir = root_dir / ".worktrees"
    worktrees_dir.mkdir(parents=True, exist_ok=True)
    
    print("正在从远端获取最新分支信息...")
    fetch_remote()
    
    if args.branch_name:
        handle_with_branch_name(args.branch_name, worktrees_dir)
    else:
        handle_without_branch_name(worktrees_dir)


if __name__ == "__main__":
    main()