#!/usr/bin/env bash
set -euo pipefail

# ============================================================================
# AILinux Mirror Repo Adder v5.0 - Intelligent Auto-Discovery
# ============================================================================
# Automatically discovers ALL repositories by analyzing the mirror directory
# structure. Works both locally (filesystem scan) and remotely (HTTP scan).
#
# Usage:
#   curl -sSL https://repo.ailinux.me/mirror/add-ailinux-repo.sh | sudo bash
#   sudo ./add-ailinux-repo.sh [OPTIONS]
#
# Options:
#   --no-update     Skip apt-get update
#   --list-repos    Show discovered repositories and exit
#   --dry-run       Show what would be done without making changes
#   --verbose       Show detailed discovery process
#   -h, --help      Show help
# ============================================================================

# === CONFIGURATION ===
DEFAULT_BASE="https://repo.ailinux.me/mirror"
BASE_URL="${AILINUX_REPO_BASE:-$DEFAULT_BASE}"
KEYRING_DIR="/usr/share/keyrings"
KEYRING_PATH="${KEYRING_DIR}/ailinux-archive-keyring.gpg"
LIST_PATH="/etc/apt/sources.list.d/ailinux-mirror.list"

# curl options: force IPv4, timeout, retry
CURL_OPTS="-4 --connect-timeout 10 --max-time 30 --retry 2 --retry-delay 1"

# Runtime flags
VERBOSE=0
DRY_RUN=0

# === COLORS ===
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
CYAN='\033[0;36m'
BOLD='\033[1m'
NC='\033[0m'

log_info()  { echo -e "${BLUE}[*]${NC} $*"; }
log_ok()    { echo -e "${GREEN}[+]${NC} $*"; }
log_warn()  { echo -e "${YELLOW}[!]${NC} $*"; }
log_error() { echo -e "${RED}[ERROR]${NC} $*" >&2; }
log_debug() { [[ $VERBOSE -eq 1 ]] && echo -e "${CYAN}[DBG]${NC} $*" >&2; }

# === HELPER FUNCTIONS ===

