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

# Migrate Cargo.toml dependencies to workspace-level declarations
# Only modifies dependency sections - preserves all other formatting

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

SKIP_GIT_CHECK=false
SKIP_DIRTY_CHECK=false

warn() { echo -e "${YELLOW}warning:${NC} $*" >&2; }
error() { echo -e "${RED}error:${NC} $*" >&2; }
success() { echo -e "${GREEN}$*${NC}"; }

usage() {
    echo "Usage: $0 [--no-git] [--allow-dirty]"
    echo "  --no-git       Skip git repository check"
    echo "  --allow-dirty  Allow dirty working directory"
    exit 0
}

while [[ $# -gt 0 ]]; do
    case $1 in
        --no-git) SKIP_GIT_CHECK=true; shift ;;
        --allow-dirty) SKIP_DIRTY_CHECK=true; shift ;;
        -h|--help) usage ;;
        *) error "Unknown option: $1"; exit 1 ;;
    esac
done

# Check required tools
for tool in tomlq jq git; do
    command -v "$tool" &>/dev/null || { error "Required tool '$tool' not found"; exit 1; }
done

# Find workspace root
find_workspace_root() {
    local dir="$PWD"
    while [[ "$dir" != "/" ]]; do
        if [[ -f "$dir/Cargo.toml" ]] && tomlq -e '.workspace' "$dir/Cargo.toml" &>/dev/null; then
            echo "$dir"
            return 0
        fi
        dir="$(dirname "$dir")"
    done
    error "No workspace root found (Cargo.toml with [workspace] section)"
    exit 1
}

WORKSPACE_ROOT="$(find_workspace_root)"
echo "Workspace root: $WORKSPACE_ROOT"
cd "$WORKSPACE_ROOT" || exit 1

# Git checks
if [[ "$SKIP_GIT_CHECK" == false ]]; then
    git rev-parse --git-dir &>/dev/null || { error "Not a git repo. Use --no-git to skip."; exit 1; }
    if [[ "$SKIP_DIRTY_CHECK" == false ]]; then
        if ! git diff --quiet 2>/dev/null || ! git diff --cached --quiet 2>/dev/null; then
            error "Dirty working directory. Use --allow-dirty to skip."
            exit 1
        fi
    fi
fi

# Temp files
COLLECTED=$(mktemp)
SKIPPED=$(mktemp)
MODIFIED_FILES=$(mktemp)
trap 'rm -f "$COLLECTED" "$SKIPPED" "$MODIFIED_FILES"' EXIT

# Existing workspace deps
EXISTING_WS_DEPS=$(tomlq -r '.workspace.dependencies // {} | keys[]' Cargo.toml 2>/dev/null || true)

# Version comparison
compare_versions() {
    local v1="$1" v2="$2"
    local v1_maj v1_min v1_pat v2_maj v2_min v2_pat
    
    IFS='.' read -r v1_maj v1_min v1_pat <<< "$v1"
    IFS='.' read -r v2_maj v2_min v2_pat <<< "$v2"
    v1_min="${v1_min:-0}"; v1_pat="${v1_pat:-0}"
    v2_min="${v2_min:-0}"; v2_pat="${v2_pat:-0}"
    
    if [[ "$v1_maj" == "0" || "$v2_maj" == "0" ]]; then
        [[ "$v1_maj" == "$v2_maj" && "$v1_min" == "$v2_min" ]] || return 1
        [[ "$v1_pat" -ge "$v2_pat" ]] && echo "$v1" || echo "$v2"
        return 0
    fi
    
    [[ "$v1_maj" == "$v2_maj" ]] || return 1
    
    if [[ "$v1_min" -gt "$v2_min" ]]; then echo "$v1"
    elif [[ "$v1_min" -lt "$v2_min" ]]; then echo "$v2"
    elif [[ "$v1_pat" -ge "$v2_pat" ]]; then echo "$v1"
    else echo "$v2"
    fi
}

is_simple_version() {
    [[ "$1" =~ ^[\^=]?[0-9]+(\.[0-9]+)*$ ]]
}

normalize_version() {
    local v="$1"
    v="${v#^}"; v="${v#=}"
    echo "$v"
}

list_crate_tomls() {
    tomlq -r '.workspace.members // [] | .[]' Cargo.toml 2>/dev/null | while read -r pattern; do
        for dir in $pattern; do
            [[ -f "$dir/Cargo.toml" ]] && echo "$dir/Cargo.toml"
        done
    done
}

