#!/usr/bin/env python3
"""Install or update AnimeWeave in a private directory; never change global Python.
Requires uv: https://docs.astral.sh/uv/getting-started/installation/
The license is entered with getpass, sent only to the chosen origin, and never saved.
Rerun to update. Use --rollback to restore the previous installation without network.
"""

import argparse
import fcntl
import getpass
import hashlib
import json
import os
import platform
import re
import shutil
import subprocess
import sys
import tempfile
import urllib.error
import urllib.parse
import urllib.request
import uuid
import zipfile
from pathlib import Path

USER_AGENT = "AnimeWeave-Installer/0.1"


class NoRedirect(urllib.request.HTTPRedirectHandler):
    def redirect_request(self, req, fp, code, msg, headers, newurl):
        raise ValueError("Redirect refused: the license must stay on the original server.")


def active_release(root):
    active = root / "venv"
    if active.is_symlink():
        resolved = active.resolve()
        if resolved.parent != root / "releases" or not (resolved / "bin/animeweave").is_file():
            raise ValueError("Invalid managed installation link; existing files were preserved.")
        return resolved
    if active.exists():
        raise ValueError(
            "This is an older installation layout. Use a new --directory; keep the old installation until the new one is ready."
        )
    return None


def point_to(root, name, target):
    temporary = root / ("." + name + "-" + uuid.uuid4().hex)
    temporary.symlink_to(target.relative_to(root), target_is_directory=True)
    os.replace(temporary, root / name)


def activate(root, target):
    previous = active_release(root)
    if previous:
        point_to(root, "previous", previous)
    point_to(root, "venv", target)


def rollback(root):
    current = active_release(root)
    previous = root / "previous"
    target = previous.resolve()
    if (
        not current
        or not previous.is_symlink()
        or target.parent != root / "releases"
        or not (target / "bin/animeweave").is_file()
    ):
        raise ValueError("No previous managed installation is available to restore.")
    # Switching the active pointer first makes rollback recoverable even if interrupted.
    point_to(root, "venv", target)
    point_to(root, "previous", current)
    print("Restored CLI:", root / "venv/bin/animeweave")


def check_render_runtime(python):
    probe = [
        str(python),
        "-c",
        "import manim, torch, soundfile; from animeweave_engine._vendor.kokoro import KModel; from animeweave_engine.english import runtime_probe; runtime_probe()",
    ]
    try:
        subprocess.run(probe, check=True, capture_output=True, text=True, timeout=60)
    except subprocess.CalledProcessError as error:
        # A cached locally built wheel can retain incompatible Intel/ARM link flags.
        # Repair only the diagnosed Cairo problem, only inside this new environment.
        cairo_pc = Path("/opt/homebrew/opt/cairo/lib/pkgconfig/cairo.pc")
        if (
            platform.system() == "Darwin"
            and platform.machine() == "arm64"
            and cairo_pc.is_file()
            and "cairo" in (error.stderr or "")
        ):
            print(
                "Rebuilding the managed Cairo binding against the installed Apple Silicon libraries."
            )
            environment = dict(os.environ)
            environment["PKG_CONFIG_PATH"] = ":".join(
                [
                    str(cairo_pc.parent),
                    "/opt/homebrew/opt/pango/lib/pkgconfig",
                    "/opt/homebrew/lib/pkgconfig",
                    environment.get("PKG_CONFIG_PATH", ""),
                ]
            )
            subprocess.run(
                [
                    "uv",
                    "pip",
                    "install",
                    "--python",
                    str(python),
                    "--reinstall",
                    "--no-cache",
                    "--no-binary",
                    "pycairo",
                    "pycairo==1.29.0",
                ],
                check=True,
                env=environment,
            )
            subprocess.run(probe, check=True, timeout=60)
        else:
            print(error.stderr or "The native rendering runtime could not load.", file=sys.stderr)
            raise


def extract_dependency_lock(wheel, temporary_directory, with_render):
    profile = "render" if with_render else "base"
    member = "animeweave/guidance/" + profile + "-requirements.txt"
    with zipfile.ZipFile(wheel) as archive:
        try:
            info = archive.getinfo(member)
        except KeyError:
            raise ValueError(
                "This release does not contain a dependency lock; obtain a current release before installing."
            ) from None
        if not 0 < info.file_size <= 128 * 1024:
            raise ValueError("Invalid dependency lock size.")
        requirements = archive.read(info).decode("utf-8")
    target = temporary_directory / (profile + "-requirements.txt")
    target.write_text(requirements)
    return target