usage() {
    cat <<'USAGE'
AILinux Mirror Repository Adder v5.0 (Intelligent Auto-Discovery)

Automatically discovers ALL repositories by analyzing the mirror directory
structure. Works both locally and remotely.

Usage:
  curl -sSL https://repo.ailinux.me/mirror/add-ailinux-repo.sh | sudo bash
  sudo ./add-ailinux-repo.sh [OPTIONS]

Options:
  --no-update     Skip apt-get update after adding repos
  --list-repos    Show all discovered repositories and exit
  --dry-run       Show what would be done without making changes
  --verbose       Show detailed discovery process
  -h, --help      Show this help

Environment Variables:
  AILINUX_REPO_BASE   Override base URL (default: https://repo.ailinux.me/mirror)
USAGE
}

need_cmd() {
    command -v "$1" >/dev/null 2>&1 || {
        log_error "Required command not found: $1"
        return 1
    }
}

detect_codename() {
    local codename=""
    if [[ -r /etc/os-release ]]; then
        . /etc/os-release
        codename="${VERSION_CODENAME:-${UBUNTU_CODENAME:-}}"
    fi
    [[ -z "$codename" ]] && command -v lsb_release >/dev/null 2>&1 && \
        codename=$(lsb_release -cs 2>/dev/null || true)

    if [[ -z "$codename" ]]; then
        log_error "Could not detect distribution codename"
        return 1
    fi
    echo "$codename"
}

# Check if running locally (script exists in mirror dir)
is_local_mode() {
    local script_path="${BASH_SOURCE[0]:-}"
    [[ -n "$script_path" ]] && [[ -f "$script_path" ]] && \
        [[ -d "$(dirname "$script_path")" ]] && \
        [[ -d "$(dirname "$script_path")/archive.ubuntu.com" ]]
}

# === KEY MANAGEMENT ===

install_key() {
    mkdir -p "${KEYRING_DIR}"
    local tmp key_url="${BASE_URL}/ailinux-archive-key.gpg"
    tmp="$(mktemp)"
    trap "rm -f '$tmp'" RETURN

    log_info "Fetching GPG key from: $key_url"
    if curl $CURL_OPTS -fsSL "$key_url" -o "$tmp" 2>/dev/null; then
        if file "$tmp" 2>/dev/null | grep -qi "PGP public key"; then
            cp -f "$tmp" "${KEYRING_PATH}"
        elif grep -q "BEGIN PGP" "$tmp" 2>/dev/null; then
            gpg --batch --yes --dearmor --output "${KEYRING_PATH}" "$tmp" 2>/dev/null
        else
            gpg --batch --yes --dearmor --output "${KEYRING_PATH}" "$tmp" 2>/dev/null || \
                cp -f "$tmp" "${KEYRING_PATH}"
        fi
        chmod 0644 "${KEYRING_PATH}"
        log_ok "Installed keyring: ${KEYRING_PATH}"
        return 0
    fi
    log_error "Failed to fetch GPG key"
    return 1
}

# === INTELLIGENT DISCOVERY ===

# Get friendly name for a repository path
get_repo_name() {
    local path="$1"
    local host="${path%%/*}"

    case "$host" in
        archive.ubuntu.com)     echo "Ubuntu Archive" ;;
        security.ubuntu.com)    echo "Ubuntu Security" ;;
        archive.neon.kde.org)   echo "KDE Neon" ;;
        dl.google.com)          echo "Google Chrome" ;;
        dl.winehq.org)          echo "WineHQ" ;;
        download.docker.com)    echo "Docker CE" ;;
        repo.steampowered.com)  echo "Steam" ;;
        deb.nodesource.com)     echo "Node.js" ;;
        packages.microsoft.com) echo "VS Code" ;;
        cli.github.com)         echo "GitHub CLI" ;;
        download.sublimetext.com) echo "Sublime Text" ;;
        apt.releases.hashicorp.com) echo "HashiCorp" ;;
        developer.download.nvidia.com) echo "NVIDIA CUDA" ;;
        nvidia.github.io)       echo "NVIDIA Container" ;;
        apt.repos.intel.com)    echo "Intel oneAPI" ;;
        updates.signal.org)     echo "Signal" ;;
        ppa.launchpadcontent.net)
            # Extract PPA name
            if [[ "$path" == *"/cappelikan/"* ]]; then echo "Mainline Kernels"
            elif [[ "$path" == *"/graphics-drivers/"* ]]; then echo "NVIDIA Drivers PPA"
            elif [[ "$path" == *"/kisak/"* ]]; then echo "Kisak Mesa"
            elif [[ "$path" == *"/oibaf/"* ]]; then echo "Oibaf Graphics"
            elif [[ "$path" == *"/fastfetch/"* ]]; then echo "Fastfetch"
            elif [[ "$path" == *"/lutris-team/"* ]]; then echo "Lutris"
            elif [[ "$path" == *"/libreoffice/"* ]]; then echo "LibreOffice"
            else echo "PPA: ${path#ppa.launchpadcontent.net/}"
            fi
            ;;
        pkgs.k8s.io)            echo "Kubernetes" ;;
        *)                      echo "$host" ;;
    esac
}

# Get section/category for a repository
get_repo_section() {
    local path="$1"

    case "$path" in
        archive.ubuntu.com*|security.ubuntu.com*) echo "UBUNTU" ;;
        *neon*) echo "DESKTOP" ;;
        *microsoft*|*github.com*|*sublime*|*nodesource*|*docker*|*k8s*|*hashicorp*) echo "DEV" ;;
        *lutris*|*wine*|*steam*|*chrome*|*google*) echo "GAMING" ;;
        *graphics-drivers*|*mesa*|*oibaf*) echo "GRAPHICS" ;;
        *cuda*|*nvidia*|*intel*) echo "AI/ML" ;;
        *cappelikan*|*fastfetch*|*libreoffice*) echo "SYSTEM" ;;
        *signal*) echo "MESSAGING" ;;
        *) echo "OTHER" ;;
    esac
}

# Check if repository supports i386 architecture
# Returns "amd64,i386" for multi-arch repos, "amd64" for amd64-only
get_repo_arch() {
    local path="$1"

    case "$path" in
        # Ubuntu repos have full i386 support
        archive.ubuntu.com*|security.ubuntu.com*)
            echo "amd64,i386" ;;
        # Wine needs i386 for 32-bit Windows apps
        *winehq*|*wine*)
            echo "amd64,i386" ;;
        # Steam needs i386 for 32-bit games
        *steampowered*|*steam*)
            echo "amd64,i386" ;;
        # Lutris may need i386 for some games
        *lutris*)
            echo "amd64,i386" ;;
        # All other repos are amd64 only
        *)
            echo "amd64" ;;
    esac
}

