#!/bin/bash

set -euo pipefail

OMARCHY_THEME_PATH_V4="$HOME/.local/state/omarchy/current/theme"
OMARCHY_THEME_PATH_V3="$HOME/.config/omarchy/current/theme"
OMARCHY_THEME_PATH="$OMARCHY_THEME_PATH_V3"
ZED_CONFIG_DIR="$HOME/.config/zed"
ZED_THEMES_DIR="$ZED_CONFIG_DIR/themes"
ZED_SETTINGS_PATH="$ZED_CONFIG_DIR/settings.json"
LOG_FILE="$HOME/.local/share/omazed/sync.log"
OMARCHY_HOOKS_DIR="$HOME/.config/omarchy/hooks"
THEME_SET_HOOK="$OMARCHY_HOOKS_DIR/theme-set"
INITIALIZED_MARKER="$HOME/.local/share/omazed/initialized"

SCRIPT_DIR="$(cd "$(dirname "$(readlink -f "${BASH_SOURCE[0]}")")" && pwd)"
GENERATOR_SCRIPT=""
TEMPLATE_FILE=""
FONT_SCRIPT=""
for candidate_dir in "$SCRIPT_DIR" "$HOME/.local/bin" "/usr/bin"; do
    if [[ -z "$GENERATOR_SCRIPT" ]]; then
        if [[ -f "$candidate_dir/omazed-generator.sh" ]]; then
            GENERATOR_SCRIPT="$candidate_dir/omazed-generator.sh"
        elif [[ -f "$candidate_dir/omazed-converter.sh" ]]; then
            GENERATOR_SCRIPT="$candidate_dir/omazed-converter.sh"
        fi
    fi
    if [[ -z "$TEMPLATE_FILE" && -f "$candidate_dir/omazed-theme.tpl" ]]; then
        TEMPLATE_FILE="$candidate_dir/omazed-theme.tpl"
    fi
    if [[ -z "$FONT_SCRIPT" && -f "$candidate_dir/omazed-font.sh" ]]; then
        FONT_SCRIPT="$candidate_dir/omazed-font.sh"
    fi
done

if [[ -z "$GENERATOR_SCRIPT" || -z "$TEMPLATE_FILE" ]]; then
    echo -e "${RED}Error: Could not find omazed-generator.sh or omazed-theme.tpl${NC}" >&2
    exit 1
fi

RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m'

mkdir -p "$(dirname "$LOG_FILE")"

log() {
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] $*" | tee -a "$LOG_FILE"
}

info() {
    echo -e "${BLUE}[INFO]${NC} $*" | tee -a "$LOG_FILE" 2>/dev/null || true
}

success() {
    echo -e "${GREEN}[SUCCESS]${NC} $*" | tee -a "$LOG_FILE" 2>/dev/null || true
}

warn() {
echo -e "${YELLOW}[WARN]${NC} $*" | tee -a "$LOG_FILE" 2>/dev/null || true
}

error() {
echo -e "${RED}[ERROR]${NC} $*" | tee -a "$LOG_FILE" 2>/dev/null || true
}

detect_omarchy_version() {
    if command -v omarchy-version >/dev/null 2>&1; then
        local version
        version=$(omarchy-version 2>/dev/null | cut -d. -f1)
        if [[ "$version" =~ ^[0-9]+$ ]] && (( version >= 4 )); then
            echo "v4"
            return
        fi
    fi
    if [[ -d "$OMARCHY_THEME_PATH_V4" ]]; then
        echo "v4"
    else
        echo "v3"
    fi
}

set_theme_from_name() {
    local theme_name="$1"

    if [[ -z "$theme_name" ]]; then
        error "No theme name provided"
        return 1
    fi

    info "Setting Zed theme for: $theme_name"

    local theme_file="$ZED_THEMES_DIR/omazed.json"

    info "Generating Zed theme..."
    if generate_zed_theme; then
        if [[ -f "$theme_file" ]]; then
            success "Generated Zed theme"
        fi
    else
        info "Could not generate theme, skipping update"
        return 0
    fi

    success "Theme generated for: $theme_name"
    return 0
}

