msr.py

import json
import os
from datetime import datetime, timezone, timedelta
from collections import defaultdict
from huggingface_hub import HfApi, hf_hub_download
from huggingface_hub.errors import HfHubHTTPError
from dotenv import load_dotenv
import duckdb
import backoff
import requests
import requests.exceptions
import traceback
import re

# Load environment variables
load_dotenv(override=True)

# =============================================================================
# CONFIGURATION
# =============================================================================

# Get script directory for relative paths
SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
BASE_DIR = os.path.dirname(SCRIPT_DIR)  # Parent directory

AGENTS_REPO = "SWE-Arena/bot_data"
AGENTS_REPO_LOCAL_PATH = os.path.join(BASE_DIR, "bot_data")  # Local git clone path
DUCKDB_CACHE_FILE = os.path.join(SCRIPT_DIR, "cache.duckdb")
GHARCHIVE_DATA_LOCAL_PATH = os.path.join(BASE_DIR, "gharchive/data")
LEADERBOARD_FILENAME = f"{os.getenv('COMPOSE_PROJECT_NAME')}.json"
LEADERBOARD_REPO = "SWE-Arena/leaderboard_data"
LEADERBOARD_TIME_FRAME_DAYS = 180

# Git sync configuration (mandatory to get latest bot data)
GIT_SYNC_TIMEOUT = 300  # 5 minutes timeout for git pull

# Streaming batch configuration
BATCH_SIZE_DAYS = 1  # Process 1 day at a time (~24 hourly files)

# Retry configuration
MAX_RETRIES = 5

# =============================================================================
# UTILITY FUNCTIONS
# =============================================================================

def load_jsonl(filename):
    """Load JSONL file and return list of dictionaries."""
    if not os.path.exists(filename):
        return []

    data = []
    with open(filename, 'r', encoding='utf-8') as f:
        for line in f:
            line = line.strip()
            if line:
                try:
                    data.append(json.loads(line))
                except json.JSONDecodeError as e:
                    print(f"Warning: Skipping invalid JSON line: {e}")
    return data


def save_jsonl(filename, data):
    """Save list of dictionaries to JSONL file."""
    with open(filename, 'w', encoding='utf-8') as f:
        for item in data:
            f.write(json.dumps(item) + '\n')


def normalize_date_format(date_string):
    """Convert date strings or datetime objects to standardized ISO 8601 format with Z suffix."""
    if not date_string or date_string == 'N/A':
        return 'N/A'

    try:
        if isinstance(date_string, datetime):
            return date_string.strftime('%Y-%m-%dT%H:%M:%SZ')

        date_string = re.sub(r'\s+', ' ', date_string.strip())
        date_string = date_string.replace(' ', 'T')

        if len(date_string) >= 3:
            if date_string[-3:-2] in ('+', '-') and ':' not in date_string[-3:]:
                date_string = date_string + ':00'

        dt = datetime.fromisoformat(date_string.replace('Z', '+00:00'))
        return dt.strftime('%Y-%m-%dT%H:%M:%SZ')
    except Exception as e:
        print(f"Warning: Could not parse date '{date_string}': {e}")
        return date_string


def get_hf_token():
    """Get HuggingFace token from environment variables."""
    token = os.getenv('HF_TOKEN')
    if not token:
        print("Warning: HF_TOKEN not found in environment variables")
    return token


# =============================================================================
# GHARCHIVE DOWNLOAD FUNCTIONS
# =============================================================================

def download_file(url):
    """Download a GHArchive file with retry logic."""
    filename = url.split("/")[-1]
    filepath = os.path.join(GHARCHIVE_DATA_LOCAL_PATH, filename)

    if os.path.exists(filepath):
        return True

    try:
        response = requests.get(url, timeout=30)
        response.raise_for_status()
        with open(filepath, "wb") as f:
            f.write(response.content)
        return True
    except Exception as e:
        print(f"   ⚠ {filename}: {e}")
        return False