def install_release(root, name, data, with_render, manifest, server):
    previous = active_release(root)
    if previous and (previous / "animeweave-install.json").is_file():
        old = json.loads((previous / "animeweave-install.json").read_text())
        with_render = with_render or old.get("with_render") is True
    release = root / "releases" / uuid.uuid4().hex
    release.parent.mkdir(parents=True, exist_ok=True)
    # Each environment keeps its final path: moving a venv breaks its entrypoint shebangs.
    subprocess.run(["uv", "venv", "--python", "3.12", str(release)], check=True)
    python = release / "bin/python"
    try:
        with tempfile.TemporaryDirectory(prefix="animeweave-install-") as temp:
            wheel = Path(temp) / name
            wheel.write_bytes(data)
            package = str(wheel) + ("[render]" if with_render else "")
            requirements = extract_dependency_lock(wheel, Path(temp), with_render)
            subprocess.run(
                [
                    "uv",
                    "pip",
                    "install",
                    "--python",
                    str(python),
                    "--requirements",
                    str(requirements),
                    package,
                ],
                check=True,
            )
        executable = release / "bin/animeweave"
        subprocess.run([str(executable), "--version"], check=True)
        subprocess.run(
            [str(executable), "guide", "workflow"], check=True, stdout=subprocess.DEVNULL
        )
        if with_render:
            check_render_runtime(python)
        (release / "animeweave-install.json").write_text(
            json.dumps(
                {
                    "schema_version": 1,
                    "filename": name,
                    "sha256": manifest["sha256"],
                    "with_render": with_render,
                    "dependency_lock": "render" if with_render else "base",
                    "server": server,
                },
                indent=2,
            )
        )
        activate(root, release)
    except (OSError, ValueError, subprocess.SubprocessError, zipfile.BadZipFile):
        print(
            "Update did not activate. Your previous installation is unchanged. Failed environment retained:",
            release,
            file=sys.stderr,
        )
        raise
    return root / "venv/bin/animeweave"


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--server", help="AnimeWeave website origin (required for downloads)")
    parser.add_argument(
        "--directory", type=Path, default=Path.home() / ".local/share/animeweave/installation"
    )
    parser.add_argument(
        "--download-only", type=Path, help="Verify and save the package without installing"
    )
    parser.add_argument(
        "--with-render",
        action="store_true",
        help="Also install large rendering dependencies (system libraries required); preserved on later updates",
    )
    parser.add_argument(
        "--rollback",
        action="store_true",
        help="Restore the previous managed installation without network or a license",
    )
    args = parser.parse_args()
    if sys.platform not in ("darwin", "linux"):
        parser.error("This installer currently supports macOS and Linux only.")
    if args.rollback and (args.download_only or args.with_render):
        parser.error("--rollback cannot be combined with download or render installation options.")
    root = args.directory.expanduser().resolve()
    root.mkdir(parents=True, exist_ok=True, mode=0o700)
    # POSIX advisory lock is released automatically, including on process termination.
    with (root / ".install.lock").open("a") as lock:
        try:
            fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB)
        except BlockingIOError:
            parser.error("Another installer is using this directory. Wait for it to finish.")
        if args.rollback:
            rollback(root)
            return
        if not args.server:
            parser.error("--server is required for downloads.")
        origin = args.server.rstrip("/")
        url = urllib.parse.urlsplit(origin)
        if (
            not url.hostname
            or url.username
            or url.password
            or url.query
            or url.fragment
            or url.path
            or (
                url.scheme != "https"
                and not (url.scheme == "http" and url.hostname in ("127.0.0.1", "localhost", "::1"))
            )
        ):
            parser.error("Use an HTTPS origin, or HTTP loopback for local preview.")
        if not args.download_only:
            active_release(root)
            if not shutil.which("uv"):
                parser.error(
                    "uv is required. Install it using https://docs.astral.sh/uv/getting-started/installation/ then retry."
                )
        print("Download server:", origin)
        key = getpass.getpass("License (hidden): ").strip()
        if not re.fullmatch(r"awe_[a-f0-9]{64}", key):
            parser.error("Invalid license format.")
        opener = urllib.request.build_opener(NoRedirect())

        def fetch(path, limit):
            req = urllib.request.Request(
                origin + path,
                headers={"Authorization": "Bearer " + key, "User-Agent": USER_AGENT},
            )
            with opener.open(req, timeout=120) as response:
                data = response.read(limit + 1)
            if len(data) > limit:
                raise ValueError("Download exceeds the expected size.")
            return data

        manifest = json.loads(fetch("/api/releases/latest", 16384))
        name = manifest.get("filename", "")
        size = manifest.get("size")
        if (
            not re.fullmatch(r"animeweave-[A-Za-z0-9_.+-]+\.whl", name)
            or type(size) is not int
            or not 0 < size <= 100 * 1024 * 1024
            or not re.fullmatch("[a-f0-9]{64}", manifest.get("sha256", ""))
        ):
            raise ValueError("Invalid release manifest.")
        data = fetch("/api/releases/download", size)
        if len(data) != size or hashlib.sha256(data).hexdigest() != manifest["sha256"]:
            raise ValueError("Package verification failed; nothing was installed.")
        if args.download_only:
            args.download_only.mkdir(parents=True, exist_ok=True)
            target = args.download_only / name
            with target.open("xb") as out:
                out.write(data)
            print("Verified package:", target)
            return
        executable = install_release(root, name, data, args.with_render, manifest, origin)
        print("\nCLI installed:", executable)
        print("Add this directory to your shell PATH:", executable.parent)
        print("Checking dependencies; missing items are reported below.")
        result = subprocess.run([str(executable), "doctor"], check=False)
        if result.returncode:
            print("The CLI is installed, but rendering is not ready. Follow doctor output.")
            print(
                "After system prerequisites, rerun this installer with --with-render, then run animeweave doctor --fetch-model and animeweave doctor."
            )
        print(
            "To update, rerun this installer with the same --directory and --server. Render libraries remain enabled once selected."
        )
        print(
            "To restore the previous version, rerun with --directory and --rollback; no server or license is needed."
        )
        print(
            "Installed versions do not need periodic license renewal. Your agent still uses its own network connection."
        )


if __name__ == "__main__":
    try:
        main()
    except urllib.error.HTTPError as error:
        print(
            "Download failed (HTTP "
            + str(error.code)
            + "). Check your license and purchase status.",
            file=sys.stderr,
        )
        sys.exit(1)
    except (ValueError, OSError, subprocess.SubprocessError, zipfile.BadZipFile) as error:
        print("Setup stopped:", type(error).__name__, str(error), file=sys.stderr)
        sys.exit(1)