# LOCAL DISCOVERY: Scan filesystem
discover_local() {
    local mirror_dir="$1"

    log_debug "Scanning local directory: $mirror_dir"

    # Find all dists directories
    find "$mirror_dir" -maxdepth 8 -type d -name "dists" 2>/dev/null | while read -r dists_dir; do
        local repo_root="${dists_dir%/dists}"
        local rel_path="${repo_root#$mirror_dir/}"

        log_debug "Found dists in: $rel_path"

        # Scan distributions in this repo
        for dist_path in "$dists_dir"/*/; do
            [[ -d "$dist_path" ]] || continue
            local dist_name
            dist_name=$(basename "$dist_path")

            # Skip special directories
            [[ "$dist_name" == "by-hash" ]] && continue

            # Check if it has a Release file
            [[ -f "${dist_path}Release" ]] || [[ -f "${dist_path}InRelease" ]] || continue

            # Find components
            local components=""
            for comp_dir in "$dist_path"/*/; do
                [[ -d "$comp_dir" ]] || continue
                local comp_name
                comp_name=$(basename "$comp_dir")

                # Skip non-component dirs
                [[ "$comp_name" == "by-hash" ]] && continue

                # Verify it's a real component (has binary-* or source)
                if [[ -d "${comp_dir}binary-amd64" ]] || \
                   [[ -d "${comp_dir}binary-i386" ]] || \
                   [[ -d "${comp_dir}source" ]]; then
                    components="${components} ${comp_name}"
                fi
            done

            components="${components# }"  # trim leading space

            # Output: path|dist|components|name|section
            local name
            name=$(get_repo_name "$rel_path")
            local section
            section=$(get_repo_section "$rel_path")

            echo "${rel_path}|${dist_name}|${components}|${name}|${section}"
        done
    done | sort -u
}

# REMOTE DISCOVERY: Scan via HTTP
discover_remote() {
    local base_url="$1"

    log_debug "Scanning remote URL: $base_url"

    # First, get list of top-level directories from index
    local index_html
    index_html=$(curl $CURL_OPTS -fsSL "${base_url}/" 2>/dev/null) || {
        log_error "Cannot fetch mirror index"
        return 1
    }

    # Parse directory links - handle both relative and absolute URLs
    # Look for patterns like href="archive.ubuntu.com/" or href="https://repo.ailinux.me/mirror/archive.ubuntu.com/"
    local top_dirs
    top_dirs=$(echo "$index_html" | grep -oE 'href="[^"]+' | sed 's/href="//' | \
               grep -E '(archive\.|dl\.|download\.|repo\.|deb\.|ppa\.|packages\.|cli\.|apt\.|developer\.|nvidia\.|updates\.|pkgs\.)' | \
               sed "s|${base_url}/||g" | sed 's|/$||' | sort -u)

    for dir in $top_dirs; do
        [[ -z "$dir" ]] && continue
        # Skip non-directory entries
        [[ "$dir" == *.gpg ]] && continue
        [[ "$dir" == *.sh ]] && continue
        [[ "$dir" == *.html ]] && continue

        log_debug "Checking top-level: $dir"

        # Recursively find dists in this hierarchy
        discover_remote_recursive "$base_url" "$dir" ""
    done | sort -u
}