def download_all_gharchive_data():
    """Download all GHArchive data files for the last LEADERBOARD_TIME_FRAME_DAYS."""
    os.makedirs(GHARCHIVE_DATA_LOCAL_PATH, exist_ok=True)

    end_date = datetime.now(timezone.utc)
    start_date = end_date - timedelta(days=LEADERBOARD_TIME_FRAME_DAYS)

    urls = []
    current_date = start_date
    while current_date <= end_date:
        date_str = current_date.strftime("%Y-%m-%d")
        for hour in range(24):
            url = f"https://data.gharchive.org/{date_str}-{hour}.json.gz"
            urls.append(url)
        current_date += timedelta(days=1)

    success = True
    for url in urls:
        if not download_file(url):
            success = False

    return success

# =============================================================================
# HUGGINGFACE API WRAPPERS
# =============================================================================

def is_retryable_error(e):
    """Check if exception is retryable (rate limit or timeout error)."""
    if isinstance(e, HfHubHTTPError):
        if e.response.status_code == 429:
            return True

    if isinstance(e, (requests.exceptions.Timeout,
                     requests.exceptions.ReadTimeout,
                     requests.exceptions.ConnectTimeout)):
        return True

    if isinstance(e, Exception):
        error_str = str(e).lower()
        if 'timeout' in error_str or 'timed out' in error_str:
            return True

    return False


@backoff.on_exception(
    backoff.expo,
    (HfHubHTTPError, requests.exceptions.Timeout, requests.exceptions.RequestException, Exception),
    max_tries=MAX_RETRIES,
    base=300,
    max_value=3600,
    giveup=lambda e: not is_retryable_error(e),
    on_backoff=lambda details: print(
        f"   {details['exception']} error. Retrying in {details['wait']/60:.1f} minutes ({details['wait']:.0f}s) - attempt {details['tries']}/5..."
    )
)
def list_repo_files_with_backoff(api, **kwargs):
    """Wrapper for api.list_repo_files() with exponential backoff."""
    return api.list_repo_files(**kwargs)


@backoff.on_exception(
    backoff.expo,
    (HfHubHTTPError, requests.exceptions.Timeout, requests.exceptions.RequestException, Exception),
    max_tries=MAX_RETRIES,
    base=300,
    max_value=3600,
    giveup=lambda e: not is_retryable_error(e),
    on_backoff=lambda details: print(
        f"   {details['exception']} error. Retrying in {details['wait']/60:.1f} minutes ({details['wait']:.0f}s) - attempt {details['tries']}/5..."
    )
)
def hf_hub_download_with_backoff(**kwargs):
    """Wrapper for hf_hub_download() with exponential backoff."""
    return hf_hub_download(**kwargs)


@backoff.on_exception(
    backoff.expo,
    (HfHubHTTPError, requests.exceptions.Timeout, requests.exceptions.RequestException, Exception),
    max_tries=MAX_RETRIES,
    base=300,
    max_value=3600,
    giveup=lambda e: not is_retryable_error(e),
    on_backoff=lambda details: print(
        f"   {details['exception']} error. Retrying in {details['wait']/60:.1f} minutes ({details['wait']:.0f}s) - attempt {details['tries']}/5..."
    )
)
def upload_file_with_backoff(api, **kwargs):
    """Wrapper for api.upload_file() with exponential backoff."""
    return api.upload_file(**kwargs)


@backoff.on_exception(
    backoff.expo,
    (HfHubHTTPError, requests.exceptions.Timeout, requests.exceptions.RequestException, Exception),
    max_tries=MAX_RETRIES,
    base=300,
    max_value=3600,
    giveup=lambda e: not is_retryable_error(e),
    on_backoff=lambda details: print(
        f"   {details['exception']} error. Retrying in {details['wait']/60:.1f} minutes ({details['wait']:.0f}s) - attempt {details['tries']}/5..."
    )
)
def upload_folder_with_backoff(api, **kwargs):
    """Wrapper for api.upload_folder() with exponential backoff."""
    return api.upload_folder(**kwargs)


