#!/usr/bin/env python3

import os
import logging
import subprocess
import dataclasses

logger = logging.getLogger(__name__)

STATUS_COLORS = {
    "untracked": 93,
    "changed": 31,
    "clean": 32,
    "unknown": 94,
}

FORMAT_STRINGS = {
    "default": "\033[{color}m{name:20.20s} {branch:10.10s} {status:10.10s} {changes}\033[0m"
}


@dataclasses.dataclass
class GitStatus:
    name: str
    branch: str = ""
    status: str = "untracked"
    modified: int = 0
    deleted: int = 0
    untracked: int = 0

    def info(self, fmt):
        changes = []
        if self.modified:
            changes.append(f"{self.modified}M")
        if self.deleted:
            changes.append(f"{self.deleted}D")
        if self.untracked:
            changes.append(f"{self.untracked}U")

        changes = f" [{', '.join(changes)}]" if changes else ""
        color = STATUS_COLORS.get(self.status, STATUS_COLORS["unknown"])
        fmt = fmt or FORMAT_STRINGS["default"]

        return fmt.format(changes=changes, color=color, **dataclasses.asdict(self))


def multi_repo_overview(directory: str, all: bool = False):
    result = []
    for f in os.scandir(directory):
        if not f.is_dir():
            continue

        try:
            result.append(repo_overview(f.path))
        except FileNotFoundError:
            if all:
                result.append(GitStatus(name=f.name, status="-"))

    return result


def repo_overview(repo):
    if not os.path.isdir(os.path.join(repo, ".git")):
        raise FileNotFoundError(f"{repo} is not a git repo")

    p = subprocess.run(("git", "-C", repo, "status"), capture_output=True)
    lines = p.stdout.decode().split("\n")

    info = GitStatus(name=os.path.basename(repo))
    for line in lines:
        if line.startswith("On branch"):
            info.branch = line[10:]

        if line.startswith("Changes "):
            info.status = "changed"

        if line.startswith("Your branch is up to date"):
            info.status = "clean"

        if line.startswith("\tmodified:"):
            info.modified += 1
        elif line.startswith("\tdeleted:"):
            info.deleted += 1
        elif line.startswith("\t"):
            info.untracked += 1

    return info


def cmd():
    import argparse

    parser = argparse.ArgumentParser(
        description="Check the work trees in a directory for their status"
    )
    parser.add_argument("--log-level", "-l", default="info")
    parser.add_argument(
        "--all",
        "-a",
        action="store_true",
        help="Include directories without a git folder",
    )
    parser.add_argument("--format", "-f", default="default")
    parser.add_argument("directory", nargs="?", default=".")
    options = parser.parse_args()

    logging.basicConfig(level=getattr(logging, options.log_level.upper(), logging.INFO))
    logger.debug(options)

    repos = multi_repo_overview(options.directory, options.all)

    fmt = FORMAT_STRINGS.get(options.format, FORMAT_STRINGS["default"])

    print(
        fmt.format(
            color=0,
            name="Workdir",
            branch="Branch",
            status="Status",
            changes="",
        )
    )
    for repo in repos:
        print(repo.info(fmt))


if __name__ == "__main__":
    cmd()