#!/bin/bash
# Bachelor Thesis Work Tracker
# Monitors active windows for thesis-related work (Unreal Engine, Unity, Nvidia Omniverse, VS Code with specific repo)
# Unlocks Steam and other distractions only after sufficient work time is accumulated
#
# This daemon runs continuously and:
# 1. Tracks active window time for approved thesis work applications
# 2. Maintains a protected state file with accumulated work time
# 3. Manages hosts file blocking/unblocking based on work quota
# 4. Provides psychological friction against circumvention

set -euo pipefail

# Configuration
# shellcheck disable=SC2034  # SCRIPT_DIR reserved for future use
SCRIPT_DIR="$(dirname "$(readlink -f "$0")")"
STATE_DIR="/var/lib/thesis-work-tracker"
STATE_FILE="$STATE_DIR/work-time.state"
LOCK_FILE="$STATE_DIR/tracker.lock"
LOG_DIR="/var/log/thesis-work-tracker"
LOG_FILE="$LOG_DIR/tracker.log"
CHECK_INTERVAL=15 # Check every 15 seconds
PROC_ROOT="${PROC_ROOT:-/proc}"
HOSTS_FILE="${HOSTS_FILE:-/etc/hosts}"

# Work requirements (in seconds)
# 2 hours of work = 7200 seconds required before Steam access
WORK_QUOTA_REQUIRED=7200 # 2 hours
WORK_DECAY_PER_HOUR=1800 # Lose 30 minutes per hour of Steam usage

# Thesis work applications - process names and window patterns
# These are the applications that count as "thesis work"
declare -A THESIS_APPS=(
	["UnrealEditor"]="Unreal Engine"
	["UE4Editor"]="Unreal Engine 4"
	["UE5Editor"]="Unreal Engine 5"
	["Unity"]="Unity Editor"
	["UnityHub"]="Unity Hub"
	["Code"]="Visual Studio Code" # Special handling for repo check
	["code"]="Visual Studio Code" # lowercase variant
	["omniverse"]="Nvidia Omniverse"
	["kit"]="Nvidia Omniverse Kit"
)

# VS Code specific repo to track
VSCODE_REQUIRED_REPO="praca_magisterska"

# Steam and distraction patterns for hosts blocking
STEAM_DOMAINS=(
	"steampowered.com"
	"steamcommunity.com"
	"steamgames.com"
	"store.steampowered.com"
	"steamcdn-a.akamaihd.net"
	"steamstatic.com"
	"steamusercontent.com"
)

# Additional distraction sites that should be blocked
DISTRACTION_DOMAINS=(
	"reddit.com"
	"twitter.com"
	"x.com"
	"facebook.com"
	"instagram.com"
	"youtube.com"
	"twitch.tv"
	"9gag.com"
	"imgur.com"
)

# Colors for logging
# shellcheck disable=SC2034  # Colors available for log formatting
RED='\033[0;31m'
# shellcheck disable=SC2034
GREEN='\033[0;32m'
# shellcheck disable=SC2034
YELLOW='\033[0;33m'
# shellcheck disable=SC2034
BLUE='\033[0;34m'
# shellcheck disable=SC2034
CYAN='\033[0;36m'
# shellcheck disable=SC2034
BOLD='\033[1m'
# shellcheck disable=SC2034
NC='\033[0m' # No Color

# Logging function
log_message() {
	local level="$1"
	shift
	local message="$*"
	local timestamp
	printf -v timestamp '%(%Y-%m-%d %H:%M:%S)T' -1
	local formatted
	formatted="[${timestamp}] [${level}] ${message}"
	printf '%s\n' "$formatted" >&2
	printf '%s\n' "$formatted" >> "$LOG_FILE" 2>/dev/null || true
}

log_info() { log_message "INFO" "$@"; }
log_warn() { log_message "WARN" "$@"; }
log_error() { log_message "ERROR" "$@"; }
log_debug() {
	if [[ ${DEBUG:-0} -eq 1 ]]; then
		log_message "DEBUG" "$@"
	fi
}

