#!/usr/bin/env python3
"""Store files to a remote git branch without affecting your current working state."""

import argparse
import os
import subprocess
import sys
from pathlib import Path


def run(cmd: list[str], *, check: bool = True, capture: bool = False, env: dict | None = None) -> subprocess.CompletedProcess:
    merged_env = {**os.environ, **(env or {})}
    return subprocess.run(
        cmd,
        check=check,
        capture_output=capture,
        text=True,
        env=merged_env,
    )


def run_output(cmd: list[str], *, check: bool = True, env: dict | None = None) -> str:
    result = run(cmd, check=check, capture=True, env=env)
    return result.stdout.strip() if result.stdout else ""


def temp_index_path(name: str) -> str:
    return run_output(["git", "rev-parse", "--path-format=absolute", "--git-path", name])


def branch_exists(remote: str, branch: str) -> bool:
    result = run(["git", "rev-parse", "--verify", f"{remote}/{branch}"], check=False, capture=True)
    return result.returncode == 0


def get_branch_files(remote: str, branch: str) -> list[str]:
    output = run_output(["git", "ls-tree", "-r", "--name-only", f"{remote}/{branch}"], check=False)
    return [f for f in output.splitlines() if f]


def cmd_status(args: argparse.Namespace) -> int:
    """Show git status for files in the remote branch."""
    remote, branch = args.remote, args.branch

    if not branch_exists(remote, branch):
        print(f"Branch {remote}/{branch} does not exist.", file=sys.stderr)
        return 1

    branch_files = get_branch_files(remote, branch)
    if not branch_files:
        print(f"No files in branch {remote}/{branch}")
        return 0

    modified_files = []
    deleted_files = []

    for file in branch_files:
        if not Path(file).exists():
            deleted_files.append(file)
            continue

        result = run(
            ["git", "--no-pager", "diff", "--quiet", f"{remote}/{branch}:./{file}", file],
            check=False,
            capture=True,
        )
        if result.returncode != 0:
            modified_files.append(file)

    has_output = False
    for file in modified_files:
        print(f"  modified:   {file}")
        has_output = True
    for file in deleted_files:
        print(f"  deleted:    {file}")
        has_output = True

    if not has_output:
        print(f"Files in {remote}/{branch} are up to date.")

    if args.diff and modified_files:
        for file in modified_files:
            print(f"--- {remote}/{branch}:{file}")
            print(f"+++ {file}")
            run(["git", "--no-pager", "diff", f"{remote}/{branch}:./{file}", file], check=False)

    return 0


def cmd_restore(args: argparse.Namespace) -> int:
    """Restore files from remote branch to working directory."""
    remote, branch = args.remote, args.branch

    if not branch_exists(remote, branch):
        print(f"Branch {remote}/{branch} does not exist.", file=sys.stderr)
        return 1

    temp_index = temp_index_path("temp_restore_index")
    env = {"GIT_INDEX_FILE": temp_index}

    try:
        # Load the branch's tree into the temporary index
        result = run(["git", "read-tree", f"{remote}/{branch}"], check=False, capture=True, env=env)
        if result.returncode != 0:
            run(["git", "read-tree", "--empty"], env=env)

        if args.files:
            run(["git", "restore", "--worktree", "--", *args.files], env=env)
        else:
            run(["git", "restore", "--worktree", "."], env=env)
    finally:
        if Path(temp_index).exists():
            Path(temp_index).unlink()

    return 0


def cmd_remove(args: argparse.Namespace) -> int:
    """Remove files from the remote branch (stops tracking them there) without deleting from workspace."""
    remote, branch = args.remote, args.branch
    files = args.files

    if not files:
        print("Error: No files specified to remove", file=sys.stderr)
        return 1

    if not branch_exists(remote, branch):
        print(f"Branch {remote}/{branch} does not exist.", file=sys.stderr)
        return 1

    branch_files = set(get_branch_files(remote, branch))
    missing = [f for f in files if f not in branch_files]
    if missing:
        for f in missing:
            print(f"Error: '{f}' is not in branch {remote}/{branch}", file=sys.stderr)
        return 1

    temp_index = temp_index_path("temp_remove_index")
    env = {"GIT_INDEX_FILE": temp_index}

    try:
        # Load the target branch's current state
        run(["git", "read-tree", f"{remote}/{branch}"], env=env)

        # Remove from index only (files stay in workspace); -f avoids staged-content check
        run(["git", "rm", "--cached", "-f", "--ignore-unmatch", "--", *files], env=env)

        tree = run_output(["git", "write-tree"], env=env)

        parent = run_output(["git", "rev-parse", "--verify", f"{remote}/{branch}"], env=env)

        commit_msg = f"Remove from store: {' '.join(files)}"

        commit_hash = run_output(
            ["git", "commit-tree", tree, "-p", parent, "-m", commit_msg],
            env=env,
        )

        print("Commit stats:")
        result = run(
            ["git", "--no-pager", "diff-tree", "--stat", parent, commit_hash],
            check=False,
            capture=True,
        )
        if result.returncode == 0 and result.stdout:
            print(result.stdout)
        print()

        try:
            response = input(f"Push to {remote}/{branch}? [y/N] ")
        except EOFError:
            response = ""

        if response.lower() != "y":
            print("Push cancelled.")
            return 1

        run(["git", "push", remote, f"{commit_hash}:refs/heads/{branch}"])
        print(f"Removed from {remote}/{branch}: {' '.join(files)}")

    finally:
        if Path(temp_index).exists():
            Path(temp_index).unlink()

    return 0


