import os
import sys
import zipfile
import tarfile
import urllib.request
import json
import shutil
import traceback
import re

# Primary PyPI base URL and mirrors for China fallback
PYPI_BASES = [
    "https://pypi.org/pypi",                 # Official PyPI
    "https://pypi.tuna.tsinghua.edu.cn/pypi", # Tsinghua (TUNA)
    "https://mirrors.aliyun.com/pypi",        # Aliyun
]

def _parse_package_spec(package_spec):
    match = re.match(r"^([a-zA-Z0-9_\-\.]+)(.*)$", package_spec.strip())
    if match:
        return match.group(1), match.group(2).strip()
    return package_spec, ""

def _fetch_pypi_data(clean_name):
    """Attempt to fetch metadata JSON from PyPI and mirrors."""
    for base_url in PYPI_BASES:
        url = f"{base_url}/{clean_name}/json"
        try:
            print(f"Fetching metadata from: {base_url}…")
            req = urllib.request.Request(url, headers={'User-Agent': 'Mozilla/5.0'})
            with urllib.request.urlopen(req, timeout=10) as response:
                return json.loads(response.read().decode("utf-8"))
        except Exception as e:
            print(f"Failed to fetch metadata from {base_url}: {e}")
    return None

def _download_file(file_info_urls, dest_path):
    """Download package file with mirror URL rewriting if official PyPI fails."""
    # List of candidate download URLs (original URL + mirror URL replacements)
    for file_info in file_info_urls:
        original_url = file_info.get("url")
        if not original_url:
            continue

        # Generate candidate download URLs for this file
        candidate_urls = [original_url]

        # Rewrite official PyPI file URLs to mirror locations if needed
        if "files.pythonhosted.org" in original_url:
            candidate_urls.append(
                original_url.replace("https://files.pythonhosted.org", "https://pypi.tuna.tsinghua.edu.cn")
            )
            candidate_urls.append(
                original_url.replace("https://files.pythonhosted.org", "https://mirrors.aliyun.com/pypi")
            )

        for dl_url in candidate_urls:
            try:
                print(f"Downloading from {dl_url}…")
                dl_req = urllib.request.Request(dl_url, headers={'User-Agent': 'Mozilla/5.0'})
                with urllib.request.urlopen(dl_req, timeout=30) as dl_resp, open(dest_path, "wb") as out_file:
                    out_file.write(dl_resp.read())
                return True
            except Exception as e:
                print(f"Download failed from {dl_url}: {e}")

    return False

def install_package(package_name, target_dir):
    try:
        clean_name, version_spec = _parse_package_spec(package_name)
        print(f"Resolving PyPI metadata for {clean_name} (spec: '{version_spec}')…")

        data = _fetch_pypi_data(clean_name)
        if not data:
            print(f"Could not retrieve PyPI metadata for {clean_name} from any mirror.")
            return False

        urls = data.get("urls", [])

        if version_spec.startswith("=="):
            req_version = version_spec.lstrip("=").strip()
            releases = data.get("releases", {})
            if req_version in releases:
                urls = releases[req_version]
            else:
                print(f"Version {req_version} not found on PyPI for {clean_name}")
                return False

        target_file_infos = []

        # 1. Prefer pure-Python wheels
        for file_info in urls:
            filename = file_info.get("filename", "")
            if filename.endswith("py3-none-any.whl") or filename.endswith("py2.py3-none-any.whl"):
                target_file_infos.append(file_info)
                break

        # 2. Fallback to source distribution
        if not target_file_infos:
            for file_info in urls:
                if file_info.get("packagetype") == "sdist":
                    target_file_infos.append(file_info)
                    break

        if not target_file_infos:
            print(f"No compatible pure-Python release found for {clean_name}")
            return False

        if not os.path.exists(target_dir):
            os.makedirs(target_dir, exist_ok=True)

        safe_filename = clean_name.replace("/", "_")
        dest_path = os.path.join(target_dir, f"_dl_{safe_filename}")

        if not _download_file(target_file_infos, dest_path):
            print(f"Failed to download payload for {clean_name} from all sources.")
            return False

        # Extract package files
        if dest_path.endswith(".whl") or dest_path.endswith(".zip") or target_file_infos[0].get("filename", "").endswith((".whl", ".zip")):
            with zipfile.ZipFile(dest_path, "r") as zip_ref:
                zip_ref.extractall(target_dir)
        else:
            with tarfile.open(dest_path, "r:gz") as tar_ref:
                tar_ref.extractall(target_dir)

        if os.path.exists(dest_path):
            os.remove(dest_path)

        # Handle 'src' layout if present
        for root, dirs, files in os.walk(target_dir):
            if os.path.basename(root) == "src":
                for d in dirs:
                    src_pkg = os.path.join(root, d)
                    dest_pkg = os.path.join(target_dir, d)
                    if not os.path.exists(dest_pkg):
                        shutil.move(src_pkg, dest_pkg)

        # Cleanup __pycache__
        for root, dirs, files in os.walk(target_dir):
            for d in list(dirs):
                if d == "__pycache__":
                    shutil.rmtree(os.path.join(root, d), ignore_errors=True)

        if target_dir not in sys.path:
            sys.path.insert(0, target_dir)

        print(f"Successfully installed {clean_name}")
        return True

    except Exception as e:
        print(f"Failed to install {package_name}: {e}")
        traceback.print_exc()
        return False