wait_seconds() {
	local timeout_s=$1
	local start_ts end_ts elapsed_s remaining_s

	printf -v start_ts '%(%s)T' -1
	IFS= read -r -t "$timeout_s" || true
	printf -v end_ts '%(%s)T' -1

	elapsed_s=$((end_ts - start_ts))
	if (( elapsed_s < timeout_s )); then
		remaining_s=$((timeout_s - elapsed_s))
		sleep "$remaining_s"
	fi
}

# Initialize directories and state file
init_state() {
	# Create directories with proper permissions
	if [[ ! -d $STATE_DIR ]]; then
		sudo mkdir -p "$STATE_DIR"
		sudo chmod 700 "$STATE_DIR"
	fi

	if [[ ! -d $LOG_DIR ]]; then
		sudo mkdir -p "$LOG_DIR"
		sudo chmod 755 "$LOG_DIR"
	fi

	# Initialize state file if it doesn't exist
	if [[ ! -f $STATE_FILE ]]; then
		local now_iso now_epoch
		printf -v now_iso '%(%Y-%m-%d %H:%M:%S)T' -1
		printf -v now_epoch '%(%s)T' -1
		sudo bash -c "cat > '$STATE_FILE'" <<EOF
# Thesis Work Tracker State File
# DO NOT EDIT MANUALLY - Managed by thesis_work_tracker daemon
# Last updated: ${now_iso}

TOTAL_WORK_SECONDS=0
LAST_UPDATE_TIMESTAMP=${now_epoch}
STEAM_ACCESS_GRANTED=0
LAST_WORK_SESSION_START=0
CURRENT_SESSION_SECONDS=0
EOF
		sudo chmod 600 "$STATE_FILE"
		if ! sudo chattr +i "$STATE_FILE" 2>/dev/null; then
			log_warn "Failed to set immutable flag on state file - protections may be weaker"
		fi
	fi
}

# Load current state from file
load_state() {
	if [[ ! -f $STATE_FILE ]]; then
		log_error "State file not found: $STATE_FILE"
		return 1
	fi

	# Temporarily remove immutable flag to read
	sudo chattr -i "$STATE_FILE" 2>/dev/null || true

	# Parse state file safely without using source or external text helpers.
	TOTAL_WORK_SECONDS=0
	STEAM_ACCESS_GRANTED=0
	CURRENT_SESSION_SECONDS=0
	LAST_WORK_SESSION_START=0

	local key value
	while IFS='=' read -r key value; do
		case $key in
		TOTAL_WORK_SECONDS)
			TOTAL_WORK_SECONDS=$value
			;;
		STEAM_ACCESS_GRANTED)
			STEAM_ACCESS_GRANTED=$value
			;;
		CURRENT_SESSION_SECONDS)
			CURRENT_SESSION_SECONDS=$value
			;;
		LAST_WORK_SESSION_START)
			LAST_WORK_SESSION_START=$value
			;;
		esac
	done < "$STATE_FILE"

	# Validate that values are numeric
	if ! [[ $TOTAL_WORK_SECONDS =~ ^[0-9]+$ ]]; then TOTAL_WORK_SECONDS=0; fi
	if ! [[ $STEAM_ACCESS_GRANTED =~ ^[01]$ ]]; then STEAM_ACCESS_GRANTED=0; fi
	if ! [[ $CURRENT_SESSION_SECONDS =~ ^[0-9]+$ ]]; then CURRENT_SESSION_SECONDS=0; fi
	if ! [[ $LAST_WORK_SESSION_START =~ ^[0-9]+$ ]]; then LAST_WORK_SESSION_START=0; fi

	# Re-apply immutable flag
	sudo chattr +i "$STATE_FILE" 2>/dev/null || true
}

