import os
import requests
import zipfile
import shutil
from urllib.parse import urlparse


def _repo_zip_url(repo: str, branch: str) -> str:
    """Return the GitHub archive zip URL for a repo and branch.

    Repo may be provided as 'owner/repo' or a full URL.
    """
    if repo.startswith("http"):
        parsed = urlparse(repo)
        path = parsed.path.strip("/")
        repo_path = path
    else:
        repo_path = repo
    return f"https://github.com/{repo_path}/archive/{branch}.zip"


def fetch_tests(
    branch: str, output_dir: str = None, repo: str = "bayesflow-org/bayesflow"
) -> str:
    """Fetch tests from a specific branch of a GitHub repo and extract them into a unique directory.

    Returns the path to the extracted directory.
    - repo: 'owner/repo' or full URL to a GitHub repo
    - output_dir: if None a temp directory name 'tests_<branch>' will be used
    """
    if output_dir is None:
        output_dir = f"tests_{branch}"

    # avoid clobbering an existing folder by creating a unique directory
    base_output = output_dir
    counter = 0
    while os.path.exists(output_dir):
        counter += 1
        output_dir = f"{base_output}_{counter}"

    # Try the requested branch, and fall back to common alternatives if 404
    # Always try the requested branch first, then 'main', then 'master'
    branches_to_try = []
    for b in (branch, "main", "master"):
        if b and b not in branches_to_try:
            branches_to_try.append(b)

    os.makedirs(output_dir, exist_ok=True)
    zip_path = os.path.join(output_dir + ".zip")

    response = None
    last_exc = None
    for br in branches_to_try:
        url = _repo_zip_url(repo, br)
        try:
            resp = requests.get(url, stream=True)
            resp.raise_for_status()
            response = resp
            # update zip_path name to include the actual branch tried
            zip_path = os.path.join(output_dir + f"_{br}.zip")
            break
        except Exception as exc:
            last_exc = exc
            # try next branch
            continue

    if response is None:
        # Try git-clone fallback (shallow clone of the branch)
        try:
            import subprocess

            git_url = (
                repo if repo.startswith("http") else f"https://github.com/{repo}.git"
            )
            # try cloning the requested branch
            for b in branches_to_try:
                try:
                    subprocess.run(
                        [
                            "git",
                            "clone",
                            "--depth",
                            "1",
                            "--branch",
                            b,
                            git_url,
                            output_dir,
                        ],
                        check=True,
                    )
                    # success
                    return output_dir
                except Exception:
                    continue
        except Exception:
            pass

        # Provide a clearer error message including branches attempted
        tried = ", ".join(branches_to_try)
        raise RuntimeError(
            f"Failed to download archive for repo={repo}; tried branches: {tried}. Last error: {last_exc}"
        )

    # Download the zip file content to disk
    with open(zip_path, "wb") as f:
        for chunk in response.iter_content(chunk_size=8192):
            f.write(chunk)

    # Extract the zip file into the output_dir
    with zipfile.ZipFile(zip_path, "r") as zip_ref:
        # Extract into a temporary folder first
        zip_ref.extractall(output_dir)

    # If the archive extracted into a single top-level folder (owner-repo-branch), move contents up
    entries = os.listdir(output_dir)
    if len(entries) == 1 and os.path.isdir(os.path.join(output_dir, entries[0])):
        inner = os.path.join(output_dir, entries[0])
        for name in os.listdir(inner):
            src = os.path.join(inner, name)
            dst = os.path.join(output_dir, name)
            shutil.move(src, dst)
        shutil.rmtree(inner)

    # Clean up zip
    try:
        os.remove(zip_path)
    except OSError:
        pass

    return output_dir