def get_duckdb_connection():
    """
    Initialize DuckDB connection with OPTIMIZED memory settings.
    Uses persistent database and reduced memory footprint.
    Automatically removes cache file if lock conflict is detected.
    """
    try:
        conn = duckdb.connect(DUCKDB_CACHE_FILE)
    except Exception as e:
        # Check if it's a locking error
        error_msg = str(e)
        if "lock" in error_msg.lower() or "conflicting" in error_msg.lower():
            print(f"   ⚠ Lock conflict detected, removing {DUCKDB_CACHE_FILE}...")
            if os.path.exists(DUCKDB_CACHE_FILE):
                os.remove(DUCKDB_CACHE_FILE)
                print(f"   ✓ Cache file removed, retrying connection...")
            # Retry connection after removing cache
            conn = duckdb.connect(DUCKDB_CACHE_FILE)
        else:
            # Re-raise if it's not a locking error
            raise

    # CORE MEMORY & THREADING SETTINGS
    conn.execute(f"SET threads TO 4;")
    conn.execute(f"SET max_memory = '50GB';")
    conn.execute("SET temp_directory = '/tmp/duckdb_temp';")

    # PERFORMANCE OPTIMIZATIONS
    conn.execute("SET preserve_insertion_order = false;")  # Disable expensive ordering
    conn.execute("SET enable_object_cache = true;")  # Cache repeatedly read files

    return conn


def generate_file_path_patterns(start_date, end_date, data_dir=GHARCHIVE_DATA_LOCAL_PATH):
    """Generate file path patterns for GHArchive data in date range (only existing files)."""
    file_patterns = []
    missing_dates = set()

    current_date = start_date.replace(hour=0, minute=0, second=0, microsecond=0)
    end_day = end_date.replace(hour=0, minute=0, second=0, microsecond=0)

    while current_date <= end_day:
        date_has_files = False
        for hour in range(24):
            pattern = os.path.join(data_dir, f"{current_date.strftime('%Y-%m-%d')}-{hour}.json.gz")
            if os.path.exists(pattern):
                file_patterns.append(pattern)
                date_has_files = True

        if not date_has_files:
            missing_dates.add(current_date.strftime('%Y-%m-%d'))

        current_date += timedelta(days=1)

    if missing_dates:
        print(f"   ⚠ Skipping {len(missing_dates)} date(s) with no data")

    return file_patterns


# =============================================================================
# STREAMING BATCH PROCESSING
# =============================================================================