# Recursive remote discovery
discover_remote_recursive() {
    local base_url="$1"
    local current_path="$2"
    local depth="$3"

    # Limit recursion depth
    [[ ${#depth} -gt 20 ]] && return

    local full_path="${current_path%/}"
    local check_url="${base_url}/${full_path}"

    # Check if this path has a dists directory
    local release_found=0
    local dist_names=""

    # Try to fetch dists listing
    local dists_html
    dists_html=$(curl $CURL_OPTS -fsSL "${check_url}/dists/" 2>/dev/null) || dists_html=""

    if [[ -n "$dists_html" ]]; then
        # Found dists! Parse distributions
        dist_names=$(echo "$dists_html" | grep -oE 'href="[^"]+/"' | sed 's/href="//;s/\/"//' | \
                     grep -v '^\.\.' | grep -v '^/' | grep -v 'http' | grep -v 'by-hash')

        for dist in $dist_names; do
            [[ -z "$dist" ]] && continue

            # Check for Release file
            if curl $CURL_OPTS -fsSI "${check_url}/dists/${dist}/Release" >/dev/null 2>&1 || \
               curl $CURL_OPTS -fsSI "${check_url}/dists/${dist}/InRelease" >/dev/null 2>&1; then

                log_debug "Found dist: ${full_path}/dists/${dist}"

                # Try to get components from Release file
                local components=""
                local release_content
                release_content=$(curl $CURL_OPTS -fsSL "${check_url}/dists/${dist}/Release" 2>/dev/null | head -50) || release_content=""

                if [[ -n "$release_content" ]]; then
                    components=$(echo "$release_content" | grep -i "^Components:" | sed 's/Components://i' | tr -d '\r')
                    components="${components# }"
                fi

                # Fallback: scan for component directories
                if [[ -z "$components" ]]; then
                    local comp_html
                    comp_html=$(curl $CURL_OPTS -fsSL "${check_url}/dists/${dist}/" 2>/dev/null) || comp_html=""
                    for comp in main restricted universe multiverse stable steam; do
                        if echo "$comp_html" | grep -q "href=\"${comp}/\""; then
                            components="${components} ${comp}"
                        fi
                    done
                    components="${components# }"
                fi

                local name
                name=$(get_repo_name "$full_path")
                local section
                section=$(get_repo_section "$full_path")

                echo "${full_path}|${dist}|${components}|${name}|${section}"
            fi
        done
        return
    fi

    # No dists here, try subdirectories
    local sub_html
    sub_html=$(curl $CURL_OPTS -fsSL "${check_url}/" 2>/dev/null) || return

    local subdirs
    subdirs=$(echo "$sub_html" | grep -oE 'href="[^"]+/"' | sed 's/href="//;s/\/"//' | \
              grep -v '^\.\.' | grep -v '^/' | grep -v 'http' | head -10)

    for subdir in $subdirs; do
        [[ -z "$subdir" ]] && continue
        discover_remote_recursive "$base_url" "${full_path}/${subdir}" "${depth}."
    done
}

# Main discovery function
discover_repositories() {
    local base="$1"
    local mode="$2"

    if [[ "$mode" == "local" ]]; then
        local script_dir
        script_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
        discover_local "$script_dir"
    else
        discover_remote "$base"
    fi
}

# === OUTPUT FUNCTIONS ===

list_repos() {
    local base="$1"
    local mode="$2"

    echo ""
    echo "==========================================="
    echo "  AILinux Mirror - Auto-Discovery Results"
    echo "==========================================="
    echo ""
    echo "  Mode: ${mode}"
    echo "  Base: ${base}"
    echo ""

    local repos
    repos=$(discover_repositories "$base" "$mode")

    if [[ -z "$repos" ]]; then
        log_error "No repositories discovered!"
        return 1
    fi

    local count=0
    local current_section=""

    while IFS='|' read -r path dist components name section; do
        [[ -z "$path" ]] && continue

        if [[ "$section" != "$current_section" ]]; then
            echo ""
            echo -e "  ${BOLD}--- ${section} ---${NC}"
            current_section="$section"
        fi

        ((count++)) || true
        echo -e "  ${CYAN}[${count}]${NC} ${GREEN}${name}${NC}"
        echo "      Path: ${path}"
        echo "      Dist: ${dist}"
        [[ -n "$components" ]] && echo "      Components: ${components}"
    done <<< "$repos"

    echo ""
    echo "==========================================="
    echo "  Total: ${count} repositories found"
    echo "==========================================="
}

generate_sources_list() {
    local base="$1"
    local mode="$2"
    local codename="$3"

    echo "# ============================================"
    echo "# AILinux Mirror Repositories"
    echo "# Auto-generated: $(date -Iseconds)"
    echo "# Base URL: ${base}"
    echo "# Mode: ${mode}"
    echo "# System: ${codename}"
    echo "# Keyring: ${KEYRING_PATH}"
    echo "# ============================================"

    local repos
    repos=$(discover_repositories "$base" "$mode")

    if [[ -z "$repos" ]]; then
        echo "# ERROR: No repositories discovered!"
        return 1
    fi

    local current_section=""

    while IFS='|' read -r path dist components name section; do
        [[ -z "$path" ]] && continue

        # Filter: only include repos matching codename or universal ones
        local include=0
        case "$dist" in
            "$codename"*) include=1 ;;  # matches noble, noble-updates, etc.
            stable|nodistro|all|xenial|/) include=1 ;;  # universal dists
        esac

        [[ $include -eq 0 ]] && continue

        if [[ "$section" != "$current_section" ]]; then
            echo ""
            echo "# --- ${section} ---"
            current_section="$section"
        fi

        # Get architecture for this repo
        local arch
        arch=$(get_repo_arch "$path")

        echo "# ${name} (${dist})"
        if [[ -z "$components" ]] || [[ "$dist" == "/" ]]; then
            echo "deb [arch=${arch} signed-by=${KEYRING_PATH}] ${base}/${path} ${dist}"
        else
            echo "deb [arch=${arch} signed-by=${KEYRING_PATH}] ${base}/${path} ${dist} ${components}"
        fi
    done <<< "$repos"
}

# === MAIN ===

main() {
    local do_update="yes"
    local list_only="no"
    local mode="auto"

    while [[ $# -gt 0 ]]; do
        case "$1" in
            --no-update)  do_update="no"; shift ;;
            --list-repos) list_only="yes"; shift ;;
            --dry-run)    DRY_RUN=1; shift ;;
            --verbose)    VERBOSE=1; shift ;;
            -h|--help)    usage; exit 0 ;;
            *)            log_error "Unknown: $1"; usage; exit 1 ;;
        esac
    done

    # Detect mode
    if is_local_mode; then
        mode="local"
        local script_dir
        script_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
        BASE_URL="file://${script_dir}"
    else
        mode="remote"
    fi

    # Detect codename
    local codename
    codename=$(detect_codename) || exit 1

    # List only mode
    if [[ "$list_only" == "yes" ]]; then
        list_repos "$BASE_URL" "$mode"
        exit 0
    fi

    # Check root (unless dry-run)
    if [[ $EUID -ne 0 ]] && [[ $DRY_RUN -eq 0 ]]; then
        log_error "This script must be run as root (use sudo)"
        exit 1
    fi

    need_cmd curl || exit 1

    echo ""
    echo "==========================================="
    echo "  AILinux Mirror Repository Installer v5"
    echo "==========================================="
    echo ""
    log_info "Mode: ${mode}"
    log_info "Detected codename: ${codename}"
    log_info "Base URL: ${BASE_URL}"
    [[ $DRY_RUN -eq 1 ]] && log_warn "DRY-RUN MODE"
    echo ""

    # Discover repositories
    log_info "Discovering repositories..."

    local sources_content
    sources_content=$(generate_sources_list "$BASE_URL" "$mode" "$codename")

    if [[ -z "$sources_content" ]] || echo "$sources_content" | grep -q "ERROR:"; then
        log_error "Discovery failed!"
        exit 1
    fi

    local repo_count
    repo_count=$(echo "$sources_content" | grep -c "^deb " || echo 0)

    if [[ $DRY_RUN -eq 1 ]]; then
        echo ""
        log_info "Would write to: ${LIST_PATH}"
        echo "--- Preview ---"
        echo "$sources_content"
        echo "--- End ---"
        log_info "Would configure ${repo_count} repositories"
        exit 0
    fi

    # Install prerequisites
    export DEBIAN_FRONTEND=noninteractive
    log_info "Installing prerequisites..."
    apt-get update -qq >/dev/null 2>&1 || true
    apt-get install -y -qq ca-certificates curl gnupg >/dev/null 2>&1 || true

    # Install GPG key
    if [[ "$mode" == "remote" ]]; then
        log_info "Installing GPG keyring..."
        install_key || {
            log_warn "GPG key installation failed, continuing anyway..."
        }
    else
        # Local mode: copy key directly
        local script_dir
        script_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
        if [[ -f "${script_dir}/ailinux-archive-key.gpg" ]]; then
            mkdir -p "${KEYRING_DIR}"
            cp -f "${script_dir}/ailinux-archive-key.gpg" "${KEYRING_PATH}"
            chmod 0644 "${KEYRING_PATH}"
            log_ok "Installed local keyring"
        fi
    fi

    # Write sources list
    mkdir -p "$(dirname "${LIST_PATH}")"
    echo "$sources_content" > "${LIST_PATH}"
    chmod 0644 "${LIST_PATH}"

    log_ok "Written: ${LIST_PATH}"
    log_ok "Configured ${repo_count} repositories"
    echo ""

    echo "--- ${LIST_PATH} ---"
    cat "${LIST_PATH}"
    echo "--- end ---"
    echo ""

    # Update apt
    if [[ "$do_update" == "yes" ]]; then
        log_info "Running apt-get update..."
        if apt-get update 2>&1; then
            log_ok "APT update completed"
        else
            log_warn "APT update had errors (see above)"
        fi
    fi

    echo ""
    echo "==========================================="
    log_ok "AILinux mirror configured!"
    echo "==========================================="
    echo ""
}

main "$@"