#!/usr/bin/env python3
"""
Generate a DOT graph of commits from multiple lakeFS branches.

Uses the lakeFS SDK to fetch commit information including parent relationships.

Usage:
    python lakefs_commits_graph.py -r <repository> -b <branch1,branch2,...> [options]

Example:
    python lakefs_commits_graph.py -r my-repo -b main,feature-1 -d 7 -o graph.dot --svg
"""

import argparse
import os
import subprocess
import sys
from datetime import datetime, timedelta, timezone
from typing import Optional

import lakefs_sdk
from lakefs_sdk.client import LakeFSClient

# Color palette for branches
COLORS = ["red", "blue", "green", "orange", "purple", "brown", "pink", "cyan", "magenta", "olive"]


def get_client() -> LakeFSClient:
    """Create a lakeFS client from environment or lakectl config."""
    # Try environment variables first
    endpoint = os.environ.get("LAKEFS_ENDPOINT") or os.environ.get("LAKECTL_SERVER_ENDPOINT_URL")
    access_key = os.environ.get("LAKEFS_ACCESS_KEY_ID") or os.environ.get("LAKECTL_CREDENTIALS_ACCESS_KEY_ID")
    secret_key = os.environ.get("LAKEFS_SECRET_ACCESS_KEY") or os.environ.get("LAKECTL_CREDENTIALS_SECRET_ACCESS_KEY")

    # Try to read from lakectl config if not in environment
    if not all([endpoint, access_key, secret_key]):
        config_path = os.path.expanduser("~/.lakectl.yaml")
        if os.path.exists(config_path):
            import yaml
            with open(config_path) as f:
                config = yaml.safe_load(f)
                if config:
                    creds = config.get("credentials", {})
                    server = config.get("server", {})
                    endpoint = endpoint or server.get("endpoint_url")
                    access_key = access_key or creds.get("access_key_id")
                    secret_key = secret_key or creds.get("secret_access_key")

    if not all([endpoint, access_key, secret_key]):
        print("Error: lakeFS credentials not found.", file=sys.stderr)
        print("Set LAKEFS_ENDPOINT, LAKEFS_ACCESS_KEY_ID, LAKEFS_SECRET_ACCESS_KEY", file=sys.stderr)
        print("Or configure ~/.lakectl.yaml", file=sys.stderr)
        sys.exit(1)

    configuration = lakefs_sdk.Configuration(
        host=endpoint,
        username=access_key,
        password=secret_key,
    )
    return LakeFSClient(configuration)


def fetch_commits(client: LakeFSClient, repository: str, branch: str, since: datetime) -> list:
    """Fetch commits from a branch since a given date."""
    commits = []
    after = None
    since_ts = int(since.timestamp())

    while True:
        # Use the refs API to get commit log
        response = client.refs_api.log_commits(
            repository=repository,
            ref=branch,
            after=after,
            amount=1000,
        )

        for commit in response.results:
            # Check if commit is within our time range
            if commit.creation_date < since_ts:
                return commits
            commits.append(commit)

        if not response.pagination.has_more:
            break
        after = response.pagination.next_offset

    return commits


def get_branch_head(client: LakeFSClient, repository: str, branch: str) -> Optional[str]:
    """Get the head commit ID of a branch."""
    try:
        ref = client.branches_api.get_branch(repository=repository, branch=branch)
        return ref.commit_id
    except Exception as e:
        print(f"    Warning: Could not get branch head for {branch}: {e}", file=sys.stderr)
        return None


def generate_dot(
    commits: dict,
    edges: set,
    branch_tips: dict,
    branch_colors: dict,
    repository: str,
    branches: list,
    since: datetime,
) -> str:
    """Generate DOT format graph."""
    now = datetime.now(timezone.utc)
    
    lines = [
        "digraph commits {",
        "    rankdir=TB;",
        '    node [shape=box, style=filled, fillcolor=lightyellow, fontname="Courier", fontsize=10];',
        "    edge [color=gray];",
        "",
        "    // Graph title",
        '    labelloc="t";',
        f'    label="lakeFS Commits Graph\\n{repository}\\nBranches: {",".join(branches)}\\nSince: {since.isoformat()}";',
        "",
    ]

    # Add commit nodes
    lines.append("    // Commit nodes")
    for commit_id, commit_info in commits.items():
        short_sha = commit_id[:8]
        date_str = commit_info["date"]
        color = commit_info.get("color", "gray")
        
        # Calculate days ago
        commit_ts = commit_info.get("timestamp")
        if commit_ts:
            days_ago = (now - commit_ts).days
            if days_ago > 0:
                date_str = f"{date_str} ({days_ago}d)"
        
        label = f"{short_sha}\\n{date_str}"
        lines.append(f'    "{commit_id}" [label="{label}", color={color}, penwidth=2];')

    # Add branch labels
    lines.append("")
    lines.append("    // Branch labels")
    for branch_name, tip_sha in branch_tips.items():
        color = branch_colors.get(branch_name, "black")
        lines.append(
            f'    "branch_{branch_name}" [label="{branch_name}", shape=ellipse, fillcolor=lightblue, color={color}, penwidth=2];'
        )
        lines.append(f'    "branch_{branch_name}" -> "{tip_sha}" [style=dashed, color={color}, penwidth=2];')

    # Add edges (child -> parent)
    lines.append("")
    lines.append("    // Commit relationships (child -> parent)")

    # Track which parent nodes we've added as placeholders
    placeholder_nodes = set()

    for child, parent in sorted(edges):
        if child in commits:
            if parent in commits:
                lines.append(f'    "{child}" -> "{parent}";')
            else:
                # Parent is outside our time window - add placeholder
                if parent not in placeholder_nodes:
                    parent_short = parent[:8]
                    lines.append(
                        f'    "{parent}" [label="{parent_short}\\n(older)", style="filled,dashed", fillcolor=lightgray];'
                    )
                    placeholder_nodes.add(parent)
                lines.append(f'    "{child}" -> "{parent}" [style=dashed];')

    lines.append("}")
    return "\n".join(lines)


