"""Custom build backend that wraps maturin and adds .pth file support."""

# Claude wrote all the code in this file

import hashlib
import os
import shutil
from pathlib import Path
from typing import Any, Dict, Tuple, Union

from maturin import build_sdist as maturin_build_sdist
from maturin import build_wheel as maturin_build_wheel
from maturin import build_editable as maturin_build_editable
from maturin import (
    get_requires_for_build_sdist,
    get_requires_for_build_wheel,
    get_requires_for_build_editable as maturin_get_requires_for_build_editable,
    prepare_metadata_for_build_wheel,
    prepare_metadata_for_build_editable as maturin_prepare_metadata_for_build_editable,
)


def compute_file_hash(file_path: Path) -> Tuple[str, int]:
    """Compute SHA256 hash and size for a file.

    Returns:
        Tuple of (base64-encoded hash, file size in bytes)
    """
    import base64

    sha256_hash = hashlib.sha256()
    file_size = 0

    with open(file_path, "rb") as f:
        while chunk := f.read(8192):
            sha256_hash.update(chunk)
            file_size += len(chunk)

    hash_digest = base64.urlsafe_b64encode(sha256_hash.digest()).rstrip(b"=").decode()
    return hash_digest, file_size


def build_wheel(
    wheel_directory: str,
    config_settings: Union[Dict[str, Any], None] = None,
    metadata_directory: Union[str, None] = None,
) -> str:
    """Build wheel using maturin, then add our .pth file."""
    # First build the wheel with maturin
    wheel_name = maturin_build_wheel(
        wheel_directory, config_settings, metadata_directory
    )

    # Now we need to add the .pth file to the wheel
    wheel_path = Path(wheel_directory) / wheel_name

    # Extract wheel, add file, rebuild
    import tempfile
    import zipfile

    with tempfile.TemporaryDirectory() as temp_dir:
        temp_path = Path(temp_dir)

        # Extract the wheel
        with zipfile.ZipFile(wheel_path, "r") as wheel_zip:
            wheel_zip.extractall(temp_path)

        # Add the kolo1.pth file to the .data/purelib directory (installs to site-packages)
        data_dir = temp_path / f"kolo-{get_version()}.data" / "purelib"
        data_dir.mkdir(parents=True, exist_ok=True)

        # Copy kolo1.pth to the purelib directory
        pth_source = Path("kolo1.pth")
        pth_added = False
        pth_dest = None
        if pth_source.exists():
            pth_dest = data_dir / "kolo1.pth"
            shutil.copy2(pth_source, pth_dest)
            pth_added = True

        # Update RECORD file to include kolo1.pth
        if pth_added and pth_dest:
            # Compute hash and size for kolo1.pth
            hash_digest, file_size = compute_file_hash(pth_dest)

            # Find the .dist-info directory
            dist_info_dirs = list(temp_path.glob("*.dist-info"))
            if dist_info_dirs:
                dist_info = dist_info_dirs[0]
                record_file = dist_info / "RECORD"
                if record_file.exists():
                    # Append kolo1.pth entry to RECORD
                    # Read existing content to check if it ends with newline
                    with open(record_file, "r") as f:
                        content = f.read()

                    # Append new entry with hash, ensuring proper line ending
                    with open(record_file, "a") as f:
                        # Add newline prefix only if file doesn't end with one
                        if content and not content.endswith("\n"):
                            f.write("\n")
                        f.write(
                            f"kolo-{get_version()}.data/purelib/kolo1.pth,"
                            f"sha256={hash_digest},{file_size}\n"
                        )

        # Rebuild the wheel
        # Use temporary file to avoid data loss if rebuild fails
        temp_wheel = wheel_path.with_suffix(".whl.tmp")
        with zipfile.ZipFile(temp_wheel, "w", zipfile.ZIP_DEFLATED) as wheel_zip:
            for file_path in temp_path.rglob("*"):
                if file_path.is_file():
                    arcname = str(file_path.relative_to(temp_path))
                    wheel_zip.write(file_path, arcname)
        os.remove(wheel_path)
        temp_wheel.rename(wheel_path)

    return wheel_name


def get_version():
    """Extract version from pyproject.toml."""
    try:
        import tomllib
    except ImportError:
        # Python < 3.11
        import tomli as tomllib

    with open("pyproject.toml", "rb") as f:
        data = tomllib.load(f)
    return data["project"]["version"]


