#!/usr/bin/env bash

# Define color variables
CSI="\e["
RESET="${CSI}0m"
BOLD_BLUE="${CSI}1;34m"
YELLOW="${CSI}0;33m"
GREEN="${CSI}0;32m"
MAGENTA="${CSI}0;35m"
CYAN="${CSI}0;36m"
DIM="${CSI}2m"

help() {
    local content=`cat << EOF
${BOLD_BLUE}tmuxss${RESET} - a simple tmux multiview session manager

${CYAN}Syntax:${RESET} tmuxss -<${YELLOW}k|a|i|c${RESET}> [-${YELLOW}g${RESET} group] [-${YELLOW}s${RESET} sid] [-${YELLOW}d${RESET}] [-${YELLOW}t${RESET} template] [-${YELLOW}p${RESET} path]

${BOLD_BLUE}Commands:${RESET}
  ${YELLOW}k${RESET}     Kill sessions associated with a SID or group
  ${YELLOW}a${RESET}     Attach to a session group with the given root SID
  ${YELLOW}i${RESET}     Interactive open mode
  ${YELLOW}c${RESET}     Create a session group in the current directory or a provided one, optionally using a template
  ${YELLOW}p${RESET}     A directory path to create the session in
  ${YELLOW}b${RESET}     Initialize the main session
  ${YELLOW}h${RESET}     Print the entire help menu

${BOLD_BLUE}Options:${RESET}
  ${YELLOW}s${RESET}     SID (Session ID)
  ${YELLOW}g${RESET}     Group name
  ${YELLOW}t${RESET}     Template name
  ${YELLOW}f${RESET}     Path to template config file
  ${YELLOW}d${RESET}     Stay detached from new session

EOF
`
printf "$content"
echo
echo
}