# Save current state to file
save_state() {
	local total_work="$1"
	local steam_access="$2"
	local current_session="$3"
	local session_start="$4"

	# Remove immutable flag
	sudo chattr -i "$STATE_FILE" 2>/dev/null || true

	local now_iso now_epoch
	printf -v now_iso '%(%Y-%m-%d %H:%M:%S)T' -1
	printf -v now_epoch '%(%s)T' -1

	if [[ -w $STATE_FILE ]]; then
		cat <<EOF > "$STATE_FILE"
# Thesis Work Tracker State File
# DO NOT EDIT MANUALLY - Managed by thesis_work_tracker daemon
# Last updated: ${now_iso}

TOTAL_WORK_SECONDS=$total_work
LAST_UPDATE_TIMESTAMP=${now_epoch}
STEAM_ACCESS_GRANTED=$steam_access
LAST_WORK_SESSION_START=$session_start
CURRENT_SESSION_SECONDS=$current_session
EOF
	else
		# Non-writable path: write in temp and copy with sudo.
		local temp_state
		temp_state=$(mktemp) || {
			log_error "Failed to create temporary state file"
			return 1
		}

		cat <<EOF > "$temp_state"
# Thesis Work Tracker State File
# DO NOT EDIT MANUALLY - Managed by thesis_work_tracker daemon
# Last updated: ${now_iso}

TOTAL_WORK_SECONDS=$total_work
LAST_UPDATE_TIMESTAMP=${now_epoch}
STEAM_ACCESS_GRANTED=$steam_access
LAST_WORK_SESSION_START=$session_start
CURRENT_SESSION_SECONDS=$current_session
EOF

		sudo cp "$temp_state" "$STATE_FILE"
		rm -f "$temp_state"
	fi

	sudo chmod 600 "$STATE_FILE"
	# Re-apply immutable flag
	if ! sudo chattr +i "$STATE_FILE" 2>/dev/null; then
		log_warn "Failed to set immutable flag on state file after save"
	fi
}

# Get active window title and process name
get_active_window_info() {
	if ! command -v xdotool &>/dev/null; then
		log_error "xdotool not installed, cannot detect active window"
		return 1
	fi

	local active_window_id
	active_window_id=$(xdotool getactivewindow 2>/dev/null || echo "")

	if [[ -z $active_window_id ]]; then
		return 1
	fi

	local window_name

	local window_pid
	window_pid=$(xdotool getwindowpid "$active_window_id" 2>/dev/null || echo "")

	local process_name=""
	if [[ -n $window_pid ]]; then
		if [[ -r "$PROC_ROOT/$window_pid/comm" ]]; then
			read -r process_name < "$PROC_ROOT/$window_pid/comm" || process_name=""
		else
			process_name=$(ps -p "$window_pid" -o comm= 2>/dev/null || echo "")
		fi
	fi

	window_name=""
	if [[ $process_name == "Code" || $process_name == "code" ]]; then
		window_name=$(xdotool getwindowname "$active_window_id" 2>/dev/null || echo "")
	fi

	echo "${process_name}|${window_name}"
}

# Check if VS Code is working on the required repository
is_vscode_on_thesis_repo() {
	local window_title="$1"

	# VS Code window titles typically contain the folder/workspace name
	# Look for the repo name in the window title
	# Window title format is usually: "filename - reponame - Visual Studio Code"
	if [[ $window_title == *"$VSCODE_REQUIRED_REPO"* ]]; then
		return 0
	fi

	return 1
}