def fetch_all_metadata_streaming(conn, identifiers, start_date, end_date):
    """
    UNIFIED QUERY: Fetch PR, review, and commit metadata using streaming batch processing.
    Single query per batch retrieves all event types and post-processes in Python:
    - PushEvent (commits)
    - PullRequestEvent (authored PRs)
    - PullRequestReviewEvent & PullRequestReviewCommentEvent (reviews)

    Args:
        conn: DuckDB connection instance
        identifiers: List of GitHub usernames/bot identifiers
        start_date: Start datetime (timezone-aware)
        end_date: End datetime (timezone-aware)

    Returns:
        Dictionary with three keys:
        - 'commits': {author: [commit_metadata]}
        - 'prs': {author: [pr_metadata]}
        - 'reviews': {reviewer: [review_metadata]}
    """
    identifier_list = ', '.join([f"'{id}'" for id in identifiers])
    identifiers_set = set(identifiers)

    commits_by_agent = defaultdict(list)
    prs_by_agent = defaultdict(list)
    reviews_by_agent = defaultdict(list)

    total_days = (end_date - start_date).days
    total_batches = (total_days // BATCH_SIZE_DAYS) + 1

    current_date = start_date
    batch_num = 0
    total_commits = 0
    total_prs = 0
    total_reviews = 0

    print(f"   Streaming {total_batches} batches of {BATCH_SIZE_DAYS}-day intervals...")
    
    while current_date <= end_date:
        batch_num += 1
        batch_end = min(current_date + timedelta(days=BATCH_SIZE_DAYS - 1), end_date)
        
        # Get file patterns for THIS BATCH ONLY (not all 180 days)
        file_patterns = generate_file_path_patterns(current_date, batch_end)
        
        if not file_patterns:
            print(f"   Batch {batch_num}/{total_batches}: {current_date.date()} to {batch_end.date()} - NO DATA")
            current_date = batch_end + timedelta(days=1)
            continue
        
        # Progress indicator
        print(f"   Batch {batch_num}/{total_batches}: {current_date.date()} to {batch_end.date()} ({len(file_patterns)} files)... ", end="", flush=True)
        
        # Build file patterns SQL for THIS BATCH
        file_patterns_sql = '[' + ', '.join([f"'{fp}'" for fp in file_patterns]) + ']'

        # UNIFIED QUERY: Fetch PRs, reviews, and commits in a single pass
        unified_query = f"""
        SELECT
            type,
            TRY_CAST(json_extract_string(to_json(actor), '$.login') AS VARCHAR) as actor_login,
            TRY_CAST(json_extract_string(to_json(payload), '$.head') AS VARCHAR) as commit_sha,
            CONCAT(
                REPLACE(repo.url, 'api.github.com/repos/', 'github.com/'),
                '/pull/',
                CAST(payload.pull_request.number AS VARCHAR)
            ) as pr_url,
            TRY_CAST(json_extract_string(to_json(payload), '$.action') AS VARCHAR) as action,
            TRY_CAST(json_extract_string(to_json(payload), '$.pull_request.user.login') AS VARCHAR) as pr_author,
            TRY_CAST(json_extract_string(to_json(payload), '$.pull_request.created_at') AS VARCHAR) as pr_created_at,
            TRY_CAST(json_extract_string(to_json(payload), '$.pull_request.merged_at') AS VARCHAR) as pr_merged_at,
            TRY_CAST(json_extract_string(to_json(payload), '$.pull_request.closed_at') AS VARCHAR) as pr_closed_at,
            TRY_CAST(json_extract_string(to_json(payload), '$.comment.user.login') AS VARCHAR) as comment_user_login,
            created_at
        FROM read_json(
            {file_patterns_sql},
            union_by_name=true,
            filename=true,
            compression='gzip',
            format='newline_delimited',
            ignore_errors=true
        )
        WHERE
            -- PushEvent: Commits by assistants
            (type = 'PushEvent'
                AND TRY_CAST(json_extract_string(to_json(payload), '$.head') AS VARCHAR) IS NOT NULL
                AND TRY_CAST(json_extract_string(to_json(actor), '$.login') AS VARCHAR) IN ({identifier_list}))
            OR
            -- PullRequestEvent: PRs authored by assistants (opened/closed)
            (type = 'PullRequestEvent'
                AND payload.pull_request.number IS NOT NULL
                AND TRY_CAST(json_extract_string(to_json(payload), '$.pull_request.created_at') AS VARCHAR) IS NOT NULL
                AND TRY_CAST(json_extract_string(to_json(payload), '$.pull_request.user.login') AS VARCHAR) IN ({identifier_list})
                AND TRY_CAST(json_extract_string(to_json(payload), '$.action') AS VARCHAR) IN ('opened', 'closed'))
            OR
            -- PullRequestReviewEvent: Reviews by assistants
            (type = 'PullRequestReviewEvent'
                AND payload.pull_request.number IS NOT NULL
                AND TRY_CAST(json_extract_string(to_json(actor), '$.login') AS VARCHAR) IN ({identifier_list}))
            OR
            -- PullRequestReviewCommentEvent: Review comments by assistants
            (type = 'PullRequestReviewCommentEvent'
                AND payload.pull_request.number IS NOT NULL
                AND TRY_CAST(json_extract_string(to_json(payload), '$.comment.user.login') AS VARCHAR) IN ({identifier_list}))
        """
        
        try:
            all_results = conn.execute(unified_query).fetchall()

            # Post-process results to separate PRs, reviews, and commits
            # Row: (type, actor_login, commit_sha, pr_url, action, pr_author,
            #        pr_created_at, pr_merged_at, pr_closed_at, comment_user_login, created_at)

            commit_rows = []
            authored_pr_rows = []
            review_rows = []

            for row in all_results:
                event_type = row[0]
                if event_type == 'PushEvent':
                    commit_rows.append(row)
                elif event_type in ('PullRequestReviewEvent', 'PullRequestReviewCommentEvent'):
                    review_rows.append(row)
                elif event_type == 'PullRequestEvent':
                    action = row[4]
                    pr_author = row[5]
                    if pr_author in identifiers_set and action in ('opened', 'closed'):
                        authored_pr_rows.append(row)

            # --- Process commits ---
            batch_commits = 0
            for row in commit_rows:
                author = row[1]
                sha = row[2]
                created_at = normalize_date_format(row[10]) if row[10] else None

                if not author or not sha:
                    continue

                commits_by_agent[author].append({
                    'sha': sha,
                    'created_at': created_at,
                })
                batch_commits += 1
                total_commits += 1

            # --- Process authored PRs ---
            # Group events by PR URL to merge opened and closed events
            pr_events = defaultdict(lambda: {'opened': None, 'closed': None})

            for row in authored_pr_rows:
                pr_url = row[3]
                action = row[4]
                pr_author = row[5]
                pr_created_at = normalize_date_format(row[6]) if row[6] else None
                pr_merged_at = normalize_date_format(row[7]) if row[7] else None
                pr_closed_at = normalize_date_format(row[8]) if row[8] else None

                if not pr_url or not action:
                    continue

                pr_events[pr_url][action] = {
                    'pr_author': pr_author,
                    'created_at': pr_created_at,
                    'merged_at': pr_merged_at,
                    'closed_at': pr_closed_at,
                }

            batch_prs = 0
            for url, events in pr_events.items():
                if not events['opened']:
                    continue

                opened_event = events['opened']
                closed_event = events['closed']

                pr_author = opened_event['pr_author']
                if not pr_author:
                    continue

                prs_by_agent[pr_author].append({
                    'html_url': url,
                    'created_at': opened_event['created_at'],
                    'merged_at': closed_event['merged_at'] if closed_event else None,
                    'closed_at': closed_event['closed_at'] if closed_event else None,
                })
                batch_prs += 1
                total_prs += 1

            # --- Process reviews ---
            batch_reviews = 0
            for row in review_rows:
                event_type = row[0]
                if event_type == 'PullRequestReviewEvent':
                    reviewer = row[1]  # actor_login
                else:  # PullRequestReviewCommentEvent
                    reviewer = row[9]  # comment_user_login

                pr_url = row[3]
                reviewed_at = normalize_date_format(row[10]) if row[10] else None

                if not reviewer or not pr_url or not reviewed_at:
                    continue

                reviews_by_agent[reviewer].append({
                    'url': pr_url,
                    'reviewed_at': reviewed_at,
                })
                batch_reviews += 1
                total_reviews += 1

            print(f"✓ {batch_prs} PRs, {batch_reviews} reviews, {batch_commits} commits found")
            
        except Exception as e:
            print(f"\n   ✗ Batch {batch_num} error: {str(e)}")
            traceback.print_exc()
        
        # Move to next batch
        current_date = batch_end + timedelta(days=1)

    # Final summary
    agents_with_commits = sum(1 for commits in commits_by_agent.values() if commits)
    agents_with_prs = sum(1 for prs in prs_by_agent.values() if prs)
    agents_with_reviews = sum(1 for reviews in reviews_by_agent.values() if reviews)
    print(f"\n   ✓ Complete: {total_prs} PRs for {agents_with_prs}/{len(identifiers)} assistants")
    print(f"   ✓ Complete: {total_reviews} reviews for {agents_with_reviews}/{len(identifiers)} assistants")
    print(f"   ✓ Complete: {total_commits} commits for {agents_with_commits}/{len(identifiers)} assistants")

    return {
        'commits': dict(commits_by_agent),
        'prs': dict(prs_by_agent),
        'reviews': dict(reviews_by_agent),
    }


def load_agents_from_hf():
    """
    Load all assistant metadata JSON files from local git repository.
    """
    assistants = []

    # Scan local directory for JSON files
    if not os.path.exists(AGENTS_REPO_LOCAL_PATH):
        raise FileNotFoundError(f"Local repository not found at {AGENTS_REPO_LOCAL_PATH}")

    # Walk through the directory to find all JSON files
    files_processed = 0
    print(f"   Loading assistant metadata from {AGENTS_REPO_LOCAL_PATH}...")

    for root, dirs, files in os.walk(AGENTS_REPO_LOCAL_PATH):
        # Skip .git directory
        if '.git' in root:
            continue

        for filename in files:
            if not filename.endswith('.json'):
                continue

            files_processed += 1
            file_path = os.path.join(root, filename)

            try:
                with open(file_path, 'r', encoding='utf-8') as f:
                    agent_data = json.load(f)

                # Only include active assistants
                if agent_data.get('status') != 'active':
                    continue

                # Extract github_identifier from filename
                github_identifier = filename.replace('.json', '')
                agent_data['github_identifier'] = github_identifier

                assistants.append(agent_data)

            except Exception as e:
                print(f"   ⚠ Error loading {filename}: {str(e)}")
                continue

    print(f"   ✓ Loaded {len(assistants)} active assistants (from {files_processed} total files)")
    return assistants


def calculate_commit_stats_from_metadata(metadata_list):
    """Calculate statistics from a list of commit metadata."""
    total_commits = len(metadata_list)

    return {
        'total_commits': total_commits,
    }


def calculate_pr_stats_from_metadata(metadata_list):
    """Calculate statistics from a list of PR metadata."""
    total_prs = len(metadata_list)
    merged = sum(1 for pr_meta in metadata_list if get_pr_status(pr_meta) == 'merged')
    closed_not_merged = sum(1 for pr_meta in metadata_list if get_pr_status(pr_meta) == 'closed')

    total_decisions = merged + closed_not_merged
    acceptance_rate = (merged / total_decisions * 100) if total_decisions > 0 else 0

    return {
        'total_prs': total_prs,
        'merged_prs': merged,
        'acceptance_rate': round(acceptance_rate, 2),
    }


def get_pr_status(meta):
    """Derive PR status from merged_at and closed_at fields."""
    if meta.get('merged_at'):
        return 'merged'
    elif meta.get('closed_at'):
        return 'closed'
    else:
        return 'open'


def calculate_review_stats_from_metadata(metadata_list):
    """Calculate statistics from a list of review metadata."""
    return {
        'total_reviews': len(metadata_list),
    }


def calculate_monthly_metrics_by_agent_commits(all_metadata_dict, assistants):
    """Calculate monthly metrics for commits for all assistants for visualization."""
    identifier_to_name = {assistant.get('github_identifier'): assistant.get('name') for assistant in assistants if assistant.get('github_identifier')}

    if not all_metadata_dict:
        return {'assistants': [], 'months': [], 'data': {}}

    agent_month_data = defaultdict(lambda: defaultdict(list))

    for agent_identifier, metadata_list in all_metadata_dict.items():
        for commit_meta in metadata_list:
            created_at = commit_meta.get('created_at')

            if not created_at:
                continue

            agent_name = identifier_to_name.get(agent_identifier, agent_identifier)

            try:
                dt = datetime.fromisoformat(created_at.replace('Z', '+00:00'))
                month_key = f"{dt.year}-{dt.month:02d}"
                agent_month_data[agent_name][month_key].append(commit_meta)
            except Exception as e:
                print(f"Warning: Could not parse date '{created_at}': {e}")
                continue

    all_months = set()
    for agent_data in agent_month_data.values():
        all_months.update(agent_data.keys())
    months = sorted(list(all_months))

    result_data = {}
    for agent_name, month_dict in agent_month_data.items():
        total_commits_list = []

        for month in months:
            commits_in_month = month_dict.get(month, [])
            total_count = len(commits_in_month)

            total_commits_list.append(total_count)

        result_data[agent_name] = {
            'total_commits': total_commits_list,
        }

    agents_list = sorted(list(agent_month_data.keys()))

    return {
        'assistants': agents_list,
        'months': months,
        'data': result_data
    }


def calculate_monthly_metrics_by_agent_prs(all_metadata_dict, assistants):
    """Calculate monthly metrics for PRs for all assistants for visualization."""
    identifier_to_name = {assistant.get('github_identifier'): assistant.get('name') for assistant in assistants if assistant.get('github_identifier')}

    if not all_metadata_dict:
        return {'assistants': [], 'months': [], 'data': {}}

    agent_month_data = defaultdict(lambda: defaultdict(list))

    for agent_identifier, metadata_list in all_metadata_dict.items():
        for pr_meta in metadata_list:
            created_at = pr_meta.get('created_at')

            if not created_at:
                continue

            agent_name = identifier_to_name.get(agent_identifier, agent_identifier)

            try:
                dt = datetime.fromisoformat(created_at.replace('Z', '+00:00'))                
                month_key = f"{dt.year}-{dt.month:02d}"
                agent_month_data[agent_name][month_key].append(pr_meta)
            except Exception as e:
                print(f"Warning: Could not parse date '{created_at}': {e}")
                continue

    all_months = set()
    for agent_data in agent_month_data.values():
        all_months.update(agent_data.keys())
    months = sorted(list(all_months))

    result_data = {}
    for agent_name, month_dict in agent_month_data.items():
        acceptance_rates = []
        total_prs_list = []
        merged_prs_list = []
        closed_not_merged_list = []

        for month in months:
            prs_in_month = month_dict.get(month, [])

            merged_count = sum(1 for pr in prs_in_month if get_pr_status(pr) == 'merged')
            closed_not_merged_count = sum(1 for pr in prs_in_month if get_pr_status(pr) == 'closed')
            total_count = len(prs_in_month)

            total_decisions = merged_count + closed_not_merged_count
            acceptance_rate = (merged_count / total_decisions * 100) if total_decisions > 0 else None

            acceptance_rates.append(acceptance_rate)
            total_prs_list.append(total_count)
            merged_prs_list.append(merged_count)
            closed_not_merged_list.append(closed_not_merged_count)

        result_data[agent_name] = {
            'acceptance_rates': acceptance_rates,
            'total_prs': total_prs_list,
            'merged_prs': merged_prs_list,
            'closed_not_merged': closed_not_merged_list
        }

    agents_list = sorted(list(agent_month_data.keys()))

    return {
        'assistants': agents_list,
        'months': months,
        'data': result_data
    }


def calculate_monthly_metrics_by_agent_reviews(all_metadata_dict, assistants):
    """Calculate monthly metrics for reviews for all assistants for visualization."""
    identifier_to_name = {assistant.get('github_identifier'): assistant.get('name') for assistant in assistants if assistant.get('github_identifier')}

    if not all_metadata_dict:
        return {'assistants': [], 'months': [], 'data': {}}

    agent_month_data = defaultdict(lambda: defaultdict(list))

    for agent_identifier, metadata_list in all_metadata_dict.items():
        for review_meta in metadata_list:
            reviewed_at = review_meta.get('reviewed_at')

            if not reviewed_at:
                continue

            agent_name = identifier_to_name.get(agent_identifier, agent_identifier)

            try:
                dt = datetime.fromisoformat(reviewed_at.replace('Z', '+00:00'))
                month_key = f"{dt.year}-{dt.month:02d}"
                agent_month_data[agent_name][month_key].append(review_meta)
            except Exception as e:
                print(f"Warning: Could not parse date '{reviewed_at}': {e}")
                continue

    all_months = set()
    for agent_data in agent_month_data.values():
        all_months.update(agent_data.keys())
    months = sorted(list(all_months))

    result_data = {}
    for agent_name, month_dict in agent_month_data.items():
        total_reviews_list = []

        for month in months:
            reviews_in_month = month_dict.get(month, [])
            total_reviews_list.append(len(reviews_in_month))

        result_data[agent_name] = {
            'total_reviews': total_reviews_list,
        }

    agents_list = sorted(list(agent_month_data.keys()))

    return {
        'assistants': agents_list,
        'months': months,
        'data': result_data
    }


def construct_leaderboard_from_metadata(commit_metadata_dict, pr_metadata_dict, review_metadata_dict, assistants):
    """Construct leaderboard from in-memory PR, review, and commit metadata.

    Args:
        commit_metadata_dict: Dictionary mapping assistant ID to list of commit metadata
        pr_metadata_dict: Dictionary mapping assistant ID to list of PR metadata
        review_metadata_dict: Dictionary mapping assistant ID to list of review metadata
        assistants: List of assistant metadata

    Returns:
        Dictionary with leaderboard data including PR, review, and commit statistics
    """
    if not assistants:
        print("Error: No assistants found")
        return {}

    cache_dict = {}

    for assistant in assistants:
        identifier = assistant.get('github_identifier')
        agent_name = assistant.get('name', 'Unknown')

        commit_metadata = commit_metadata_dict.get(identifier, [])
        pr_metadata = pr_metadata_dict.get(identifier, [])
        review_data = review_metadata_dict.get(identifier, [])

        commit_stats = calculate_commit_stats_from_metadata(commit_metadata)
        pr_stats = calculate_pr_stats_from_metadata(pr_metadata)
        review_stats = calculate_review_stats_from_metadata(review_data)

        cache_dict[identifier] = {
            'name': agent_name,
            'website': assistant.get('website', 'N/A'),
            'github_identifier': identifier,
            **commit_stats,
            **pr_stats,
            **review_stats
        }

    return cache_dict


def save_leaderboard_data_to_hf(leaderboard_dict, commit_monthly_metrics, pr_monthly_metrics, review_monthly_metrics):
    """Save leaderboard data and all monthly metrics to HuggingFace dataset."""
    try:
        token = get_hf_token()
        if not token:
            raise Exception("No HuggingFace token found")

        api = HfApi(token=token)

        combined_data = {
            'metadata': {
                'last_updated': datetime.now(timezone.utc).isoformat(),
                'leaderboard_time_frame_days': LEADERBOARD_TIME_FRAME_DAYS
            },
            'leaderboard': leaderboard_dict,
            'commit_monthly_metrics': commit_monthly_metrics,
            'pr_monthly_metrics': pr_monthly_metrics,
            'review_monthly_metrics': review_monthly_metrics
        }

        with open(LEADERBOARD_FILENAME, 'w') as f:
            json.dump(combined_data, f, indent=2)

        try:
            upload_file_with_backoff(
                api=api,
                path_or_fileobj=LEADERBOARD_FILENAME,
                path_in_repo=LEADERBOARD_FILENAME,
                repo_id=LEADERBOARD_REPO,
                repo_type="dataset"
            )
            return True
        finally:
            if os.path.exists(LEADERBOARD_FILENAME):
                os.remove(LEADERBOARD_FILENAME)

    except Exception as e:
        print(f"Error saving leaderboard data: {str(e)}")
        traceback.print_exc()
        return False


# =============================================================================
# MINING FUNCTION
# =============================================================================

def mine_all_agents():
    """
    Mine PR, review, and commit metadata for all assistants using STREAMING batch processing.
    Downloads GHArchive data, then uses BATCH-based DuckDB queries.
    """
    print(f"\n[1/4] Downloading GHArchive data...")

    if not download_all_gharchive_data():
        print("Warning: Download had errors, continuing with available data...")

    print(f"\n[2/4] Loading assistant metadata...")

    assistants = load_agents_from_hf()
    if not assistants:
        print("Error: No assistants found")
        return

    identifiers = [assistant['github_identifier'] for assistant in assistants if assistant.get('github_identifier')]
    if not identifiers:
        print("Error: No valid assistant identifiers found")
        return

    current_time = datetime.now(timezone.utc)
    end_date = current_time.replace(hour=0, minute=0, second=0, microsecond=0)
    start_date = end_date - timedelta(days=LEADERBOARD_TIME_FRAME_DAYS)

    try:
        conn = get_duckdb_connection()
    except Exception as e:
        print(f"Failed to initialize DuckDB connection: {str(e)}")
        return

    print(f"\n[3/4] Mining PR, review, and commit metadata ({len(identifiers)} assistants, {LEADERBOARD_TIME_FRAME_DAYS} days)...")

    try:
        results = fetch_all_metadata_streaming(
            conn, identifiers, start_date, end_date
        )
        commit_metadata = results['commits']
        pr_metadata = results['prs']
        review_metadata = results['reviews']
    except Exception as e:
        print(f"Error during metadata fetch: {str(e)}")
        traceback.print_exc()
        return
    finally:
        conn.close()

    print(f"\n[4/4] Saving leaderboard...")

    try:
        leaderboard_dict = construct_leaderboard_from_metadata(
            commit_metadata, pr_metadata, review_metadata, assistants
        )

        commit_monthly_metrics = calculate_monthly_metrics_by_agent_commits(
            commit_metadata, assistants
        )
        pr_monthly_metrics = calculate_monthly_metrics_by_agent_prs(
            pr_metadata, assistants
        )
        review_monthly_metrics = calculate_monthly_metrics_by_agent_reviews(
            review_metadata, assistants
        )

        save_leaderboard_data_to_hf(
            leaderboard_dict, commit_monthly_metrics, pr_monthly_metrics, review_monthly_metrics
        )
    except Exception as e:
        print(f"Error saving leaderboard: {str(e)}")
        traceback.print_exc()
    finally:
        if os.path.exists(DUCKDB_CACHE_FILE):
            try:
                os.remove(DUCKDB_CACHE_FILE)
                print(f"   ✓ Cache file removed: {DUCKDB_CACHE_FILE}")
            except Exception as e:
                print(f"   ⚠ Failed to remove cache file: {str(e)}")

# =============================================================================
# ENTRY POINT
# =============================================================================

if __name__ == "__main__":
    mine_all_agents()