extended_help() {
    content=`cat << EOF
${BOLD_BLUE}Config Format:${RESET}
${CYAN}{${RESET}
  ${YELLOW}"default"${RESET}: ${GREEN}"<env-name>"${RESET}, ${DIM}# default: \$PWD${RESET}
  ${YELLOW}"envs"${RESET}: ${CYAN}{${RESET}
    ${GREEN}"<env-name>"${RESET}: ${CYAN}{${RESET}
      ${YELLOW}"path"${RESET}: ${GREEN}"<cwd>"${RESET}, ${DIM}# default: \$PWD${RESET}
      ${YELLOW}"group"${RESET}: ${GREEN}"<optional session-group>"${RESET}, ${DIM}# default: basename \$PWD${RESET}
      ${YELLOW}"focused"${RESET}: ${CYAN}<window-index>${RESET}, ${DIM}# default: 0${RESET}
      ${YELLOW}"windows"${RESET}: ${CYAN}[${RESET}
        ${CYAN}{${RESET}
          ${YELLOW}"name"${RESET}: ${GREEN}"<window-name>"${RESET}, ${DIM}# default: index${RESET}
          ${YELLOW}"path"${RESET}: ${GREEN}"<optional path>"${RESET}, ${DIM}# default: \$PWD${RESET}
          ${YELLOW}"command_run"${RESET}: ${GREEN}"<shell command>"${RESET}, ${DIM}# default: ""${RESET}
          ${YELLOW}"command_prepare"${RESET}: ${GREEN}"<setup command>"${RESET}, ${DIM}# default: ""${RESET}
          ${YELLOW}"read_only"${RESET}: ${MAGENTA}true${RESET} | ${MAGENTA}false${RESET}, ${DIM}# default: false${RESET}
          ${YELLOW}"hist_file"${RESET}: ${GREEN}"<path-to-hist-file>"${RESET}, ${DIM}# default: std Hist${RESET}
          ${YELLOW}"hsplit"${RESET}: ${CYAN}[ <pane> , ...]${RESET},
          ${YELLOW}"vsplit"${RESET}: ${CYAN}[ <pane> , ...]${RESET}
        ${CYAN}}, ...]${RESET}
    ${CYAN}}${RESET}
  ${CYAN}}${RESET}
${CYAN}}${RESET}

${BOLD_BLUE}Pane:${RESET}
${CYAN}{${RESET}
  ${YELLOW}"path"${RESET}: ${GREEN}"<optional path>"${RESET}, ${DIM}# default: \$PWD${RESET}
  ${YELLOW}"command_run"${RESET}: ${GREEN}"<shell command>"${RESET}, ${DIM}# default: ""${RESET}
  ${YELLOW}"command_prepare"${RESET}: ${GREEN}"<setup command>"${RESET}, ${DIM}# default: ""${RESET}
  ${YELLOW}"read_only"${RESET}: ${MAGENTA}true${RESET} | ${MAGENTA}false${RESET}, ${DIM}# default: false${RESET}
  ${YELLOW}"hist_file"${RESET}: ${GREEN}"<path-to-hist-file>"${RESET}, ${DIM}# default: std Hist${RESET}
  ${YELLOW}"hsplit"${RESET}?: ${CYAN}[ <pane> , ...]${RESET},
  ${YELLOW}"vsplit"${RESET}?: ${CYAN}[ <pane> , ...]${RESET}
${CYAN}}${RESET}

${BOLD_BLUE}Notes:${RESET}
  - All paths are relative to the environment path unless absolute
  - Only either hsplit or vsplit may be defined

${BOLD_BLUE}Automatic Template Resolution (resolve key):${RESET}
${CYAN}[${RESET}
  ${CYAN}{${RESET}
    ${YELLOW}"path"${RESET}: ${GREEN}"<directory path>"${RESET}, ${DIM}# Path to search for this template${RESET}
    ${YELLOW}"template"${RESET}: ${GREEN}"<env-name>"${RESET}, ${DIM}# Template to automatically select if path matches${RESET}
  ${CYAN}}, ...${RESET}
${CYAN}]${RESET}

${BOLD_BLUE}Notes:${RESET}
  - Paths can be absolute or relative to \$HOME
  - When creating a session, tmuxss will check if the current path matches any 'resolve' entry
  - If a match is found, the corresponding template is selected automatically
  - Multiple entries are checked in order; first match wins

${BOLD_BLUE}Example:${RESET}
${CYAN}[${RESET}
  ${CYAN}{${RESET} "path": "~/src/autopurger", "template": "autopurger" ${CYAN}},${RESET}
  ${CYAN}{${RESET} "path": "~/src/project-x", "template": "projectx" ${CYAN}}${RESET}
${CYAN}]${RESET}

${BOLD_BLUE}Examples:${RESET}
  ${YELLOW}Basic layout:${RESET}
    ${CYAN}{${RESET}
      ${YELLOW}"default"${RESET}: ${GREEN}"default"${RESET},
      ${YELLOW}"envs"${RESET}: ${CYAN}{${RESET}
        ${GREEN}"default"${RESET}: ${CYAN}{${RESET}
          ${YELLOW}"path"${RESET}: ${GREEN}"."${RESET},
          ${YELLOW}"focused"${RESET}: ${CYAN}0${RESET},
          ${YELLOW}"windows"${RESET}: ${CYAN}[${RESET}
            ${CYAN}{${RESET} ${YELLOW}"name"${RESET}: ${GREEN}"code"${RESET}, ${YELLOW}"command_run"${RESET}: ${GREEN}"nvim"${RESET} ${CYAN}},${RESET}
            ${CYAN}{${RESET} ${YELLOW}"name"${RESET}: ${GREEN}"util"${RESET}, ${YELLOW}"command_run"${RESET}: ${GREEN}"htop"${RESET} ${CYAN}}${RESET}
          ${CYAN}]${RESET}
        ${CYAN}}${RESET}
      ${CYAN}}${RESET}
    ${CYAN}}${RESET}

  ${YELLOW}Nested splits with preparation:${RESET}
    ${CYAN}{${RESET}
      ${YELLOW}"envs"${RESET}: ${CYAN}{${RESET}
        ${GREEN}"dev"${RESET}: ${CYAN}{${RESET}
          ${YELLOW}"path"${RESET}: ${GREEN}"~/src/project"${RESET},
          ${YELLOW}"windows"${RESET}: ${CYAN}[${RESET}
            ${CYAN}{${RESET} ${YELLOW}"name"${RESET}: ${GREEN}"code"${RESET}, ${YELLOW}"command_run"${RESET}: ${GREEN}"nvim"${RESET} ${CYAN}},${RESET}
            ${CYAN}{${RESET}
              ${YELLOW}"name"${RESET}: ${GREEN}"runtime"${RESET},
              ${YELLOW}"hsplit"${RESET}: ${CYAN}[${RESET}
                ${CYAN}{${RESET}
                  ${YELLOW}"vsplit"${RESET}: ${CYAN}[${RESET}
                    ${CYAN}{${RESET} ${YELLOW}"command_prepare"${RESET}: ${GREEN}"yarn start"${RESET}, ${YELLOW}"path"${RESET}: ${GREEN}"./frontend"${RESET} ${CYAN}},${RESET}
                    ${CYAN}{${RESET} ${YELLOW}"command_prepare"${RESET}: ${GREEN}"cargo run"${RESET}, ${YELLOW}"path"${RESET}: ${GREEN}"./backend"${RESET} ${CYAN}}${RESET}
                  ${CYAN}]${RESET}
                ${CYAN}},${RESET}
                ${CYAN}{${RESET} ${YELLOW}"command_run"${RESET}: ${GREEN}"htop"${RESET} ${CYAN}}${RESET}
              ${CYAN}]${RESET}
            ${CYAN}}${RESET}
          ${CYAN}]${RESET}
        ${CYAN}}${RESET}
      ${CYAN}}${RESET}
    ${CYAN}}${RESET}

EOF
`
printf "$content"
echo
}