# Process one dependency (collection phase - read only)
process_dep() {
    local crate_toml="$1" dep_name="$2" dep_json="$3"
    
    echo "$dep_json" | jq -e '.workspace == true' &>/dev/null && return 0
    echo "$EXISTING_WS_DEPS" | grep -qxF "$dep_name" && return 0
    
    local version="" path="" git="" branch="" tag="" rev=""
    
    if echo "$dep_json" | jq -e 'type == "string"' &>/dev/null; then
        version=$(echo "$dep_json" | jq -r '.')
    else
        version=$(echo "$dep_json" | jq -r '.version // empty')
        path=$(echo "$dep_json" | jq -r '.path // empty')
        git=$(echo "$dep_json" | jq -r '.git // empty')
        branch=$(echo "$dep_json" | jq -r '.branch // empty')
        tag=$(echo "$dep_json" | jq -r '.tag // empty')
        rev=$(echo "$dep_json" | jq -r '.rev // empty')
    fi
    
    # Check for package rename
    local package=""
    if ! echo "$dep_json" | jq -e 'type == "string"' &>/dev/null; then
        package=$(echo "$dep_json" | jq -r '.package // empty')
    fi
    
    local ws_value="" ws_type=""
    
    if [[ -n "$path" ]]; then
        local crate_dir abs_path rel_path
        crate_dir=$(dirname "$crate_toml")
        abs_path=$(cd "$crate_dir" && cd "$path" && pwd)
        rel_path=$(realpath --relative-to="$WORKSPACE_ROOT" "$abs_path")
        ws_value="{ path = \"$rel_path\""
        [[ -n "$package" ]] && ws_value="$ws_value, package = \"$package\""
        ws_value="$ws_value }"
        ws_type="path"
    elif [[ -n "$git" ]]; then
        ws_value="{ git = \"$git\""
        [[ -n "$branch" ]] && ws_value="$ws_value, branch = \"$branch\""
        [[ -n "$tag" ]] && ws_value="$ws_value, tag = \"$tag\""
        [[ -n "$rev" ]] && ws_value="$ws_value, rev = \"$rev\""
        [[ -n "$package" ]] && ws_value="$ws_value, package = \"$package\""
        ws_value="$ws_value }"
        ws_type="git"
    elif [[ -n "$version" ]]; then
        if ! is_simple_version "$version"; then
            warn "skipping '$dep_name' in $crate_toml: complex version '$version'"
            echo "$dep_name" >> "$SKIPPED"
            return 0
        fi
        version=$(normalize_version "$version")
        if [[ -n "$package" ]]; then
            ws_value="{ version = \"$version\", package = \"$package\" }"
        else
            ws_value="\"$version\""
        fi
        ws_type="version"
    else
        warn "skipping '$dep_name' in $crate_toml: no version/path/git"
        return 0
    fi
    
    local existing
    existing=$(grep "^${dep_name}	" "$COLLECTED" 2>/dev/null || true)
    
    if [[ -n "$existing" ]]; then
        local existing_value existing_type
        existing_value=$(echo "$existing" | cut -f2)
        existing_type=$(echo "$existing" | cut -f3)
        
        if [[ "$ws_type" == "version" && "$existing_type" == "version" ]]; then
            local existing_ver new_ver unified_ver
            existing_ver=$(echo "$existing_value" | tr -d '"')
            new_ver="$version"
            
            if unified_ver=$(compare_versions "$existing_ver" "$new_ver"); then
                ws_value="\"$unified_ver\""
                grep -v "^${dep_name}	" "$COLLECTED" > "$COLLECTED.tmp" || true
                mv "$COLLECTED.tmp" "$COLLECTED"
                printf '%s\t%s\t%s\n' "$dep_name" "$ws_value" "$ws_type" >> "$COLLECTED"
            else
                warn "skipping '$dep_name': incompatible versions '$existing_ver' vs '$new_ver'"
                echo "$dep_name" >> "$SKIPPED"
            fi
            return 0
        fi
        
        if [[ "$ws_value" != "$existing_value" ]]; then
            warn "skipping '$dep_name': conflicting definitions"
            echo "$dep_name" >> "$SKIPPED"
        fi
        return 0
    fi
    
    printf '%s\t%s\t%s\n' "$dep_name" "$ws_value" "$ws_type" >> "$COLLECTED"
}