def main():
    parser = argparse.ArgumentParser(
        description="Generate a DOT graph visualization of commits from lakeFS branches."
    )
    parser.add_argument(
        "-r", "--repository", required=True, help="lakeFS repository name (e.g., my-repo)"
    )
    parser.add_argument(
        "-b",
        "--branches",
        required=True,
        help="Comma-separated list of branches (e.g., main,dev,feature)",
    )
    parser.add_argument(
        "-d", "--days", type=int, default=30, help="Number of days to look back (default: 30)"
    )
    parser.add_argument(
        "-o", "--output", default="commits_graph.dot", help="Output DOT file (default: commits_graph.dot)"
    )
    parser.add_argument("--svg", action="store_true", help="Generate SVG file (requires 'dot' from graphviz)")
    parser.add_argument("--png", action="store_true", help="Generate PNG file (requires 'dot' from graphviz)")

    args = parser.parse_args()

    # Parse branches
    branches = [b.strip() for b in args.branches.split(",")]

    # Calculate since date
    since = datetime.now(timezone.utc) - timedelta(days=args.days)

    print(f"Fetching commits since: {since.isoformat()}")
    print(f"Processing branches: {', '.join(branches)}")

    # Initialize client
    client = get_client()

    # Data structures
    all_commits = {}  # commit_id -> {date, color, parents}
    all_edges = set()  # (child, parent) tuples
    branch_tips = {}  # branch_name -> commit_id
    branch_colors = {}  # branch_name -> color
    commit_to_branch = {}  # commit_id -> first branch that contains it

    # Process each branch
    for idx, branch in enumerate(branches):
        color = COLORS[idx % len(COLORS)]
        branch_colors[branch] = color

        print(f"  Fetching commits from branch: {branch}")

        try:
            commits = fetch_commits(client, args.repository, branch, since)
        except Exception as e:
            print(f"    Warning: Failed to fetch commits from {branch}: {e}", file=sys.stderr)
            continue

        print(f"    Found {len(commits)} commits")

        for commit in commits:
            commit_id = commit.id
            creation_date = datetime.fromtimestamp(commit.creation_date, tz=timezone.utc)
            date_str = creation_date.strftime("%Y-%m-%d %H:%M:%S")

            # Store commit info (only if not already seen)
            if commit_id not in all_commits:
                all_commits[commit_id] = {
                    "date": date_str,
                    "timestamp": creation_date,
                    "color": color,
                    "parents": commit.parents,
                    "message": commit.message,
                }
                commit_to_branch[commit_id] = branch

            # Add edges for all parents
            for parent_id in commit.parents:
                all_edges.add((commit_id, parent_id))

        # Get branch tip
        tip = get_branch_head(client, args.repository, branch)
        if tip:
            branch_tips[branch] = tip

    print(f"\nFound {len(all_commits)} unique commits across all branches")
    print(f"Found {len(all_edges)} edges (parent relationships)")

    # Generate DOT file
    print(f"\nGenerating DOT file: {args.output}")
    dot_content = generate_dot(
        all_commits,
        all_edges,
        branch_tips,
        branch_colors,
        args.repository,
        branches,
        since,
    )

    with open(args.output, "w") as f:
        f.write(dot_content)

    print(f"  - Output: {args.output}")

    # Generate SVG if requested
    if args.svg:
        svg_file = args.output.rsplit(".", 1)[0] + ".svg"
        print(f"\nGenerating SVG: {svg_file}")
        try:
            subprocess.run(["dot", "-Tsvg", args.output, "-o", svg_file], check=True)
            print(f"  - SVG output: {svg_file}")
        except FileNotFoundError:
            print("Error: 'dot' command not found. Please install graphviz.", file=sys.stderr)
            print("  macOS: brew install graphviz", file=sys.stderr)
            print("  Ubuntu: apt-get install graphviz", file=sys.stderr)
        except subprocess.CalledProcessError as e:
            print(f"Error generating SVG: {e}", file=sys.stderr)

    # Generate PNG if requested
    if args.png:
        png_file = args.output.rsplit(".", 1)[0] + ".png"
        print(f"\nGenerating PNG: {png_file}")
        try:
            subprocess.run(["dot", "-Tpng", args.output, "-o", png_file], check=True)
            print(f"  - PNG output: {png_file}")
        except FileNotFoundError:
            print("Error: 'dot' command not found. Please install graphviz.", file=sys.stderr)
        except subprocess.CalledProcessError as e:
            print(f"Error generating PNG: {e}", file=sys.stderr)

    print("\nDone!")
    print("\nTo generate images manually:")
    base_name = args.output.rsplit(".", 1)[0]
    print(f"  dot -Tpng {args.output} -o {base_name}.png")
    print(f"  dot -Tsvg {args.output} -o {base_name}.svg")


if __name__ == "__main__":
    main()