#!/bin/bash
# Krasis installer — downloads pre-built wheels from GitHub releases.
# No PyPI, no pipx, no sudo. Works out of the box.
#
# Install / upgrade (latest stable):
#   curl -sSf https://raw.githubusercontent.com/brontoguana/krasis/main/install.sh | bash
#   krasis update
#
# Install latest pre-release:
#   curl -sSf https://raw.githubusercontent.com/brontoguana/krasis/main/install.sh | bash -s -- prerelease
#   krasis prerelease
#
# Uninstall:
#   curl -sSf https://raw.githubusercontent.com/brontoguana/krasis/main/install.sh | bash -s -- --uninstall
set -euo pipefail

# Wrap everything in a function so bash reads the entire script from the pipe
# before executing anything. Without this, subprocesses can consume stdin
# (which IS the pipe) and steal the rest of the script, causing silent exit.
do_install() {

trap 'echo -e "\n\033[0;31m\033[1mInstall failed\033[0m at line $LINENO. Run with \"bash -x\" for details." >&2' ERR

REPO="brontoguana/krasis"
VENV_DIR="$HOME/.krasis/venv"
BIN_DIR="$HOME/.local/bin"
COMMANDS="krasis krasis-chat krasis-setup"
OBSOLETE_VENV_PACKAGES=(
    # Stale from the old SGLang/KTransformers path; Krasis no longer uses it.
    sglang
    # Not a current Krasis install dependency. Old copies constrain fsspec and
    # create unrelated pip resolver warnings during Krasis upgrades.
    datasets
)

BOLD="\033[1m"
DIM="\033[2m"
RED="\033[0;31m"
GREEN="\033[0;32m"
YELLOW="\033[1;33m"
CYAN="\033[0;36m"
NC="\033[0m"

info()  { echo -e "${CYAN}${BOLD}=>${NC} $1"; }
ok()    { echo -e "${GREEN}${BOLD}OK${NC} $1"; }
warn()  { echo -e "${YELLOW}${BOLD}!!${NC} $1"; }
err()   { echo -e "${RED}${BOLD}ERROR${NC} $1"; exit 1; }

cleanup_obsolete_venv_packages() {
    [[ -x "$VENV_DIR/bin/python" ]] || return 0
    [[ ${#OBSOLETE_VENV_PACKAGES[@]} -gt 0 ]] || return 0

    local stale=()
    local package
    for package in "${OBSOLETE_VENV_PACKAGES[@]}"; do
        if "$VENV_DIR/bin/python" -m pip show "$package" &>/dev/null; then
            stale+=("$package")
        fi
    done

    if [[ ${#stale[@]} -gt 0 ]]; then
        info "Removing obsolete packages from Krasis-managed venv: ${stale[*]}"
        "$VENV_DIR/bin/python" -m pip uninstall -y "${stale[@]}" >/dev/null \
            || err "Failed to remove obsolete packages from $VENV_DIR."
    fi
}

# ── Channel (stable / prerelease) ────────────────────────────────────
CHANNEL="stable"
if [[ "${1:-}" == "prerelease" || "${1:-}" == "--prerelease" ]]; then
    CHANNEL="prerelease"
    shift
fi

# ── Uninstall ──────────────────────────────────────────────────────────
if [[ "${1:-}" == "--uninstall" ]]; then
    info "Uninstalling Krasis..."
    for cmd in $COMMANDS; do
        rm -f "$BIN_DIR/$cmd"
    done
    rm -rf "$VENV_DIR"
    ok "Krasis removed. Model files in ~/.krasis/models/ were kept."
    exit 0
fi

# ── Platform check ────────────────────────────────────────────────────
[[ "$(uname -s)" == "Linux" ]] || err "Krasis only supports Linux. On Windows, use WSL2."

# ── Check curl ────────────────────────────────────────────────────────
command -v curl &>/dev/null || err "curl is required but not found.\n  sudo apt install curl"

# ── Find all available Python 3.10+ ──────────────────────────────────
# Order matches our build matrix. Newest first so users get the best
# available wheel.  python3 is last as a catch-all — its version is
# checked, so it won't match 3.9 or older.
N_PY=0
PY_BINS=()
PY_VERS=()
PY_DOTS=()
SEEN_VERS=""

for candidate in python3.13 python3.12 python3.11 python3.10 python3; do
    command -v "$candidate" &>/dev/null || continue
    ver=$("$candidate" -c "import sys; print(f'{sys.version_info.major}.{sys.version_info.minor}')" 2>/dev/null) || continue
    major="${ver%%.*}"; minor="${ver##*.}"
    [[ "$major" -eq 3 && "$minor" -ge 10 ]] || continue
    pyver="${major}${minor}"
    # Deduplicate (python3 may resolve to same as python3.12)
    [[ "$SEEN_VERS" == *" $pyver "* ]] && continue
    SEEN_VERS="$SEEN_VERS $pyver "
    PY_BINS[$N_PY]="$candidate"
    PY_VERS[$N_PY]="$pyver"
    PY_DOTS[$N_PY]="$ver"
    N_PY=$((N_PY + 1))
done

if [[ $N_PY -eq 0 ]]; then
    warn "Python 3.10+ not found."
    if command -v apt-get &>/dev/null; then
        echo -en "  Install Python 3 now? [Y/n] "
        read -r answer < /dev/tty 2>/dev/null || answer="y"
        if [[ "${answer,,}" =~ ^(y|yes|)$ ]]; then
            SUDO=()
            [[ "$(id -u)" -ne 0 ]] && SUDO=(sudo)
            info "Installing Python 3..."
            "${SUDO[@]}" apt-get update -qq </dev/null 2>/dev/null
            "${SUDO[@]}" apt-get install -y python3 python3-venv </dev/null \
                || err "Failed to install Python 3.\n  Try manually: sudo apt install python3 python3-venv"
            # Re-scan for Python
            for candidate in python3.13 python3.12 python3.11 python3.10 python3; do
                command -v "$candidate" &>/dev/null || continue
                ver=$("$candidate" -c "import sys; print(f'{sys.version_info.major}.{sys.version_info.minor}')" 2>/dev/null) || continue
                major="${ver%%.*}"; minor="${ver##*.}"
                [[ "$major" -eq 3 && "$minor" -ge 10 ]] || continue
                pyver="${major}${minor}"
                [[ "$SEEN_VERS" == *" $pyver "* ]] && continue
                SEEN_VERS="$SEEN_VERS $pyver "
                PY_BINS[$N_PY]="$candidate"
                PY_VERS[$N_PY]="$pyver"
                PY_DOTS[$N_PY]="$ver"
                N_PY=$((N_PY + 1))
            done
            [[ $N_PY -eq 0 ]] && err "Python 3.10+ still not found after install."
        else
            err "Python 3.10+ is required.\n  Install it: sudo apt install python3"
        fi
    elif command -v dnf &>/dev/null; then
        echo -en "  Install Python 3 now? [Y/n] "
        read -r answer < /dev/tty 2>/dev/null || answer="y"
        if [[ "${answer,,}" =~ ^(y|yes|)$ ]]; then
            SUDO=()
            [[ "$(id -u)" -ne 0 ]] && SUDO=(sudo)
            info "Installing Python 3..."
            "${SUDO[@]}" dnf install -y python3 </dev/null \
                || err "Failed to install Python 3.\n  Try manually: sudo dnf install python3"
            for candidate in python3.13 python3.12 python3.11 python3.10 python3; do
                command -v "$candidate" &>/dev/null || continue
                ver=$("$candidate" -c "import sys; print(f'{sys.version_info.major}.{sys.version_info.minor}')" 2>/dev/null) || continue
                major="${ver%%.*}"; minor="${ver##*.}"
                [[ "$major" -eq 3 && "$minor" -ge 10 ]] || continue
                pyver="${major}${minor}"
                [[ "$SEEN_VERS" == *" $pyver "* ]] && continue
                SEEN_VERS="$SEEN_VERS $pyver "
                PY_BINS[$N_PY]="$candidate"
                PY_VERS[$N_PY]="$pyver"
                PY_DOTS[$N_PY]="$ver"
                N_PY=$((N_PY + 1))
            done
            [[ $N_PY -eq 0 ]] && err "Python 3.10+ still not found after install."
        else
            err "Python 3.10+ is required.\n  Install it: sudo dnf install python3"
        fi
    else
        err "Python 3.10+ not found.\n  Install it with your package manager (e.g. sudo apt install python3)"
    fi
fi

# Use the first Python we found for JSON parsing
PARSE_PY="${PY_BINS[0]}"

# ── Fetch release from GitHub ────────────────────────────────────────
ARCH=$(uname -m)
if [[ "$CHANNEL" == "prerelease" ]]; then
    info "Checking latest pre-release..."
    # /releases returns all releases sorted newest first; pick the first pre-release
    ALL_JSON=$(curl -sSf "https://api.github.com/repos/$REPO/releases?per_page=10" 2>/dev/null) \
        || err "Cannot reach GitHub API. Check your internet connection."
    RELEASE_JSON=$("$PARSE_PY" -c '
import json, sys
releases = json.loads(sys.stdin.read())
if isinstance(releases, dict) and "message" in releases:
    print(json.dumps(releases))
    sys.exit(0)
# Sort by published_at descending — GitHub API order is unreliable
releases.sort(key=lambda r: r.get("published_at", ""), reverse=True)
for r in releases:
    if r.get("prerelease", False):
        print(json.dumps(r))
        sys.exit(0)
print(json.dumps({"message": "No pre-release found. Create one on GitHub first."}))
' <<< "$ALL_JSON")
else
    info "Checking latest release..."
    RELEASE_JSON=$(curl -sSf "https://api.github.com/repos/$REPO/releases/latest" 2>/dev/null) \
        || err "Cannot reach GitHub API. Check your internet connection."
fi

# ── Match best Python + wheel from release assets ────────────────────
# Tries each installed Python version (newest first) against the
# available wheels. Falls back to sdist (build from source) if no
# pre-built wheel matches any installed Python.
export KRASIS_ARCH="$ARCH"
export KRASIS_PY_VERS="${PY_VERS[*]}"
export KRASIS_PY_BINS="${PY_BINS[*]}"
export KRASIS_PY_DOTS="${PY_DOTS[*]}"

RESULT=$("$PARSE_PY" -c '
import json, sys, os

data = json.loads(sys.stdin.read())

# Check for API errors (rate limit, not found, etc.)
if "message" in data and "assets" not in data:
    print("ERROR")
    print(data["message"])
    sys.exit(0)

tag = data.get("tag_name", "unknown")
arch = os.environ["KRASIS_ARCH"]
py_vers = os.environ["KRASIS_PY_VERS"].split()
py_bins = os.environ["KRASIS_PY_BINS"].split()
py_dots = os.environ["KRASIS_PY_DOTS"].split()

# Try each installed Python, newest first — pick the first with a wheel
for i, pyver in enumerate(py_vers):
    cpver = f"cp{pyver}"
    for asset in data.get("assets", []):
        name = asset["name"]
        parts = name.removesuffix(".whl").split("-")
        if (name.endswith(".whl") and len(parts) >= 5
                and parts[0] == "krasis" and parts[-3] == cpver
                and parts[-2] == cpver
                and any(platform.startswith(("manylinux", "linux_"))
                        and platform.endswith("_" + arch)
                        for platform in parts[-1].split("."))):
            print("WHEEL")
            print(tag)
            print(py_bins[i])
            print(pyver)
            print(py_dots[i])
            print(asset["browser_download_url"])
            sys.exit(0)

# No wheel for any installed Python — try sdist (build from source)
for asset in data.get("assets", []):
    name = asset["name"]
    if (name.startswith("krasis-") and name.endswith(".tar.gz")
            and name[len("krasis-"):][:1].isdigit()):
        print("SDIST")
        print(tag)
        print(py_bins[0])
        print(py_vers[0])
        print(py_dots[0])
        print(asset["browser_download_url"])
        sys.exit(0)

print("NONE")
print(tag)
' <<< "$RELEASE_JSON")

unset KRASIS_ARCH KRASIS_PY_VERS KRASIS_PY_BINS KRASIS_PY_DOTS

# ── Parse result ─────────────────────────────────────────────────────
mapfile -t LINES <<< "$RESULT"
RTYPE="${LINES[0]}"

case "$RTYPE" in
    ERROR)
        err "GitHub API: ${LINES[1]}\n  If rate-limited, wait a minute and try again."
        ;;
    WHEEL)
        VERSION="${LINES[1]}"
        PYTHON="${LINES[2]}"
        PY_VER="${LINES[3]}"
        PY_VER_DOT="${LINES[4]}"
        DOWNLOAD_URL="${LINES[5]}"
        info "Latest release: $VERSION (Python $PY_VER_DOT, $ARCH)"
        ;;
    SDIST)
        VERSION="${LINES[1]}"
        PYTHON="${LINES[2]}"
        PY_VER="${LINES[3]}"
        PY_VER_DOT="${LINES[4]}"
        DOWNLOAD_URL="${LINES[5]}"
        warn "No pre-built wheel for your Python version(s) on $ARCH."
        info "Building from source (requires Rust toolchain + C compiler)..."
        command -v cargo &>/dev/null \
            || err "Rust toolchain not found. Install it first:\n  curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh"
        ;;
    NONE)
        err "No packages found in release ${LINES[1]} for $ARCH.\n  Visit https://github.com/$REPO/releases"
        ;;
    *)
        err "Unexpected response from GitHub API."
        ;;
esac

info "Using $PYTHON (Python $PY_VER_DOT)"

# ── Ensure venv module ───────────────────────────────────────────────
if ! "$PYTHON" -c "import venv; import ensurepip" &>/dev/null; then
    warn "Python venv module not found. Installing python${PY_VER_DOT}-venv..."
    if command -v apt-get &>/dev/null; then
        SUDO=()
        [[ "$(id -u)" -ne 0 ]] && SUDO=(sudo)
        "${SUDO[@]}" apt-get update -qq </dev/null 2>/dev/null
        "${SUDO[@]}" apt-get install -y "python${PY_VER_DOT}-venv" </dev/null \
            || err "Failed to install python${PY_VER_DOT}-venv.\n  Try manually: sudo apt install python${PY_VER_DOT}-venv"
    elif command -v dnf &>/dev/null; then
        SUDO=()
        [[ "$(id -u)" -ne 0 ]] && SUDO=(sudo)
        "${SUDO[@]}" dnf install -y "python3-libs" </dev/null 2>/dev/null \
            || err "Failed to install python venv. Try: sudo dnf install python3-libs"
    else
        err "Python venv module not found.\n  Install it for your distro (e.g. sudo apt install python${PY_VER_DOT}-venv)"
    fi
    # Verify it worked
    "$PYTHON" -c "import venv" &>/dev/null \
        || err "Python venv module still not available after install attempt."
fi

# ── Setup venv ───────────────────────────────────────────────────────
# If the venv exists, check it's healthy and uses the right Python.
# If Python was upgraded or removed, the venv will be broken — recreate.
NEED_VENV=false
if [[ -d "$VENV_DIR" ]]; then
    VENV_PYVER=$("$VENV_DIR/bin/python" -c \
        "import sys; print(f'{sys.version_info.major}{sys.version_info.minor}')" 2>/dev/null || echo "0")
    if [[ "$VENV_PYVER" != "$PY_VER" ]]; then
        warn "Venv was Python ${VENV_PYVER:0:1}.${VENV_PYVER:1}, now using $PY_VER_DOT. Recreating..."
        rm -rf "$VENV_DIR"
        NEED_VENV=true
    elif ! "$VENV_DIR/bin/python" -c "import pip" &>/dev/null; then
        warn "Existing venv is broken. Recreating..."
        rm -rf "$VENV_DIR"
        NEED_VENV=true
    fi
else
    NEED_VENV=true
fi

if [[ "$NEED_VENV" == true ]]; then
    info "Creating environment at $VENV_DIR..."
    mkdir -p "$(dirname "$VENV_DIR")"
    "$PYTHON" -m venv "$VENV_DIR" 2>&1 || {
        # venv creation failed — likely missing ensurepip despite our earlier check
        warn "venv creation failed. Attempting to install python${PY_VER_DOT}-venv..."
        if command -v apt-get &>/dev/null; then
            SUDO=()
            [[ "$(id -u)" -ne 0 ]] && SUDO=(sudo)
            "${SUDO[@]}" apt-get update -qq </dev/null 2>/dev/null
            "${SUDO[@]}" apt-get install -y "python${PY_VER_DOT}-venv" </dev/null \
                || err "Failed to install python${PY_VER_DOT}-venv.\n  Try: sudo apt install python${PY_VER_DOT}-venv"
        elif command -v dnf &>/dev/null; then
            SUDO=()
            [[ "$(id -u)" -ne 0 ]] && SUDO=(sudo)
            "${SUDO[@]}" dnf install -y "python3-libs" </dev/null 2>/dev/null \
                || err "Failed to install python venv. Try: sudo dnf install python3-libs"
        else
            err "venv creation failed.\n  Install: sudo apt install python${PY_VER_DOT}-venv"
        fi
        # Retry
        rm -rf "$VENV_DIR"
        "$PYTHON" -m venv "$VENV_DIR" \
            || err "venv creation still failing after installing python${PY_VER_DOT}-venv."
    }
    "$VENV_DIR/bin/pip" install --upgrade pip -q \
        || warn "pip upgrade failed (non-fatal, continuing...)"
fi

cleanup_obsolete_venv_packages

# ── Create models directory ──────────────────────────────────────────
mkdir -p "$HOME/.krasis/models"

# ── Download and install ─────────────────────────────────────────────
TMPDIR=$(mktemp -d)
trap 'rm -rf "$TMPDIR"' EXIT

FILENAME="${DOWNLOAD_URL##*/}"
info "Downloading $FILENAME..."
curl -sSfL -o "$TMPDIR/$FILENAME" "$DOWNLOAD_URL" \
    || err "Download failed. Check your internet connection."

info "Installing Krasis ${VERSION#v}..."
PIP_INSTALL_ARGS=(install)
if [[ "$CHANNEL" == "prerelease" ]]; then
    PIP_INSTALL_ARGS+=(--force-reinstall --no-cache-dir)
fi
"$VENV_DIR/bin/pip" "${PIP_INSTALL_ARGS[@]}" "$TMPDIR/$FILENAME" \
    || err "pip install failed. Check the output above for details."

# ── Symlink commands ─────────────────────────────────────────────────
mkdir -p "$BIN_DIR"
for cmd in $COMMANDS; do
    if [[ -f "$VENV_DIR/bin/$cmd" ]]; then
        ln -sf "$VENV_DIR/bin/$cmd" "$BIN_DIR/$cmd"
    fi
done
ok "Commands installed: $COMMANDS"

# ── Ensure PATH ──────────────────────────────────────────────────────
ensure_path() {
    local rc="$1"
    local line='export PATH="$HOME/.local/bin:$PATH"'
    # Skip if .local/bin is already referenced
    if [[ -f "$rc" ]] && grep -q '\.local/bin' "$rc" 2>/dev/null; then
        return
    fi
    if [[ -f "$rc" ]]; then
        echo "" >> "$rc"
        echo "# Added by Krasis installer" >> "$rc"
        echo "$line" >> "$rc"
    fi
}

if ! echo "$PATH" | grep -q "$HOME/.local/bin"; then
    current_shell="$(basename "${SHELL:-bash}")"
    case "$current_shell" in
        zsh)
            ensure_path "$HOME/.zshrc"
            ;;
        bash)
            ensure_path "$HOME/.bashrc"
            [[ -f "$HOME/.profile" ]] && ensure_path "$HOME/.profile"
            ;;
        *)
            ensure_path "$HOME/.bashrc"
            ensure_path "$HOME/.profile"
            ;;
    esac
    # Make available in THIS session immediately
    export PATH="$HOME/.local/bin:$PATH"
fi

# ── Verify ───────────────────────────────────────────────────────────
if command -v krasis &>/dev/null; then
    ok "Krasis $(krasis --version 2>/dev/null || echo "${VERSION#v}") is ready!"
    echo ""
    info "Next steps:"
    echo -e "  1. Restart your terminal (or run ${BOLD}source ~/.bashrc${NC})"
    echo -e "  2. Run ${BOLD}krasis-setup${NC}   — checks CUDA tools and installs CUDA PyTorch"
    echo -e "  3. Run ${BOLD}krasis${NC}         — launch the interactive TUI"
    echo ""
else
    ok "Installed Krasis ${VERSION#v}."
    echo ""
    info "To start using it in this terminal, run:"
    echo ""
    echo -e "  ${BOLD}export PATH=\"\$HOME/.local/bin:\$PATH\"${NC}"
    echo ""
    info "Then:"
    echo -e "  1. Restart your terminal (or run the export above)"
    echo -e "  2. Run ${BOLD}krasis-setup${NC}   — checks CUDA tools and installs CUDA PyTorch"
    echo -e "  3. Run ${BOLD}krasis${NC}         — launch the interactive TUI"
    echo ""
fi
echo -e "${DIM}Upgrade:    krasis update${NC}"
echo -e "${DIM}Pre-release: krasis prerelease${NC}"
echo -e "${DIM}Uninstall:  curl -sSf https://raw.githubusercontent.com/brontoguana/krasis/main/install.sh | bash -s -- --uninstall${NC}"

} # end do_install

do_install "$@"