def build_editable(
    wheel_directory: str,
    config_settings: Union[Dict[str, Any], None] = None,
    metadata_directory: Union[str, None] = None,
) -> str:
    """Build an editable wheel using maturin, then add our .pth file."""
    # First build the editable wheel with maturin (this compiles Rust!)
    wheel_name = maturin_build_editable(
        wheel_directory, config_settings, metadata_directory
    )

    # Now we need to add the kolo1.pth file to the wheel for KOLO=1 auto-activation
    # IMPORTANT: maturin creates kolo.pth with the path to src/, we must NOT overwrite it!
    # We use kolo1.pth for our KOLO=1 activation code instead.
    wheel_path = Path(wheel_directory) / wheel_name

    import tempfile
    import zipfile

    with tempfile.TemporaryDirectory() as temp_dir:
        temp_path = Path(temp_dir)

        # Extract the wheel
        with zipfile.ZipFile(wheel_path, "r") as wheel_zip:
            wheel_zip.extractall(temp_path)

        # Copy kolo1.pth for the KOLO=1 auto-activation functionality
        pth_source = Path("kolo1.pth")
        pth_added = False
        pth_dest = None
        if pth_source.exists():
            pth_dest = temp_path / "kolo1.pth"
            shutil.copy2(pth_source, pth_dest)
            pth_added = True

        # Update RECORD file to include kolo1.pth
        if pth_added and pth_dest:
            # Compute hash and size for kolo1.pth
            hash_digest, file_size = compute_file_hash(pth_dest)

            # Find the .dist-info directory
            dist_info_dirs = list(temp_path.glob("*.dist-info"))
            if dist_info_dirs:
                dist_info = dist_info_dirs[0]
                record_file = dist_info / "RECORD"
                if record_file.exists():
                    # Append kolo1.pth entry to RECORD
                    with open(record_file, "r") as f:
                        content = f.read()

                    with open(record_file, "a") as f:
                        if content and not content.endswith("\n"):
                            f.write("\n")
                        f.write(f"kolo1.pth,sha256={hash_digest},{file_size}\n")

        # Rebuild the wheel
        temp_wheel = wheel_path.with_suffix(".whl.tmp")
        with zipfile.ZipFile(temp_wheel, "w", zipfile.ZIP_DEFLATED) as wheel_zip:
            for file_path in temp_path.rglob("*"):
                if file_path.is_file():
                    arcname = str(file_path.relative_to(temp_path))
                    wheel_zip.write(file_path, arcname)
        os.remove(wheel_path)
        temp_wheel.rename(wheel_path)

    return wheel_name


def get_requires_for_build_editable(config_settings=None):
    """Return requirements for building an editable install."""
    return maturin_get_requires_for_build_editable(config_settings)


def prepare_metadata_for_build_editable(
    metadata_directory: str,
    config_settings: Union[Dict[str, Any], None] = None,
) -> str:
    """Prepare metadata for editable install."""
    return maturin_prepare_metadata_for_build_editable(
        metadata_directory, config_settings
    )


def build_sdist(
    sdist_directory: str,
    config_settings: Union[Dict[str, Any], None] = None,
) -> str:
    """Build sdist using maturin, then add our .pth file."""
    # First build the sdist with maturin
    sdist_name = maturin_build_sdist(sdist_directory, config_settings)

    # Now we need to add the .pth file to the sdist
    import tarfile
    import tempfile

    sdist_path = Path(sdist_directory) / sdist_name

    # Extract tarball, add file, rebuild
    with tempfile.TemporaryDirectory() as temp_dir:
        temp_path = Path(temp_dir)

        # Extract the sdist
        with tarfile.open(sdist_path, "r:gz") as tar:
            tar.extractall(temp_path)

        # Find the extracted directory (should be kolo-{version})
        extracted_dirs = [d for d in temp_path.iterdir() if d.is_dir()]
        if not extracted_dirs:
            raise RuntimeError("No directory found in sdist")
        if len(extracted_dirs) > 1:
            raise RuntimeError(f"Multiple directories found in sdist: {extracted_dirs}")
        extracted_dir = extracted_dirs[0]

        # Copy kolo1.pth to the extracted directory
        pth_source = Path("kolo1.pth")
        if pth_source.exists():
            shutil.copy2(pth_source, extracted_dir / "kolo1.pth")

        # Copy _build_backend.py to the extracted directory (needed for building from sdist)
        backend_source = Path("_build_backend.py")
        if backend_source.exists():
            shutil.copy2(backend_source, extracted_dir / "_build_backend.py")

        # Rebuild the sdist
        # Use temporary file to avoid data loss if rebuild fails
        temp_sdist = sdist_path.with_suffix(".tar.gz.tmp")
        with tarfile.open(temp_sdist, "w:gz") as tar:
            tar.add(extracted_dir, arcname=extracted_dir.name)
        os.remove(sdist_path)
        temp_sdist.rename(sdist_path)

    return sdist_name


# Re-export other required functions from maturin
__all__ = [
    "build_wheel",
    "build_sdist",
    "build_editable",
    "get_requires_for_build_wheel",
    "get_requires_for_build_sdist",
    "get_requires_for_build_editable",
    "prepare_metadata_for_build_wheel",
    "prepare_metadata_for_build_editable",
]
