commit_msg.py

#!/usr/bin/env python3
import subprocess
import json
import time
import pathlib
import os
import sys
import argparse
import re
import module
import yaml
import logging
from typing import Tuple, Dict, Any, Optional, List

import utils
utils.import_or_install("tiktoken")
utils.import_or_install("rich")

from tiktoken._educational import SimpleBytePairEncoding
from rich.console import Console
from rich.progress import Progress
from rich.prompt import Confirm
from rich.logging import RichHandler


TOKENIZER = SimpleBytePairEncoding.from_tiktoken("cl100k_base")  # gpt4 encoder

# Rich console for better output
console = Console()

# Setup logging
logging.basicConfig(
    level=logging.ERROR,
    format="%(asctime)s - %(levelname)s - %(message)s",
    handlers=[RichHandler(console=console)],
)
logger = logging.getLogger(__name__)
logger.setLevel(logging.ERROR)

assert logger.level == logging.ERROR
module.configure_logging(logger.level, _console=console)


search_shortcode = re.compile(r":([a-z_]+):")
SHORTCODE_TO_EMOJI = {}
emoji_path = pathlib.Path(__file__).parent / "data/emojis.json"
if emoji_path.exists():
    with open(emoji_path) as f:
        try:
            SHORTCODE_TO_EMOJI = json.load(f)
        except json.JSONDecodeError:
            logger.error(f"Failed to parse data/emojis.json. Using empty emoji dictionary.")

parser = argparse.ArgumentParser(
    description="A tool to generate and manage git commit messages."
)

# Add the optional 'prompt' argument
parser.add_argument(
    "prompt",
    nargs="*",
    help="An optional prompt to guide the generation of the commit message.",
)
parser.add_argument(
    "-H",
    "--history-commit",
    action="store_true",
    help="Include recent repository commits in the prompt.",
    default=False,
)
parser.add_argument(
    "-U",
    "--no-auto-update",
    action="store_false",
    help="Disable auto-update.",
    default=True,
)
parser.add_argument(
    "-s",
    "--shortcodes",
    action="store_true",
    help="Show shortcodes instead of emojis.",
    default=False,
)
parser.add_argument(
    "--no-staged-changes",
    action="store_false",
    help="Do not include staged changes in the prompt.",
    default=True,
)
parser.add_argument(
    "-y",
    "--yes",
    action="store_true",
    help="Use the generated commit message without confirmation.",
    default=False,
)
parser.add_argument(
    "-m", "--model", help="Specify which AI model to use.", default=None
)
parser.add_argument(
    "-v", "--verbose", action="store_true", help="Enable verbose output", default=False
)
parser.add_argument("-c", "--config", help="Path to custom config file", default=None)

# Parse the arguments
parsed_args = parser.parse_args()

# Set log level based on verbose flag
if parsed_args.verbose:
    logger.setLevel(logging.DEBUG)
    module.configure_logging(logging.DEBUG, _console=console)

# check if the .env.json file exists
AUTO_UPDATE = parsed_args.no_auto_update

model_name = parsed_args.model or os.environ.get("COMMIT_MODEL") or "gpt-4o" or "o3-mini"
COPILOT = module.Copilot(model_name)


class bcolors:
    HEADER = "\033[95m"
    OKBLUE = "\033[94m"
    OKCYAN = "\033[96m"
    OKGREEN = "\033[92m"
    WARNING = "\033[93m"
    FAIL = "\033[91m"
    ENDC = "\033[0m"
    BOLD = "\033[1m"
    UNDERLINE = "\033[4m"


def load_config(config_path: Optional[str] = None) -> Dict[str, Any]:
    """Load commit convention from YAML with fallback paths"""
    paths_to_try = [
        pathlib.Path(config_path) if config_path else None,
        pathlib.Path(__file__).parent / "commit_convention.yml",
        pathlib.Path.home() / ".config" / "git_commit" / "commit_convention.yml",
    ]

    for path in paths_to_try:
        if not path:
            continue
        if path.exists():
            logger.debug(f"Loading config from {path}")
            try:
                with open(path, "r") as f:
                    return yaml.safe_load(f) or {}
            except Exception as e:
                logger.error(f"Error loading config from {path}: {e}")

    logger.warning("No config file found, using defaults")
    return {}


