import json
import pathlib
import re
import tarfile
import zipfile
import requests
import argparse
from concurrent.futures import ThreadPoolExecutor, as_completed
from tqdm import tqdm

pattern = r"(cpython-\d+)"


def get_source_from_cache(filepath: str) -> str:
    """Guess the source based on Python specification."""
    # remove pycache path
    filepath = filepath.replace("__pycache__/", "")
    # remove cpython and pytest versions and add 'py' as suffix
    filepath = filepath.split("cpython")[0] + "py"

    return filepath


def identify_caches(memberlist: list[str]) -> tuple[list[dict[str, str]], bool]:
    """Find cache files and identify caches without source."""
    cache_files = []
    phantom_caches = False
    for member in memberlist:
        if "__pycache__/" in str(member):
            source = get_source_from_cache(member)
            if source not in memberlist:
                phantom_caches = True
            cache_files.append(
                {
                    "source": source if source in memberlist else "",
                    "cache": member,
                    "version": re.search(pattern, member).group(0),
                }
            )
    return cache_files, phantom_caches


def get_file_list(file_location: pathlib.Path) -> list[str]:
    """Get the name of files within an artifact."""
    if file_location.suffix == ".gz":
        with tarfile.open(file_location, "r:gz") as f:
            return f.getnames()
    if file_location.suffix in [".whl", ".egg", ".zip"]:
        with zipfile.ZipFile(file_location) as f:
            return f.namelist()
    return []


def suffix_to_dist(filename: pathlib.Path) -> str:
    """Translate filename suffices to artifact types."""
    match filename.suffix:
        case ".gz":
            return "sdist"
        case ".whl":
            return "wheel"
        case ".zip":
            return "zip"
        case ".egg":
            return "egg"
        case _:
            return "other"


def scan_package(package: pathlib.Path, storage: pathlib.Path) -> dict:
    """Perform the large scale scan on all packages."""

    url = f"https://pypi.org/pypi/{package.name.split('-')[0]}/json"

    # Define the base structure
    results = {
        package.name: {
            "meta": {
                "base_path": str(package.absolute()),
                "api_url": url,
                "analysis": {
                    dist: {"status": "not run", "success": False}
                    for dist in ["api", "sdist", "wheel", "egg", "zip", "other"]
                },
            }
        }
    }

    # Shortcuts for cleaner access
    pkg_results = results[package.name]
    analysis_meta = pkg_results["meta"]["analysis"]

    # 1. Fetch Metadata from PyPI
    try:
        response = requests.get(url, timeout=10)
        response.raise_for_status()
        meta_data = response.json()

        latest_version = meta_data["info"]["version"]
        pkg_results["meta"]["version"] = latest_version
        analysis_meta["api"].update({"success": True, "status": "success"})
    except Exception as e:
        analysis_meta["api"].update({"success": False, "status": str(e)})
        return results

    # 2. Iterate over distribution artifacts
    releases = meta_data.get("releases", {}).get(latest_version, [])

    for release in releases:
        filename = pathlib.Path(release["filename"])
        dist_type = suffix_to_dist(filename)
        meta = analysis_meta[dist_type]

        if dist_type == "other":
            meta.update(
                {"status": f"Unsupported suffix: {filename.suffix}", "success": False}
            )
            continue

        file_location = storage / filename
        pkg_results[dist_type] = {
            "filepath": str(file_location.absolute()),
            "files": [],
        }

        try:
            # Get list of cache files and their source files
            file_list = get_file_list(file_location)
            cache_files, has_phantom = identify_caches(file_list)

            pkg_results[dist_type]["files"] = cache_files
            meta.update(
                {
                    "status": f"found {len(cache_files)} cache files",
                    "found": len(cache_files),
                    "success": True,
                    "phantom_caches": has_phantom,
                }
            )
        except Exception as e:
            # Record the reasong if it failed
            meta.update(
                {
                    "status": str(e),
                    "success": False,
                    "found": None,
                    "phantom_caches": None,
                }
            )

    return results


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Scan Python packages for `__pycache__` files."
    )

    parser.add_argument(
        "storage_path",
        type=pathlib.Path,
        help="The path to the directory containing the Python packages.",
    )

    parser.add_argument(
        "--workers",
        type=int,
        default=8,
        help="The number of concurrent threads to use for scanning. (default: 8)",
    )

    args = parser.parse_args()

    storage_path = pathlib.Path(args.storage_path)

    if not storage_path.is_dir():
        print(f"Error: The provided path '{storage_path}' is not a valid directory.")
        exit(1)

    packages = list(storage_path.iterdir())
    results = {}

    with ThreadPoolExecutor(max_workers=args.workers) as executor:
        futures = {
            executor.submit(scan_package, package, storage_path): package
            for package in packages
        }

        for future in tqdm(
            as_completed(futures), total=len(packages), desc="Scanning packages"
        ):
            result = future.result()
            results.update(result)

    with open("results.json", "w") as f:
        json.dump(results, f, indent=2)