check_omarchy_hook_support() {
    if command -v omarchy-hook >/dev/null 2>&1; then
        return 0
    fi

    if [[ -d "$HOME/.config/omarchy/hooks/theme-set.d" ]]; then
        return 0
    fi

    if [[ -d "$HOME/.config/omarchy/hooks" ]]; then
        return 0
    fi

    return 1
}

setup_omarchy_hook() {

    local hook_marker_start="# >>> omazed hook - do not edit >>>"
    local hook_marker_end="# <<< omazed hook - do not edit <<<"
    local hook_command='omazed set "$1"'
    local version
    version=$(detect_omarchy_version)

    if [[ "$version" == "v4" ]]; then
        # v4 uses .d directory
        local hook_dir="$OMARCHY_HOOKS_DIR/theme-set.d"
        mkdir -p "$hook_dir"
        local v4_hook_file="$hook_dir/omazed"

        cat > "$v4_hook_file" << 'EOF'
#!/bin/bash
omazed set "$1"
EOF
        chmod +x "$v4_hook_file"
        local font_hook_dir="$OMARCHY_HOOKS_DIR/font-set.d"
        mkdir -p "$font_hook_dir"
        local v4_font_hook_file="$font_hook_dir/omazed"

        cat > "$v4_font_hook_file" << 'EOF'
#!/bin/bash
omazed set-font "$1"
EOF
        chmod +x "$v4_font_hook_file"
        info "Created v4 omarchy font hook: $v4_font_hook_file"

        info "Created v4 omarchy hook: $v4_hook_file"

        # Cleanup old v3 hook if it exists
        if [[ -f "$THEME_SET_HOOK" ]] && grep -q "$hook_marker_start" "$THEME_SET_HOOK" 2>/dev/null; then
            sed -i "/$hook_marker_start/,/$hook_marker_end/d" "$THEME_SET_HOOK"
            info "Removed legacy v3 omazed hook"
        fi
        success "✓ Omarchy hook configured (v4)"
    else
        # v3 uses single file injection
        if [[ ! -f "$THEME_SET_HOOK" ]]; then
            cat > "$THEME_SET_HOOK" << 'EOF'
#!/bin/bash
# This hook is called with the snake-cased name of the theme that has just been set.
EOF
            chmod +x "$THEME_SET_HOOK"
            info "Created theme-set hook file"
        fi

        if grep -q "$hook_marker_start" "$THEME_SET_HOOK" 2>/dev/null; then
            sed -i "/$hook_marker_start/,/$hook_marker_end/d" "$THEME_SET_HOOK"
            info "Removed old omazed hook"
        fi

        cat >> "$THEME_SET_HOOK" << EOF

$hook_marker_start
$hook_command
$hook_marker_end
EOF

        chmod +x "$THEME_SET_HOOK"
        success "✓ Omarchy hook configured (v3)"
    fi
}

setup_user() {
    info "Setting up Omazed for user: $USER"
    echo ""

    mkdir -p "$ZED_THEMES_DIR"

    if check_omarchy_hook_support; then
        setup_omarchy_hook
    else
        error "Omarchy hook system not available. Please update omarchy or install a older version of omazed < 1.2"
        return 1
    fi

    sync_all

    echo ""
    success "Setup complete!"
    echo ""
    success "🎉 Omazed is ready!"
    info "Change your Omarchy theme and watch Zed follow along automatically!"
    return 0
}

get_current_theme() {
    local theme=""
    local version
    version=$(detect_omarchy_version)

    local theme_name_path="$HOME/.config/omarchy/current/theme.name"
    if [[ "$version" == "v4" ]]; then
        theme_name_path="$HOME/.local/state/omarchy/current/theme.name"
        OMARCHY_THEME_PATH="$OMARCHY_THEME_PATH_V4"
    else
        OMARCHY_THEME_PATH="$OMARCHY_THEME_PATH_V3"
    fi

    if [[ -f "$theme_name_path" ]]; then
        theme=$(cat "$theme_name_path" 2>/dev/null | tr -d '\n\r' || echo "")
    fi

    if [[ -z "$theme" && -f "$OMARCHY_THEME_PATH" ]]; then
        theme=$(cat "$OMARCHY_THEME_PATH" 2>/dev/null | tr -d '\n\r' || echo "")
    fi

    if [[ -z "$theme" && -L "$OMARCHY_THEME_PATH" ]]; then
        local link_target
        link_target=$(readlink "$OMARCHY_THEME_PATH" 2>/dev/null || echo "")
        if [[ -n "$link_target" ]]; then
            theme=$(basename "$link_target" | sed 's/\.theme$//' || echo "")
        fi
    fi

    echo "$theme"
}