def cmd_ls(args: argparse.Namespace) -> int:
    remote, branch = args.remote, args.branch

    if not branch_exists(remote, branch):
        print(f"Branch {remote}/{branch} does not exist.", file=sys.stderr)
        return 1

    branch_files = get_branch_files(remote, branch)
    if not branch_files:
        print(f"No files in branch {remote}/{branch}")
        return 0

    for file in branch_files:
        print(file)

    return 0


def cmd_store(args: argparse.Namespace) -> int:
    remote, branch = args.remote, args.branch
    files = args.files

    if not files:
        print("Error: No files specified to store", file=sys.stderr)
        return 1

    temp_index = temp_index_path("temp_backup_index")
    env = {"GIT_INDEX_FILE": temp_index}

    try:
        # Load the target branch's current state
        result = run(["git", "read-tree", f"{remote}/{branch}"], check=False, capture=True, env=env)
        if result.returncode != 0:
            run(["git", "read-tree", "--empty"], env=env)

        # Forced so ignored files can be stored too
        run(["git", "add", "-f", "--", *files], env=env)

        tree = run_output(["git", "write-tree"], env=env)

        parent = run_output(["git", "rev-parse", "--verify", f"{remote}/{branch}"], check=False)

        commit_msg = f"Add to store: {' '.join(files)}"

        if parent:
            commit_hash = run_output(
                ["git", "commit-tree", tree, "-p", parent, "-m", commit_msg],
                env=env,
            )
        else:
            commit_hash = run_output(
                ["git", "commit-tree", tree, "-m", commit_msg],
                env=env,
            )

        print("Commit stats:")
        if parent:
            result = run(
                ["git", "--no-pager", "show", "--format=", "--stat", commit_hash],
                check=False,
                capture=True,
            )
            if result.returncode == 0 and result.stdout:
                print(result.stdout)
            else:
                run(["git", "--no-pager", "diff-tree", "--stat", parent, commit_hash], check=False)
        else:
            run(["git", "--no-pager", "show", "--format=", "--stat", commit_hash], check=False)
        print()

        try:
            response = input(f"Push to {remote}/{branch}? [y/N] ")
        except EOFError:
            response = ""

        if response.lower() != "y":
            print("Push cancelled.")
            return 1

        run(["git", "push", remote, f"{commit_hash}:refs/heads/{branch}"])

        print(f"Added to {remote}/{branch}: {' '.join(files)}")

    finally:
        if Path(temp_index).exists():
            Path(temp_index).unlink()

    return 0


COMMANDS = ("store", "status", "restore", "remove", "ls")


def with_default_command(argv: list[str]) -> list[str]:
    """Insert `store` before the first positional argument unless it names a subcommand."""
    i = 0
    while i < len(argv):
        arg = argv[i]
        if arg in ("-b", "--branch", "-r", "--remote"):
            i += 2
        elif arg.startswith("-") and arg != "--":
            i += 1
        elif arg in COMMANDS:
            return argv
        else:
            return argv[:i] + ["store"] + argv[i:]
    return argv + ["store"]


def main() -> int:
    parser = argparse.ArgumentParser(
        description="Store files to a remote git branch without affecting your current working state.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  %(prog)s config.txt                    # store file (default command)
  %(prog)s file1.txt file2.txt           # store multiple files
  %(prog)s --branch backup config.txt    # store with custom branch
  %(prog)s status                        # check status
  %(prog)s status -d                     # check status with diff
  %(prog)s restore config.txt            # restore specific file
  %(prog)s restore                       # restore all files from branch
  %(prog)s remove config.txt             # remove file from branch (keep in workspace)
  %(prog)s ls                            # list all files in branch
""",
    )

    default_remote = "origin"
    result = run(["git", "remote"], check=False, capture=True)
    if result.returncode == 0 and "_store" in result.stdout.splitlines():
        default_remote = "_store"

    parser.add_argument("-b", "--branch", default="_store", help="Branch name (default: _store)")
    parser.add_argument("-r", "--remote", default=default_remote, help=f"Remote name (default: {default_remote})")

    subparsers = parser.add_subparsers(dest="command", help="Command to run")

    store_parser = subparsers.add_parser("store", help="Store files to remote branch (default)")
    store_parser.add_argument("files", nargs="*", help="Files to store")

    status_parser = subparsers.add_parser("status", help="Show git status for files in the remote branch")
    status_parser.add_argument("-d", "--diff", action="store_true", help="Show diff for modified files")

    restore_parser = subparsers.add_parser("restore", help="Restore files from remote branch to working directory")
    restore_parser.add_argument("files", nargs="*", help="Files to restore (all if none specified)")

    remove_parser = subparsers.add_parser("remove", help="Remove files from remote branch (keep in workspace)")
    remove_parser.add_argument("files", nargs="+", help="Files to remove from the branch")

    subparsers.add_parser("ls", help="List all files in the remote branch")

    args = parser.parse_args(with_default_command(sys.argv[1:]))

    if args.command == "store":
        return cmd_store(args)
    elif args.command == "status":
        return cmd_status(args)
    elif args.command == "restore":
        return cmd_restore(args)
    elif args.command == "remove":
        return cmd_remove(args)
    elif args.command == "ls":
        return cmd_ls(args)
    else:
        parser.print_help()
        return 1


if __name__ == "__main__":
    sys.exit(main())