# Copyright (c) 2024 Bytedance Ltd. and/or its affiliates

#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at

#      http://www.apache.org/licenses/LICENSE-2.0

#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.

import os
import csv
import argparse
import shutil
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
from typing import List, Tuple, Dict, Union

from multi_swe_bench.collect.util import get_tokens, optional_int
from multi_swe_bench.collect.get_pipeline import run_pipeline


def process_repository(
    out_base_dir: str,
    org: str,
    repo: str,
    token: str,
    delay_on_error: Union[int, None],
    retry_attempts: int,
    skip_commit_message: bool,
) -> bool:
    """Process a single repository using the get_pipeline logic directly"""
    dir_name = f"{org}__{repo}".replace("/", "__")
    out_dir = os.path.join(out_base_dir, dir_name)
    os.makedirs(out_dir, exist_ok=True)

    try:
        print(f"🔄 Starting: {org}/{repo}")
        run_pipeline(
            out_dir=Path(out_dir),
            tokens=[token],
            org=org,
            repo=repo,
            delay_on_error=delay_on_error,
            retry_attempts=retry_attempts,
            skip_commit_message=skip_commit_message,
        )
        print(f"✅ Success: {org}/{repo}")
        return True
    except Exception as e:
        print(f"❌ Failed: {org}/{repo}\nError: {str(e)}")
        if os.path.exists(out_dir) and not os.listdir(out_dir):
            shutil.rmtree(out_dir)
        return False


def process_repository_wrapper(args: Tuple[str, str, str, str]) -> bool:
    """Wrapper for ThreadPoolExecutor"""
    return process_repository(*args)


def read_repositories_from_csv(csv_file: str) -> List[Tuple[str, str]]:
    """Read repositories from CSV with validation"""
    repositories: List[Tuple[str, str]] = []
    with open(csv_file, mode="r", encoding="utf-8") as file:
        reader = csv.DictReader(file)
        for row in reader:
            if "Name" not in row:
                raise ValueError("CSV must contain 'Name' column")
            full_name = row["Name"]
            if "/" not in full_name:
                print(f"Skipping invalid repository name: {full_name}")
                continue
            org, repo = full_name.split("/", 1)
            repositories.append((org, repo))
    return repositories


def distribute_repositories(
    repositories: List[Tuple[str, str]],
    tokens: List[str],
    strategy: str = "round",
) -> List[Dict[str, Union[str, List[Tuple[str, str]]]]]:
    """
    Distribute repositories among tokens using specified strategy
    """
    if strategy == "chunk":
        return _distribute_chunk(repositories, tokens)
    return _distribute_round(repositories, tokens)


def _distribute_chunk(
    repositories: List[Tuple[str, str]], tokens: List[str]
) -> List[Dict[str, Union[str, List[Tuple[str, str]]]]]:
    """Equal chunks distribution"""
    num_repos = len(repositories)
    num_tokens = len(tokens)
    base_chunk_size = num_repos // num_tokens
    remainder = num_repos % num_tokens

    distributed: List[Dict[str, Union[str, List[Tuple[str, str]]]]] = []
    position = 0
    for i in range(num_tokens):
        chunk_size = base_chunk_size + (1 if i < remainder else 0)
        distributed.append(
            {
                "token": tokens[i],
                "repos": repositories[position : position + chunk_size],
            }
        )
        position += chunk_size
    return distributed


def _distribute_round(
    repositories: List[Tuple[str, str]], tokens: List[str]
) -> List[Dict[str, Union[str, List[Tuple[str, str]]]]]:
    """Round-robin distribution"""
    num_tokens = len(tokens)
    token_repos: Dict[str, List[Tuple[str, str]]] = {token: [] for token in tokens}

    for i, repo in enumerate(repositories):
        token_index = i % num_tokens
        token_repos[tokens[token_index]].append(repo)

    return [{"token": token, "repos": repos} for token, repos in token_repos.items()]


def process_token_group(
    out_dir: str,
    repos: List[Tuple[str, str]],
    token: str,
    delay_on_error: Union[int, None],
    retry_attempts: int,
    skip_commit_message: bool,
) -> List[bool]:
    """Process repos of a token group"""
    results: List[bool] = []
    for org, repo in repos:
        results.append(
            process_repository(
                out_dir,
                org,
                repo,
                token,
                delay_on_error,
                retry_attempts,
                skip_commit_message,
            )
        )
    return results


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Concurrently process GitHub repositories from CSV",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument(
        "--csv_file", required=True, help='Path to CSV file with "Name" column'
    )
    parser.add_argument(
        "--out_dir", required=True, help="Base output directory for repository data"
    )
    parser.add_argument(
        "--max_workers",
        type=int,
        default=None,
        help="Max concurrent processes (default: token count)",
    )
    parser.add_argument(
        "--distribute",
        choices=["round", "chunk"],
        default="round",
        help="Distribution strategy: round-robin (default) or equal chunks",
    )
    parser.add_argument(
        "--tokens",
        type=str,
        nargs="*",
        default=None,
        help="API token(s) or path to token file.",
    )
    parser.add_argument(
        "--delay-on-error",
        type=optional_int,
        default=300,
        help="Delay in seconds before retrying on error. If none, exit on error.",
    )
    parser.add_argument(
        "--retry-attempts",
        type=int,
        default=3,
        help="Number of attempts to retry on error.",
    )
    parser.add_argument(
        "--skip-commit-message",
        type=bool,
        default=False,
        help="Skip commit message.",
    )
    args = parser.parse_args()

    try:
        tokens: List[str] = get_tokens(args.tokens)
        print(f"Loaded {len(tokens)} GitHub tokens")

        repositories = read_repositories_from_csv(args.csv_file)
        if not repositories:
            print("No valid repositories found in CSV")
            return

        print(f"Processing {len(repositories)} repositories...")
        os.makedirs(args.out_dir, exist_ok=True)

        distribution = distribute_repositories(
            repositories, tokens, strategy=args.distribute
        )
        for i, group in enumerate(distribution, 1):
            print(f"Token {i} will process {len(group['repos'])} repositories")

        max_workers: int = args.max_workers or len(tokens)
        with ThreadPoolExecutor(max_workers=max_workers) as executor:
            tasks = []
            for group in distribution:
                tasks.append(
                    executor.submit(
                        process_token_group,
                        args.out_dir,
                        group["repos"],  # type: ignore
                        group["token"],  # type: ignore
                        args.delay_on_error,
                        args.retry_attempts,
                        args.skip_commit_message,
                    )
                )

            completed = 0
            for future in tasks:
                future.result()
                completed += 1
                print(
                    f"Progress: {completed}/{len(tasks)} ({completed / len(tasks):.1%})"
                )

        print("🎉 All repositories processed!")
    except Exception as e:
        print(f"💥 Error: {str(e)}")
        exit(1)


if __name__ == "__main__":
    main()