# Load the configuration
COMMIT_CONVENTION = load_config(parsed_args.config)


def update() -> None:
    """Update the script from git repository if it's been more than an hour since last update"""
    my_path = os.path.abspath(os.path.dirname(__file__))
    if not pathlib.Path(my_path + "/.git").exists():
        logger.debug("Not a git repository, skipping update")
        return

    try:
        last_up = subprocess.check_output(
            f'stat -c %Y "{my_path}/.git/FETCH_HEAD"', shell=True, text=True
        ).strip()
        if last_up:
            if time.time() - int(last_up) < 3600:
                return
    except subprocess.SubprocessError:
        logger.debug("Couldn't determine last update time, continuing with update")

    # Check if git command exists
    if (
        subprocess.run(
            "git --version",
            shell=True,
            text=True,
            stderr=subprocess.DEVNULL,
            stdout=subprocess.DEVNULL,
        ).returncode
        != 0
    ):
        logger.warning("Git is not installed, skipping the update")
        return

    try:
        # Run git pull
        result = subprocess.check_output(
            f"cd {my_path} && git pull", shell=True, text=True
        ).strip()
        if result != "Already up to date.":
            console.print(f"[green]{result}[/green]")
            console.print("Restarting the script")
            os.execl(sys.executable, sys.executable, *sys.argv)
            sys.exit(0)
    except subprocess.SubprocessError as e:
        logger.error(f"Failed to update: {e}")


if AUTO_UPDATE:
    update()


def get_context() -> Tuple[str, str]:
    """Get git diff context and recent commits"""
    try:
        git_diff = subprocess.check_output(
            """git diff --staged -M -D -C --minimal --name-status | awk '{if ($1 == "A") print "ADDED\t"$2; else if ($1 == "M") print "MODIFIED\t"$2; else if ($1 == "D") print "DELETED\t"$2;}'""",
            shell=True,
            text=True,
        )
        if not git_diff.strip():
            console.print("[red]No changes to commit.[/red]")
            sys.exit(1)
        git_diff = f"EDITS STATUS\n```diff\n{git_diff}\n```"

        detail_git_diff = subprocess.check_output(
            "git diff --staged -M -D -C --minimal --diff-filter=MTUXB",
            shell=True,
            text=True,
        )
        if detail_git_diff.strip():
            git_diff += f"\n\nEDITED FILES:\n```diff\n{detail_git_diff}\n```"

        # Limit git_diff to 30,000 tokens
        git_diff = TOKENIZER.decode(TOKENIZER.encode(git_diff, visualise=None)[:30_000])

        try:
            last_commits = subprocess.check_output(
                'git log --format="===================================\n%B" -n 5',
                shell=True,
                text=True,
            ).split("\n\n")
            last_commits_texts = []
            for commit_msg in last_commits:
                if not commit_msg.strip():
                    continue
                last_commits_texts.append(f"```text\n{commit_msg}\n```")
            last_commits_text = "\n\n".join(last_commits_texts)
            return git_diff, last_commits_text
        except subprocess.CalledProcessError:
            return git_diff, ""
    except subprocess.CalledProcessError as e:
        console.print(f"[red]Error getting git context: {e}[/red]")
        sys.exit(1)