generate_zed_theme() {
    # Resolve the correct theme path based on detected version
    local version
    version=$(detect_omarchy_version)
    if [[ "$version" == "v4" ]]; then
        OMARCHY_THEME_PATH="$OMARCHY_THEME_PATH_V4"
    else
        OMARCHY_THEME_PATH="$OMARCHY_THEME_PATH_V3"
    fi

    local colors_toml="$OMARCHY_THEME_PATH/colors.toml"
    local alacritty_toml="$OMARCHY_THEME_PATH/alacritty.toml"
    local source_file="$colors_toml"
    local appearance="dark"

    if [[ ! -f "$GENERATOR_SCRIPT" ]]; then
        error "Theme generator script not found: $GENERATOR_SCRIPT"
        return 1
    fi

    if [[ ! -f "$colors_toml" ]]; then
        if [[ -f "$alacritty_toml" ]]; then
            warn "colors.toml not found: $colors_toml"
            info "Falling back to Alacritty config: $alacritty_toml"
            source_file="$alacritty_toml"
        else
            warn "colors.toml not found: $colors_toml"
            warn "Alacritty config not found: $alacritty_toml"
            return 1
        fi
    fi

    if [[ ! -f "$TEMPLATE_FILE" ]]; then
        warn "Template file not found: $TEMPLATE_FILE"
        return 1
    fi

    # Check for light mode — v4 uses mode= in colors.toml, v3 uses light.mode file
    if [[ -f "$OMARCHY_THEME_PATH/light.mode" ]]; then
        appearance="light"
    elif grep -qE '^mode\s*=\s*"light"' "$colors_toml" 2>/dev/null; then
        appearance="light"
    fi

    info "Generating Zed theme from colors..."
    info "Source: $source_file"
    info "Template: $TEMPLATE_FILE"
    info "Target: $ZED_THEMES_DIR/omazed.json"
    info "Appearance: $appearance"

    mkdir -p "$ZED_THEMES_DIR"

    local generation_output
    if generation_output=$(bash "$GENERATOR_SCRIPT" "$source_file" "$ZED_THEMES_DIR" "$TEMPLATE_FILE" "$appearance" 2>&1); then
        success "Successfully generated Zed theme"
        info "Generation output: $generation_output"
        return 0
    else
        error "Failed to generate theme"
        error "Generation error: $generation_output"
        return 1
    fi
}

set_initial_theme_once() {
    local theme_name="Omazed"

    if [[ -f "$INITIALIZED_MARKER" ]]; then
        return 0
    fi

    mkdir -p "$ZED_CONFIG_DIR"

    if [[ ! -f "$ZED_SETTINGS_PATH" ]]; then
        echo '{"theme": "'"$theme_name"'"}' > "$ZED_SETTINGS_PATH"
        touch "$INITIALIZED_MARKER"
        info "Zed theme set to: $theme_name"
        return 0
    fi

    if grep -q '"theme"' "$ZED_SETTINGS_PATH"; then
        temp_file=$(mktemp)
        awk -v theme="$theme_name" '
        BEGIN {
            in_theme_object = 0
            brace_count = 0
        }
        /"theme"[[:space:]]*:[[:space:]]*"[^"]*"/ {
            gsub(/"theme"[[:space:]]*:[[:space:]]*"[^"]*"/, "\"theme\": \"" theme "\"")
            print
            next
        }
        /"theme"[[:space:]]*:[[:space:]]*\{/ {
            gsub(/"theme"[[:space:]]*:[[:space:]]*\{.*/, "\"theme\": \"" theme "\",")
            in_theme_object = 1
            brace_count = 1
            print
            next
        }
        in_theme_object == 1 {
            brace_count += gsub(/\{/, "{")
            brace_count -= gsub(/\}/, "}")
            if (brace_count <= 0) {
                in_theme_object = 0
            }
            next
        }
        { print }
        ' "$ZED_SETTINGS_PATH" > "$temp_file"

        mv "$temp_file" "$ZED_SETTINGS_PATH"
        touch "$INITIALIZED_MARKER"
        info "Zed theme set to: $theme_name"
        return 0
    fi

    temp_file=$(mktemp)
    awk -v theme="$theme_name" '
    {
        lines[NR] = $0
    }
    END {
        close_index = 0
        for (i = NR; i >= 1; i--) {
            if (lines[i] ~ /^[[:space:]]*}[[:space:]]*$/) {
                close_index = i
                break
            }
        }

        if (close_index == 0) {
            for (i = 1; i <= NR; i++) print lines[i]
            print "  \"theme\": \"" theme "\""
            exit
        }

        prev = close_index - 1
        while (prev > 0 && lines[prev] ~ /^[[:space:]]*$/) prev--
        if (prev > 0 && lines[prev] !~ /,[[:space:]]*$/) {
            lines[prev] = lines[prev] ","
        }

        for (i = 1; i < close_index; i++) print lines[i]
        print "  \"theme\": \"" theme "\""
        print lines[close_index]
        for (i = close_index + 1; i <= NR; i++) print lines[i]
    }
    ' "$ZED_SETTINGS_PATH" > "$temp_file"

    mv "$temp_file" "$ZED_SETTINGS_PATH"
    touch "$INITIALIZED_MARKER"
    info "Zed theme set to: $theme_name"
}