# Check if current active window is thesis work
is_thesis_work_active() {
	local window_info
	window_info=$(get_active_window_info)

	if [[ -z $window_info ]]; then
		return 1
	fi

	local process_name
	local window_title
	IFS='|' read -r process_name window_title <<<"$window_info"

	log_debug "Active window: process='$process_name' title='$window_title'"

	# Check each thesis application
	for proc_pattern in "${!THESIS_APPS[@]}"; do
		local app_name="${THESIS_APPS[$proc_pattern]}"

		# Check window title for application name (more reliable than process name)
		if [[ $window_title == *"$app_name"* ]]; then
			# Special handling for VS Code - must be on thesis repo
			if [[ $proc_pattern == "Code" ]] || [[ $proc_pattern == "code" ]]; then
				if is_vscode_on_thesis_repo "$window_title"; then
					log_debug "Thesis work detected: VS Code on $VSCODE_REQUIRED_REPO"
					return 0
				else
					log_debug "VS Code detected but not on thesis repo"
					continue
				fi
			fi

			log_debug "Thesis work detected: $app_name"
			return 0
		fi

		# Also check process name with exact match
		if [[ $process_name == "$proc_pattern" ]]; then
			# Special handling for VS Code - must be on thesis repo
			if [[ $proc_pattern == "Code" ]] || [[ $proc_pattern == "code" ]]; then
				if is_vscode_on_thesis_repo "$window_title"; then
					log_debug "Thesis work detected: VS Code on $VSCODE_REQUIRED_REPO"
					return 0
				else
					log_debug "VS Code detected but not on thesis repo"
					continue
				fi
			fi

			log_debug "Thesis work detected: $app_name"
			return 0
		fi
	done

	return 1
}