resolve_path() {
    local base=$1
    local pair=$2

    [[ -z $pair ]] && pair="."
    [[ $pair == ~* ]] && pair="${pair/#\~/$HOME}"

    if [[ $pair = /* ]]; then
        realpath "$pair"
    else
        realpath "$base/$pair"
    fi
}

SID=""
GROUP=""
SUBCOMMAND=""
TEMPLATE=""
DETACH=0
BASE_PATH=$PWD
YQ_AVAILABLE=$(command -v yq >/dev/null 2>&1 && echo true || echo false)
TEMPLATE_SOURCE="${XDG_DATA_HOME:-$HOME/.config}/tmux/tmuxss.json"

while getopts "p:f:g:t:s:bhkaicd" option; do
    case $option in
    h)
        help
        extended_help
        exit 0
        ;;
    s)
        if [[ $OPTARG =~ "#" ]]; then
            echo "\"#\" is a reserved character"
            exit 3
        fi
        SID=$OPTARG ;;
    g)
        if [[ $OPTARG =~ "#" ]]; then
            echo "\"#\" is a reserved character"
            exit 3
        fi
        GROUP=$OPTARG ;;
    t)
        if [[ ! $OPTARG =~ ^[a-zA-Z]+$ ]]; then
            echo "Template names must be only upper and lowercase."
            exit 3
        fi
        # GROUP=$OPTARG
        TEMPLATE=$OPTARG ;;
    f) 
        if [[ ! $OPTARG =~ ^(/)?([^/\0]+(/)?)+$ ]]; then
            echo "Template source must be a valid path."
            exit 3
        fi
        TEMPLATE_SOURCE=$OPTARG ;;
    p) 
        if [[ ! $OPTARG =~ ^(/)?([^/\0]+(/)?)+$ ]]; then
            echo "Path must be valid."
            exit 3
        fi
        BASE_PATH=$(resolve_path $BASE_PATH $OPTARG);;
    i)
        SUBCOMMAND="i" ;;
    a)
        SUBCOMMAND="a" ;;
    k)
        SUBCOMMAND="k" ;;
    b)
        SUBCOMMAND="b" ;;
    c)
        SUBCOMMAND="c" ;;
    d)
        DETACH=1 ;;
    ?)
        echo "Invalid option"
        echo "Pull up the help menu with tmuxss -h"
        exit 3
        ;;
    esac
done

ATTACHED_TO=""
SIDINFERED=0

infer() {
if [[ -z $SID ]]; then
    if [[ -n "$TMUX" ]]; then
        local SESSION=$(tmux display-message -p '#S')

        ATTACHED_TO=$SESSION

        SID=${SESSION#*#}

        if [[ -z $SID ]]; then
            echo "No ID found in the current tmux session name: $SESSION"
            exit 66
        fi

        SIDINFERED=1
    fi
fi
}
infer

attachSession() {
    local group=$1
    local sid=$2
    local session="$group#$sid"

    # Check if a session exists in the group with the current SID
    if [[ -z $(tmux list-sessions -F "#{session_name}" | grep "^$session$") ]]; then
        # If no session exists with the current SID, create a new one
        tmux new-session -ds "$session" -t "$group"
    fi

    if [[ $TMUX ]]; then
        tmux switch-client -t "$session"
    else
        tmux attach -t "$session"
    fi
}

sanitize_path() {
    local input=$1

    local sanitized=$(echo "$input" | tr -d '.' | sed 's/[^a-zA-Z0-9]/_/g' | tr '[:upper:]' '[:lower:]')

    echo "$sanitized"
}

resolve_template_from_path() {
    # explicit -t always wins
    [[ -n $TEMPLATE ]] && return

    local best_template=""
    local best_len=0
    local i=0

    while :; do
        local path_var="resolve_${i}_path"
        local tmpl_var="resolve_${i}_template"

        # stop when index no longer exists in parsed DATA
        grep -q "^${path_var}=" <<< "$DATA" || break

        local raw_path=${!path_var}
        local tmpl=${!tmpl_var}

        # normalize candidate path
        local abs_path
        abs_path=$(resolve_path "$HOME" "$raw_path")

        # prefix match
        if [[ $BASE_PATH == "$abs_path"* ]]; then
            local len=${#abs_path}
            if (( len > best_len )); then
                best_len=$len
                best_template=$tmpl
            fi
        fi

        ((i++))
    done

    [[ -n $best_template ]] && TEMPLATE="$best_template"
}

# This gets the last bit of performance
setup_pane() {
    local key="$1"
    local pane_id="$2"

    if [[ $SHELL != "/bin/fish" ]]; then
        local hist_file="${key}_hist_file"
        hist_file=${!hist_file}
        # Set HISTFILE only if a specific file was provided
        if [[ -n $hist_file ]]; then
            hist_file=$(resolve_path "$BASE_PATH" "$hist_file")
            tmux send-keys -t "$pane_id" "export HISTFILE=\"$hist_file\"; export PROMPT_COMMAND='history -a; history -c; history -r'; history -d \$(history 1)" C-m
        fi

        local read_only="${key}_read_only"
        read_only=${!read_only}
        if [[ $read_only == "true" ]]; then
            tmux send-keys -t "$pane_id" "shopt -ou history; history -d \$(history 1)" C-m
        fi

        tmux send-keys -t "$pane_id" "clear; history -d \$(history 1)" C-m
        tmux clear-history -t "$pane_id"
    fi

    local command_run="${key}_command_run"
    command_run=${!command_run}
    [[ -n $command_run ]] && tmux send-keys -t "$pane_id" "$command_run" C-m
    local command_prepare="${key}_command_prepare"
    command_prepare=${!command_prepare}
    [[ -n $command_prepare ]] && tmux send-keys -t "$pane_id" "$command_prepare"
}

visit_pane() {
    local key="$1"
    local window_index="$2"
    local pane_index="${PANE_INDEX[$window_index]:-0}"

    local path_var="${key}_path"
    local path="${!path_var}"
    path=$(resolve_path "$BASE_PATH" "$path")

    local pane_id="$GROUP:$window_index.$pane_index"

    local hsplit="${key}_hsplit"
    local vsplit="${key}_vsplit"

    if grep -q "$hsplit" <<< "$DATA"; then
        tmux split-window -h -t "$pane_id" -c "$path"

        local i=0
        while grep -q "${hsplit}_${i}" <<< "$DATA"; do
            visit_pane "${hsplit}_${i}" "$window_index" "$pane_index"
            ((i++))
        done

    elif grep -q "$vsplit" <<< "$DATA"; then
        tmux split-window -v -t "$pane_id" -c "$path"

        local i=0
        while grep -q "${vsplit}_${i}" <<< "$DATA"; do
            visit_pane "${vsplit}_${i}" "$window_index" "$pane_index"
            ((i++))
        done

    else
        # Ensure the pane exists at the correct cwd before setup
        tmux respawn-pane -k -t "$pane_id" -c "$path"

        setup_pane "$key" "$pane_id" &
        PANE_INDEX[$window_index]=$((pane_index + 1))
    fi
}

DATA=""
declare -A PANE_INDEX

build_template() {
    local path="${template_key}_path"
    path=${!path}
    
    BASE_PATH=$(resolve_path $BASE_PATH $path)

    cd "$BASE_PATH" && tmux new-session -ds "$GROUP" -c "$BASE_PATH"

    local i=0
    while :; do
       local key="${template_key}_windows_${i}"
        if [ $i -gt 0 ] && ! grep -q "$key" <<< "$DATA"; then break; fi

        local name="${key}_name"
        name=${!name:-$i}
        [[ $i -gt 0 ]] && tmux new-window -d -t "$GROUP" -c "$BASE_PATH" -n "$name"

        visit_pane "$key" "$i" &
        i=$((i + 1))
    done
    wait

    local focused_window="${template_key}_focused"
    focused_window=${!focused_window}
    [[ -z $focused_window ]] && focused_window="0"

    tmux select-window -t "$GROUP:$focused_window"
}

case $SUBCOMMAND in
k)
    if [[ -n $GROUP ]]; then
        SESSION_LIST=$(tmux list-sessions -F "#{session_name}")

        for SESSION in $SESSION_LIST; do
            if [[ $SESSION == "$GROUP#"* || $SESSION == $GROUP ]]; then
                if [[ "$GROUP#$SID" == $ATTACHED_TO && $GROUP != "main" ]]; then
                    attachSession main $SID
                fi

                tmux kill-session -t "$SESSION"
            fi
        done
    fi
    if [[ -n $SID && $SIDINFERED -eq 0 ]]; then
        SESSION_LIST=$(tmux list-sessions -F "#{session_name}")

        for SESSION in $SESSION_LIST; do
            if [[ $SESSION == *#$SID ]]; then
                tmux kill-session -t "$SESSION"
            fi
        done
    fi

    if [[ ( -n $SID && $SIDINFERED -eq 0 ) || -n $GROUP ]]; then
        exit 0
    fi
    ;;
a)
    if [[ -z $GROUP ]]; then
        echo "No group specified"
        echo "Check out tmuxss -h"
        exit 3
    fi
    [[ -z $SID ]] && SID=$$
    attachSession $GROUP $SID
    exit 0
    ;;
i)
    file=$(<"$TEMPLATE_SOURCE")
    if [[ -n $file ]] && $YQ_AVAILABLE; then
        DATA=$(yq -p=json -o shell <<< "$file") || {
            echo "Invalid template file"
            exit 3
        }
    fi

    eval "$DATA"

    SEARCH_PATHS=()
    mapfile -t SEARCH_PATHS < <(yq -r '.paths[]?' <<< "$file")

    for i in "${!SEARCH_PATHS[@]}"; do
        p="${SEARCH_PATHS[$i]}"
        [[ $p == ~* ]] && p="${p/#\~/$HOME}"
        [[ $p != /* ]] && p="$PWD/$p"
        SEARCH_PATHS[$i]="$p"
    done

    mapfile -d '' DIRS < <(
      for path in "${SEARCH_PATHS[@]}"; do
        [[ -d "$path" ]] || continue
        find "$path" -mindepth 1 -maxdepth 1 -type d -print0
      done
    )

    [[ ${#DIRS[@]} -eq 0 ]] && { echo "No directories found in SEARCH_PATHS"; exit 0; }

    selected=$(printf '%s\n' "${DIRS[@]}" | fzf --height 40% --reverse)

    [[ -z "${selected:-}" ]] && exit 0

    selected=${selected%$'\0'}

    exec tmuxss -c -p "$selected"
    ;;
b)
    [[ -z $SID ]] && SID=$$

    tmuxss -c -d -t "main" -g "main" -s "$SID"

    trap "tmuxss -k -s $SID" EXIT SIGHUP && tmuxss -a -g "main" -s "$SID"
    exit 0
    ;;
c)
    file=$(<"$TEMPLATE_SOURCE")
    if [[ -n $file ]] && $YQ_AVAILABLE; then
        DATA=$(yq -p=json -o shell <<< "$file") || {
            echo "Invalid template file"
            exit 3
        }
    fi
    if [[ -z $DATA ]]; then
        if [[ -n $TEMPLATE && $TEMPLATE != "main" ]]; then
            ! $YQ_AVAILABLE && { echo "Using templates requires yq to be installed"; exit 1; }
            echo "Could not locate config file"
            exit 1
        fi
        DATA="default='default' envs_default_path='.' envs_main_path='~' envs_main_group='main'"
    fi

    eval "$DATA"

    resolve_template_from_path

    if [[ -z $TEMPLATE ]]; then
        default_template="default"
        default_template=${!default_template}
        [[ -n $default_template ]] && TEMPLATE=$default_template
    fi

    template_key="envs_${TEMPLATE}"
    if ! grep -q "$template_key" <<< "$DATA"; then { echo "Session template not found"; exit 3; }; fi

    template_group="${template_key}_group"
    template_group=${!template_group}
    [[ -n $template_group ]] && GROUP=$template_group

    if [[ -z $GROUP ]]; then
        GROUP=$(sanitize_path "$(basename "$BASE_PATH")")

        # If the current directory is root, set GROUP to "root"
        if [[ $GROUP == "/" ]]; then
            GROUP="root"
        fi
    fi
    [[ -z $SID ]] && SID=$$

    if [[ -z $(tmux list-sessions -F "#{session_name}" | grep "^$GROUP$") ]]; then
        build_template
        # nohup build_template >/dev/null 2>&1 &
    fi
    
    if [[ $DETACH -eq 0 ]]; then
        attachSession $GROUP $SID
    fi 
    exit 0
    ;;
esac


if [[ -n $TMUX ]]; then 
    if [[ -z $SUBCOMMAND || $SUBCOMMAND == "a" ]]; then
        SUBCOMMAND="a"
        TITLE="Select to attach"
    fi

    if [[ $SUBCOMMAND == "k" ]]; then
        TITLE="Select to kill"
    fi

    ITEMS=""
    index=1
    SESSION_GROUPS=$(tmux list-sessions -F "#{session_name}" | awk -F '#' '{print $1}' | sort -u)
    while read -r group; do
    if [[ -n "$group" ]]; then
        ITEMS+="$group $index 'run-shell \"tmuxss -$SUBCOMMAND -g \"$group\"\"' "
        ((index++))
    fi
    done <<< "$SESSION_GROUPS"

    if [[ -n $ITEMS ]]; then
        eval "tmux display-menu -T \"$TITLE\" $ITEMS"
    else
        tmux display-message "No session groups found."
    fi
else
    help
fi