process_section() {
    local crate_toml="$1" section="$2"
    local deps_json
    deps_json=$(tomlq ".\"$section\" // {}" "$crate_toml" 2>/dev/null)
    
    echo "$deps_json" | jq -r 'keys[]' 2>/dev/null | while read -r dep_name; do
        [[ -z "$dep_name" ]] && continue
        local dep_json
        dep_json=$(echo "$deps_json" | jq ".\"$dep_name\"")
        process_dep "$crate_toml" "$dep_name" "$dep_json"
    done
}

check_target_deps() {
    local crate_toml="$1"
    local targets
    targets=$(tomlq -r '.target // {} | keys[]' "$crate_toml" 2>/dev/null || true)
    
    for target in $targets; do
        [[ -z "$target" ]] && continue
        local deps
        deps=$(tomlq -r ".target.\"$target\" | (.dependencies // {}) + (.[\"dev-dependencies\"] // {}) + (.[\"build-dependencies\"] // {}) | keys[]" "$crate_toml" 2>/dev/null || true)
        for dep in $deps; do
            [[ -n "$dep" ]] && warn "skipping target-specific '$dep' in $crate_toml [$target]"
        done
    done
}

# ============================================================
# SURGICAL EDITING FUNCTIONS
# ============================================================

# Append deps to [workspace.dependencies] section
append_workspace_deps() {
    local ws_toml="$1"
    shift
    local deps_to_add=("$@")
    
    [[ ${#deps_to_add[@]} -eq 0 ]] && return 0
    
    # Check if [workspace.dependencies] section exists
    if ! grep -q '^\[workspace\.dependencies\]' "$ws_toml"; then
        echo "" >> "$ws_toml"
        echo "[workspace.dependencies]" >> "$ws_toml"
    fi
    
    # Find line number of [workspace.dependencies]
    local section_line
    section_line=$(grep -n '^\[workspace\.dependencies\]' "$ws_toml" | cut -d: -f1)
    
    # Find the next section (or end of file)
    local next_section_line
    next_section_line=$(tail -n +"$((section_line + 1))" "$ws_toml" | grep -n '^\[' | head -1 | cut -d: -f1)
    
    local insert_line
    if [[ -n "$next_section_line" ]]; then
        insert_line=$((section_line + next_section_line - 1))
    else
        insert_line=$(wc -l < "$ws_toml")
    fi
    
    # Build the text to insert
    local insert_text=""
    for dep_entry in "${deps_to_add[@]}"; do
        insert_text="${insert_text}${dep_entry}"$'\n'
    done
    
    # Insert using head/tail
    if [[ -n "$insert_text" ]]; then
        head -n "$insert_line" "$ws_toml" > "$ws_toml.tmp"
        printf '%s' "$insert_text" >> "$ws_toml.tmp"
        tail -n +"$((insert_line + 1))" "$ws_toml" >> "$ws_toml.tmp" 2>/dev/null || true
        mv "$ws_toml.tmp" "$ws_toml"
        echo "$ws_toml" >> "$MODIFIED_FILES"
    fi
}

# Replace a dependency line in a crate Cargo.toml
replace_dep_with_workspace() {
    local crate_toml="$1"
    local dep_name="$2"
    local dep_json="$3"
    local section="$4"
    
    # Get preserved attributes
    local features="" optional="" default_features=""
    
    if ! echo "$dep_json" | jq -e 'type == "string"' &>/dev/null; then
        if echo "$dep_json" | jq -e '.features' &>/dev/null; then
            features=$(echo "$dep_json" | jq -c '.features' | sed 's/,/, /g')
        fi
        echo "$dep_json" | jq -e '.optional == true' &>/dev/null && optional="true"
        echo "$dep_json" | jq -e '.["default-features"] == false' &>/dev/null && default_features="false"
    fi
    
    # Build new value
    local new_value="{ workspace = true"
    [[ -n "$features" ]] && new_value="$new_value, features = $features"
    [[ -n "$optional" ]] && new_value="$new_value, optional = true"
    [[ -n "$default_features" ]] && new_value="$new_value, default-features = false"
    new_value="$new_value }"
    
    # Escape dep_name for regex
    local dep_name_escaped
    # shellcheck disable=SC2016 # Single quotes intentional - literal regex pattern
    dep_name_escaped=$(printf '%s' "$dep_name" | sed 's/[][\.*^$()+?{}|]/\\&/g')
    
    # Also create a pattern for quoted keys: "dep-name" or 'dep-name'
    local dep_pattern="${dep_name_escaped}"
    local dep_pattern_quoted="[\"']${dep_name_escaped}[\"']"
    
    # Pattern 1: Simple string - dep = "version" (with optional quotes around key)
    if grep -qE "^${dep_pattern}[[:space:]]*=[[:space:]]*\"" "$crate_toml"; then
        sed -i "s/^${dep_pattern}[[:space:]]*=[[:space:]]*\"[^\"]*\"/${dep_name} = ${new_value}/" "$crate_toml"
        echo "$crate_toml" >> "$MODIFIED_FILES"
        return 0
    fi
    if grep -qE "^${dep_pattern_quoted}[[:space:]]*=[[:space:]]*\"" "$crate_toml"; then
        sed -i "s/^${dep_pattern_quoted}[[:space:]]*=[[:space:]]*\"[^\"]*\"/${dep_name} = ${new_value}/" "$crate_toml"
        echo "$crate_toml" >> "$MODIFIED_FILES"
        return 0
    fi
    
    # Pattern 2: Inline table - dep = { ... } (with optional quotes around key)
    if grep -qE "^${dep_pattern}[[:space:]]*=[[:space:]]*\{" "$crate_toml"; then
        sed -i "s/^${dep_pattern}[[:space:]]*=[[:space:]]*{[^}]*}/${dep_name} = ${new_value}/" "$crate_toml"
        echo "$crate_toml" >> "$MODIFIED_FILES"
        return 0
    fi
    if grep -qE "^${dep_pattern_quoted}[[:space:]]*=[[:space:]]*\{" "$crate_toml"; then
        sed -i "s/^${dep_pattern_quoted}[[:space:]]*=[[:space:]]*{[^}]*}/${dep_name} = ${new_value}/" "$crate_toml"
        echo "$crate_toml" >> "$MODIFIED_FILES"
        return 0
    fi
    
    # Pattern 3: Multi-line table [section.dep_name]
    local section_patterns=("dependencies" "dev-dependencies" "build-dependencies")
    for sec in "${section_patterns[@]}"; do
        local table_header_escaped
        # shellcheck disable=SC2016 # Single quotes intentional - literal regex pattern
        table_header_escaped=$(printf '%s' "[$sec.$dep_name]" | sed 's/[][\.*^$()+?{}|]/\\&/g')
        
        if grep -q "^${table_header_escaped}$" "$crate_toml"; then
            local start_line end_line
            start_line=$(grep -n "^${table_header_escaped}$" "$crate_toml" | cut -d: -f1)
            
            end_line=$(tail -n +"$((start_line + 1))" "$crate_toml" | grep -n '^\[' | head -1 | cut -d: -f1)
            
            if [[ -n "$end_line" ]]; then
                end_line=$((start_line + end_line - 1))
            else
                end_line=$(wc -l < "$crate_toml")
            fi
            
            # Step 1: Delete the multi-line table entirely
            head -n "$((start_line - 1))" "$crate_toml" > "$crate_toml.tmp"
            tail -n +"$((end_line + 1))" "$crate_toml" >> "$crate_toml.tmp" 2>/dev/null || true
            mv "$crate_toml.tmp" "$crate_toml"
            
            # Step 2: Find or create [section] and insert inline dep there
            local sec_escaped
            # shellcheck disable=SC2016 # Single quotes intentional - literal regex pattern
            sec_escaped=$(printf '%s' "[$sec]" | sed 's/[][\.*^$()+?{}|]/\\&/g')
            
            if grep -q "^${sec_escaped}$" "$crate_toml"; then
                # Insert right after [section] header
                local sec_line
                sec_line=$(grep -n "^${sec_escaped}$" "$crate_toml" | cut -d: -f1)
                head -n "$sec_line" "$crate_toml" > "$crate_toml.tmp"
                echo "${dep_name} = ${new_value}" >> "$crate_toml.tmp"
                tail -n +"$((sec_line + 1))" "$crate_toml" >> "$crate_toml.tmp" 2>/dev/null || true
                mv "$crate_toml.tmp" "$crate_toml"
            else
                # Section doesn't exist, create it at end of file
                {
                    echo ""
                    echo "[$sec]"
                    echo "${dep_name} = ${new_value}"
                } >> "$crate_toml"
            fi
            
            echo "$crate_toml" >> "$MODIFIED_FILES"
            return 0
        fi
    done
    
    # Pattern 4: Dotted keys - dep.version = "...", dep.path = "...", etc.
    if grep -qE "^${dep_name_escaped}\." "$crate_toml"; then
        # Delete all lines starting with dep_name.
        sed -i "/^${dep_name_escaped}\./d" "$crate_toml"
        
        # Find or create [section] and insert inline dep there
        local sec_escaped
        # shellcheck disable=SC2016 # Single quotes intentional - literal regex pattern
        sec_escaped=$(printf '%s' "[$section]" | sed 's/[][\.*^$()+?{}|]/\\&/g')
        
        if grep -q "^${sec_escaped}$" "$crate_toml"; then
            # Insert right after [section] header
            local sec_line
            sec_line=$(grep -n "^${sec_escaped}$" "$crate_toml" | cut -d: -f1)
            head -n "$sec_line" "$crate_toml" > "$crate_toml.tmp"
            echo "${dep_name} = ${new_value}" >> "$crate_toml.tmp"
            tail -n +"$((sec_line + 1))" "$crate_toml" >> "$crate_toml.tmp" 2>/dev/null || true
            mv "$crate_toml.tmp" "$crate_toml"
        else
            # Section doesn't exist, create it at end of file
            {
                echo ""
                echo "[$section]"
                echo "${dep_name} = ${new_value}"
            } >> "$crate_toml"
        fi
        
        echo "$crate_toml" >> "$MODIFIED_FILES"
        return 0
    fi
    
    warn "Could not locate '$dep_name' in $crate_toml for replacement"
}

# ============================================================
# MAIN LOGIC
# ============================================================

echo "Scanning crates..."

while read -r crate_toml; do
    [[ -z "$crate_toml" ]] && continue
    echo "  $crate_toml"
    
    # Check if tomlq can parse the file (fails on TOML 1.1 features like multi-line inline tables)
    if ! tomlq '.' "$crate_toml" &>/dev/null; then
        warn "could not fully parse $crate_toml (may contain TOML 1.1 features like multi-line inline tables)"
    fi
    
    process_section "$crate_toml" "dependencies"
    process_section "$crate_toml" "dev-dependencies"
    process_section "$crate_toml" "build-dependencies"
    check_target_deps "$crate_toml"
done < <(list_crate_tomls)

SKIP_LIST=$(sort -u "$SKIPPED" 2>/dev/null || true)

echo ""
echo "Adding to workspace Cargo.toml..."

# Build array of deps to add
declare -a ws_deps_to_add=()

while IFS=$'\t' read -r dep_name ws_value ws_type; do
    [[ -z "$dep_name" ]] && continue
    echo "$SKIP_LIST" | grep -qxF "$dep_name" 2>/dev/null && continue
    
    echo "  $dep_name = $ws_value"
    ws_deps_to_add+=("${dep_name} = ${ws_value}")
done < "$COLLECTED"

append_workspace_deps "Cargo.toml" "${ws_deps_to_add[@]}"

echo ""
echo "Updating crate Cargo.toml files..."

update_crate_section() {
    local crate_toml="$1" section="$2"
    local deps_json
    deps_json=$(tomlq ".\"$section\" // {}" "$crate_toml" 2>/dev/null)
    
    echo "$deps_json" | jq -r 'keys[]' 2>/dev/null | while read -r dep_name; do
        [[ -z "$dep_name" ]] && continue
        echo "$SKIP_LIST" | grep -qxF "$dep_name" 2>/dev/null && continue
        
        local dep_json
        dep_json=$(echo "$deps_json" | jq ".\"$dep_name\"")
        
        echo "$dep_json" | jq -e '.workspace == true' &>/dev/null && continue
        
        grep -q "^${dep_name}	" "$COLLECTED" 2>/dev/null || \
            echo "$EXISTING_WS_DEPS" | grep -qxF "$dep_name" || continue
        
        replace_dep_with_workspace "$crate_toml" "$dep_name" "$dep_json" "$section"
    done
}

while read -r crate_toml; do
    [[ -z "$crate_toml" ]] && continue
    echo "  $crate_toml"
    update_crate_section "$crate_toml" "dependencies"
    update_crate_section "$crate_toml" "dev-dependencies"
    update_crate_section "$crate_toml" "build-dependencies"
done < <(list_crate_tomls)

echo ""
success "Done! Run 'cargo check' to verify."