# Block Steam and distractions in /etc/hosts
block_distractions() {
	log_info "Blocking Steam and distractions in /etc/hosts"

	# Remove immutable flag temporarily
	sudo chattr -i "$HOSTS_FILE" 2>/dev/null || true

	# Scan the file once to build a set of already-blocked domains (no grep fork)
	local -A blocked_set=()
	local scan_line scan_domain
	while IFS= read -r scan_line; do
		if [[ $scan_line == "0.0.0.0 "* || $scan_line == $'0.0.0.0\t'* ]]; then
			read -r _ scan_domain <<<"$scan_line" 2>/dev/null || true
			blocked_set[$scan_domain]=1
		fi
	done < "$HOSTS_FILE"

	# Collect entries not yet present
	local new_entries=() domain
	for domain in "${STEAM_DOMAINS[@]}" "${DISTRACTION_DOMAINS[@]}"; do
		if [[ -z ${blocked_set[$domain]+x} ]]; then
			new_entries+=("0.0.0.0 $domain")
		fi
	done

	if (( ${#new_entries[@]} > 0 )); then
		printf '%s\n' "${new_entries[@]}" | sudo bash -c "cat >> '$HOSTS_FILE'"
		log_info "Added distraction blocks to /etc/hosts"
	fi

	# Re-apply immutable flag
	sudo chattr +i "$HOSTS_FILE" 2>/dev/null || true
}

# Unblock Steam and distractions from /etc/hosts
unblock_distractions() {
	log_info "Unblocking Steam and distractions in /etc/hosts"

	# Remove immutable flag temporarily
	sudo chattr -i "$HOSTS_FILE" 2>/dev/null || true

	# Filter out blocked entries using bash (no sed/mktemp forks)
	local new_content="" hosts_line skip domain
	while IFS= read -r hosts_line; do
		skip=0
		for domain in "${STEAM_DOMAINS[@]}" "${DISTRACTION_DOMAINS[@]}"; do
			if [[ $hosts_line == "0.0.0.0 $domain" || $hosts_line == "0.0.0.0	$domain" ]]; then
				skip=1
				break
			fi
		done
		if [[ $skip -eq 0 ]]; then
			new_content+="${hosts_line}"$'\n'
		fi
	done < "$HOSTS_FILE"

	printf '%s' "$new_content" | sudo bash -c "cat > '$HOSTS_FILE'"
	sudo chmod 644 "$HOSTS_FILE"

	# Re-apply immutable flag
	sudo chattr +i "$HOSTS_FILE" 2>/dev/null || true

	log_info "Removed distraction blocks from /etc/hosts"
}

# Check if Steam is currently running (to track decay)
is_steam_running() {
	local comm_file proc_name
	for comm_file in "$PROC_ROOT"/[0-9]*/comm; do
		[[ -r $comm_file ]] || continue
		read -r proc_name < "$comm_file" || continue
		if [[ $proc_name == "steam" ]]; then
			return 0
		fi
	done
	return 1
}

# Main tracking loop
main_loop() {
	log_info "Starting thesis work tracker daemon"

	# Initialize state
	init_state

	# Load initial state
	load_state

	local total_work_seconds=${TOTAL_WORK_SECONDS:-0}
	local steam_access=${STEAM_ACCESS_GRANTED:-0}
	local session_start=${LAST_WORK_SESSION_START:-0}
	local session_seconds=${CURRENT_SESSION_SECONDS:-0}

	# Apply initial blocking state
	if [[ $steam_access -eq 0 ]]; then
		block_distractions
	fi

	local last_status_log
	printf -v last_status_log '%(%s)T' -1
	local last_decay_check
	printf -v last_decay_check '%(%s)T' -1

	while true; do
		local current_time
		printf -v current_time '%(%s)T' -1

		# Check if thesis work is active
		if is_thesis_work_active; then
			# Track work time
			if [[ $session_start -eq 0 ]]; then
				session_start=$current_time
				log_info "Thesis work session started"
			fi

			# Increment session time
			session_seconds=$((session_seconds + CHECK_INTERVAL))
			total_work_seconds=$((total_work_seconds + CHECK_INTERVAL))

			# Check if we've reached the quota
			if [[ $total_work_seconds -ge $WORK_QUOTA_REQUIRED ]] && [[ $steam_access -eq 0 ]]; then
				log_info "Work quota reached! Granting Steam access."
				steam_access=1
				unblock_distractions
			fi

		else
			# No thesis work active
			if [[ $session_start -ne 0 ]]; then
				log_info "Thesis work session ended. Session duration: $((session_seconds / 60)) minutes"
				session_start=0
				session_seconds=0
			fi

			# Check for Steam usage and apply decay
			if [[ $steam_access -eq 1 ]] && is_steam_running; then
				local time_since_decay=$((current_time - last_decay_check))
				if [[ $time_since_decay -ge 3600 ]]; then # Every hour
					total_work_seconds=$((total_work_seconds - WORK_DECAY_PER_HOUR))
					if [[ $total_work_seconds -lt 0 ]]; then
						total_work_seconds=0
					fi
					last_decay_check=$current_time
					log_info "Steam usage detected. Applied decay. Remaining work time: $((total_work_seconds / 60)) minutes"

					# Revoke access if below quota
					if [[ $total_work_seconds -lt $WORK_QUOTA_REQUIRED ]]; then
						log_info "Work quota depleted. Revoking Steam access."
						steam_access=0
						block_distractions
					fi
				fi
			fi
		fi

		# Save state periodically
		save_state "$total_work_seconds" "$steam_access" "$session_seconds" "$session_start"

		# Log status every 5 minutes
		if [[ $((current_time - last_status_log)) -ge 300 ]]; then
			local work_minutes=$((total_work_seconds / 60))
			local quota_minutes=$((WORK_QUOTA_REQUIRED / 60))
			local remaining_minutes=$((quota_minutes - work_minutes))
			if [[ $remaining_minutes -lt 0 ]]; then
				remaining_minutes=0
			fi

			log_info "Status: Total work time: ${work_minutes}m / ${quota_minutes}m | Steam access: $steam_access | Need: ${remaining_minutes}m more"
			last_status_log=$current_time
		fi

		wait_seconds "$CHECK_INTERVAL"
	done
}

# Handle signals for graceful shutdown
cleanup() {
	log_info "Received shutdown signal, saving state and exiting"
	rm -f "$LOCK_FILE"
	exit 0
}

if [[ ${THESIS_WORK_TRACKER_SKIP_MAIN:-0} -ne 1 ]]; then
	trap cleanup SIGTERM SIGINT

	# Check for lock file to prevent multiple instances
	if [[ -f $LOCK_FILE ]]; then
		log_error "Another instance is already running (lock file exists: $LOCK_FILE)"
		exit 1
	fi

	# Create lock file
	touch "$LOCK_FILE"

	# Run main loop
	main_loop
fi