sync_all() {
    local result=0
    sync_theme || result=1
    if [[ -f "$FONT_SCRIPT" ]]; then
        "$FONT_SCRIPT" || result=1
    fi
    return "$result"
}

sync_theme() {
    local current_theme
    current_theme=$(get_current_theme)

    if [[ -z "$current_theme" ]]; then
        warn "Could not determine current Omarchy theme"
        return 1
    fi

    local theme_file="$ZED_THEMES_DIR/omazed.json"
    if [[ -f "$theme_file" ]]; then
        rm -f "$theme_file"
    fi

    info "Attempting to generate Zed theme..."
    if generate_zed_theme; then
        if [[ -f "$theme_file" ]]; then
            success "Generated Zed theme: $theme_file"
        else
            warn "Theme generation completed but file not found at expected location: $theme_file"
        fi
    else
        info "Skipping Zed theme update - generation failed"
        return 0
    fi

    success "Theme sync completed: $current_theme"
    return 0
}

cleanup() {
    return 0
}

trap cleanup EXIT
trap 'error "Interrupted"; exit 1' INT TERM

usage() {
    cat << EOF
Omazed - Live theme switching for zed in omarchy

USAGE:
    $0 [COMMAND] [OPTIONS]

COMMANDS:
    setup       Set up themes and automatic theme sync hook
    sync        Sync theme once and exit
    set         Set theme by name (used by omarchy hooks)
    set-font    Set font by name (used by omarchy hooks)

OPTIONS:
    -h, --help  Show this help

EXAMPLES:
    $0 set "tokyo-night"    # Set Zed theme to Tokyo Night
    $0 sync   # Set Zed theme with system theme

EOF
}

main() {
    local command=""
    local theme_arg=""

    if [[ $# -eq 0 ]]; then
        usage
        exit 0
    fi

    while [[ $# -gt 0 ]]; do
        case $1 in
            sync)
                command="sync"
                shift
                ;;
            set)
                command="set"
                shift
                theme_arg="$1"
                shift
                ;;
            set-font)
                command="set-font"
                shift
                theme_arg="${1:-}"
                [[ $# -gt 0 ]] && shift
                ;;
            setup)
                command="setup"
                shift
                ;;
            -h|--help)
                usage
                exit 0
                ;;
            *)
                error "Unknown option: $1"
                usage
                exit 1
                ;;
        esac
    done

    set_initial_theme_once

    case $command in
        sync)
            sync_all
            ;;

        set)
            if [[ -z "$theme_arg" ]]; then
                error "Theme name required for 'set' command"
                error "Usage: omazed set \"theme-name\""
                exit 1
            fi
            set_theme_from_name "$theme_arg"
            ;;
        set-font)
            if [[ -f "$FONT_SCRIPT" ]]; then
                "$FONT_SCRIPT" "$theme_arg"
            else
                error "Font script not found at $FONT_SCRIPT"
                exit 1
            fi
            ;;
        setup)
            setup_user
            ;;
    esac
}

main "$@"