def generate_commit_message(prompt: Optional[str], args: argparse.Namespace) -> str:
    """Generate a commit message using the Copilot API"""
    git_diff, last_commits_text = get_context()

    # Get system prompt from config and format with parameters
    system_prompt = COMMIT_CONVENTION.get(
        "prompt", "Generate a concise git commit message for the following changes."
    )
    parameters = COMMIT_CONVENTION.get("parameters", {})

    for param_name, param_value in parameters.items():
        system_prompt = system_prompt.replace(f"{{{param_name}}}", str(param_value))

    messages = [
        {
            "role": "system",
            "content": system_prompt,
        }
    ]

    if args.no_staged_changes:
        messages.append(
            {
                "role": "user",
                "content": f"# CODE CHANGES:\n```\n{git_diff}\n```",
                "name": "changes",
            }
        )

    if args.history_commit and last_commits_text:
        messages.append(
            {
                "role": "user",
                "content": f"# RECENT REPOSITORY COMMITS:\n{last_commits_text}",
                "name": "recent-commits",
            }
        )

    if prompt:
        prompt_text = (
            "Some explanation from the dev about the modification for more context:\n"
            + prompt.strip()
        )
    else:
        prompt_text = ""

    messages.append(
        {
            "role": "user",
            "content": f"""
Generate the commit message.
{prompt_text.strip()}
""".strip(),
            "name": "prompt",
        }
    )

    with Progress(console=console) as progress:
        task = progress.add_task("[cyan]Generating commit message...", total=1)
        try:
            response = COPILOT.question(messages)
            progress.update(task, completed=1)
            return response
        except Exception as e:
            progress.update(task, completed=1)
            logger.error(f"Error generating commit message: {e}")
            raise e


def message_formatter(message: str) -> str:
    """Format the commit message by removing code blocks and fixing emoji spacing"""
    message = message.replace("```text", "")
    message = message.replace("```", "")
    # need 1 space between the shortcode and the (..) part
    message = message.replace(":(", ": (")
    message = message.replace(") :", "):")
    message = message.replace("::", ": :")
    while "  " in message:  # remove double spaces
        message = message.replace("  ", " ")
    return message.strip()


def to_emoji(input: str) -> str:
    """Convert shortcodes to unicode emojis"""
    return search_shortcode.sub(
        lambda x: SHORTCODE_TO_EMOJI.get(x.group(1), f":{x.group(1)}:"), input
    )


def check_git_repo() -> bool:
    """Check if the current directory is a git repository"""
    return (
        subprocess.run(
            "git rev-parse --is-inside-work-tree",
            shell=True,
            text=True,
            stderr=subprocess.DEVNULL,
            stdout=subprocess.DEVNULL,
        ).returncode
        == 0
    )


def check_unpushed_commits() -> None:
    """Check if there are too many unpushed commits and suggest pushing"""
    try:
        unpushed = subprocess.check_output(
            "git log --branches --not --remotes --oneline", shell=True, text=True
        ).split("\n")
        if len(unpushed) > 4:
            console.print(
                "[yellow]You have more than 3 commits not pushed, pushing now...[/yellow]"
            )
            try:
                subprocess.run("git push", shell=True, check=True)
                console.print("[green]Push successful![/green]")
            except subprocess.CalledProcessError:
                console.print("[red]Failed to push commits.[/red]")
    except subprocess.SubprocessError:
        logger.debug("Failed to check unpushed commits")


if __name__ == "__main__":
    # Check if folder is a git repository
    if not check_git_repo():
        console.print("[red]Not a git repository.[/red]")
        sys.exit(1)

    # Check for unpushed commits
    check_unpushed_commits()

    # Build prompt from arguments
    prompt = " ".join(parsed_args.prompt).strip() or None

    # Generate commit message
    try:
        response = generate_commit_message(prompt, parsed_args)
    except Exception as e:
        console.print(f"[red]Error generating commit message: {e}[/red]")
        sys.exit(1)

    # Format the response
    response = message_formatter(response)
    if not response:
        console.print("[red]No response from Copilot.[/red]")
        sys.exit(1)

    # Print the generated commit message
    console.print("[blue]===========================================[/blue]")
    if parsed_args.shortcodes:
        console.print(response)
    else:
        console.print(to_emoji(response))
    console.print("[blue]===========================================[/blue]")

    # Confirm and commit
    if parsed_args.yes or Confirm.ask(
        f"Do you want to use this commit message?",
        console=console,
        case_sensitive=False,
        default=True,
    ):
        try:
            subprocess.run(["git", "commit", "-m", response], check=True)
            console.print("[green]Commit successful![/green]")
        except subprocess.CalledProcessError:
            console.print("[red]Failed to commit changes.[/red]")
            sys.exit(1)
    else:
        console.print("[yellow]Commit aborted.[/yellow]")