import os
import sys
import shutil
import subprocess
import tempfile
import zipfile
from pathlib import Path
from setuptools import setup, Extension
from Cython.Build import cythonize

# Dateien oder Ordner, die nicht mit Cython kompiliert werden sollen
# (werden stattdessen direkt an PyArmor übergeben)
EXCLUDE_FROM_CYTHON = {
    "manage.py",
    "wsgi.py",
    "asgi.py",
}

def is_django_migration(path: Path) -> bool:
    """Prüft, ob eine Datei eine Django-Migrationsdatei ist."""
    return "migrations" in path.parts and path.name != "__init__.py"

def prepare_source(source_input: str) -> Path:
    """Entpackt eine Zip-Datei oder kopiert den Quellordner in ein Arbeitsverzeichnis."""
    src_path = Path(source_input).resolve()
    temp_dir = Path(tempfile.mkdtemp(prefix="build_src_"))

    if src_path.is_file() and src_path.suffix.lower() == ".zip":
        print(f"[*] Entpacke Zip-Archiv '{src_path.name}' nach temporärem Pfad...")
        with zipfile.ZipFile(src_path, 'r') as zip_ref:
            zip_ref.extractall(temp_dir)
        # Falls das Zip einen einzelnen Hauptordner enthält
        entries = [p for p in temp_dir.iterdir() if p.is_dir() and not p.name.startswith(".")]
        if len(entries) == 1:
            return entries[0]
        return temp_dir
    elif src_path.is_dir():
        print(f"[*] Kopiere Quellverzeichnis (Read-Only-Schutz der Originaldaten)...")
        work_dir = temp_dir / src_path.name
        shutil.copytree(src_path, work_dir, symlinks=True, ignore=shutil.ignore_patterns("dist", ".git", ".venv*", "__pycache__"))
        return work_dir
    else:
        raise ValueError(f"Ungültiger Quellpfad: {source_input}")

def run_build(project_dir: Path, output_dist: Path):
    if output_dist.exists():
        shutil.rmtree(output_dist)
    output_dist.mkdir(parents=True, exist_ok=True)

    py_files = []
    cython_candidates = []
    pyarmor_candidates = []

    # 1. Alle Dateien scannen
    for root, dirs, files in os.walk(project_dir):
        # Cache- und Build-Ordner überspringen
        dirs[:] = [d for d in dirs if d not in ("__pycache__", ".git", "dist", "build")]
        for file in files:
            file_path = Path(root) / file
            rel_path = file_path.relative_to(project_dir)

            if file.endswith(".py"):
                # manage.py und Migrationen direkt zu PyArmor leiten
                if file in EXCLUDE_FROM_CYTHON or is_django_migration(rel_path):
                    pyarmor_candidates.append(rel_path)
                else:
                    cython_candidates.append(rel_path)
            else:
                # Nicht-Python-Dateien (HTML, CSS, JSON etc.) direkt nach dist kopieren
                target_file = output_dist / rel_path
                target_file.parent.mkdir(parents=True, exist_ok=True)
                shutil.copy2(file_path, target_file)

    print(f"[*] Gefundene Python-Dateien: {len(cython_candidates) + len(pyarmor_candidates)}")
    print(f"    - Cython-Kandidaten: {len(cython_candidates)}")
    print(f"    - Direkte PyArmor-Ziele: {len(pyarmor_candidates)}")

    # 2. Cython-Kompilierung durchführen
    compiled_success = []
    failed_cython = []

    build_temp = Path(tempfile.mkdtemp(prefix="cython_build_"))
    original_cwd = os.getcwd()

    try:
        os.chdir(project_dir)
        for rel_file in cython_candidates:
            ext_name = str(rel_file.with_suffix("")).replace(os.sep, ".")
            try:
                ext = Extension(ext_name, [str(rel_file)])
                ext_modules = cythonize(
                    [ext],
                    compiler_directives={
                        "language_level": "3",
                        "always_allow_keywords": True,
                        "annotation_typing": False, # Verhinder Typkonflikte
                    },
                    quiet=True,
                )
                setup(
                    name=ext_name,
                    ext_modules=ext_modules,
                    script_args=["build_ext", "--inplace", f"--build-temp={build_temp}"],
                )
                compiled_success.append(rel_file)
            except Exception as err:
                print(f"[!] Cython-Fehler bei '{rel_file}', fällt auf PyArmor zurück ({err})")
                failed_cython.append(rel_file)
    finally:
        os.chdir(original_cwd)
        if build_temp.exists():
            shutil.rmtree(build_temp, ignore_errors=True)

    # 3. Kompilierte Binärdateien (.so / .pyd) nach dist/ übertragen
    for rel_file in compiled_success:
        source_folder = project_dir / rel_file.parent
        target_folder = output_dist / rel_file.parent
        target_folder.mkdir(parents=True, exist_ok=True)

        base_name = rel_file.stem
        binaries = list(source_folder.glob(f"{base_name}.*.so")) + list(source_folder.glob(f"{base_name}.*.pyd"))
        for b in binaries:
            shutil.move(str(b), str(target_folder / b.name))

    # 4. Verbleibende Dateien (Fehlgeschlagene + Ausnahmen) via PyArmor schützen
    all_pyarmor_files = pyarmor_candidates + failed_cython
    if all_pyarmor_files:
        print(f"[*] Verschlüssele {len(all_pyarmor_files)} Dateien mit PyArmor...")
        for rel_file in all_pyarmor_files:
            src_file = project_dir / rel_file
            target_file = output_dist / rel_file
            target_file.parent.mkdir(parents=True, exist_ok=True)

            # PyArmor gen Befehl
            cmd = [
                sys.executable, "-m", "pyarmor.cli", "gen",
                "-O", str(target_file.parent),
                str(src_file)
            ]
            res = subprocess.run(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.PIPE, text=True)
            if res.returncode != 0:
                # Falls PyArmor fehlschlägt, Quelltext als Fallback kopieren
                print(f"[!] PyArmor Warnung für {rel_file}: {res.stderr.strip()[:100]} - Kopiere Originaldatei")
                shutil.copy2(src_file, target_file)

    print(f"\n[+] Fertig! Das geschützte Projekt befindet sich in: {output_dist.resolve()}")

def main():
    if len(sys.argv) < 2:
        print("Verwendung: python build_dist.py <Pfad-zum-Projekt-oder-Zip-Datei>")
        sys.exit(1)

    input_source = sys.argv[1]
    work_dir = None
    try:
        work_dir = prepare_source(input_source)
        dist_dir = Path(input_source).parent / "dist" if Path(input_source).is_file() else Path(input_source) / "dist"
        run_build(work_dir, dist_dir)
    finally:
        # Temporäres Arbeitsverzeichnis bereinigen
        if work_dir and work_dir.parent.name.startswith("build_src_"):
            shutil.rmtree(work_dir.parent, ignore_errors=True)

if __name__ == "__main__":
    main()