#!/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=5 # Check every 5 seconds

# 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
	timestamp=$(date '+%Y-%m-%d %H:%M:%S')
	echo "[${timestamp}] [${level}] ${message}" | tee -a "$LOG_FILE"
}

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
}

# 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
		cat <<EOF | sudo tee "$STATE_FILE" >/dev/null
# Thesis Work Tracker State File
# DO NOT EDIT MANUALLY - Managed by thesis_work_tracker daemon
# Last updated: $(date)

TOTAL_WORK_SECONDS=0
LAST_UPDATE_TIMESTAMP=$(date +%s)
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
	# Only extract the numeric values we need
	TOTAL_WORK_SECONDS=$(grep "^TOTAL_WORK_SECONDS=" "$STATE_FILE" 2>/dev/null | cut -d= -f2 || echo "0")
	STEAM_ACCESS_GRANTED=$(grep "^STEAM_ACCESS_GRANTED=" "$STATE_FILE" 2>/dev/null | cut -d= -f2 || echo "0")
	CURRENT_SESSION_SECONDS=$(grep "^CURRENT_SESSION_SECONDS=" "$STATE_FILE" 2>/dev/null | cut -d= -f2 || echo "0")
	LAST_WORK_SESSION_START=$(grep "^LAST_WORK_SESSION_START=" "$STATE_FILE" 2>/dev/null | cut -d= -f2 || echo "0")
	# shellcheck disable=SC2034  # Written back to state file in save_state
	LAST_UPDATE_TIMESTAMP=$(grep "^LAST_UPDATE_TIMESTAMP=" "$STATE_FILE" 2>/dev/null | cut -d= -f2 || echo "0")

	# 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

	# Write new state
	cat <<EOF | sudo tee "$STATE_FILE" >/dev/null
# Thesis Work Tracker State File
# DO NOT EDIT MANUALLY - Managed by thesis_work_tracker daemon
# Last updated: $(date)

TOTAL_WORK_SECONDS=$total_work
LAST_UPDATE_TIMESTAMP=$(date +%s)
STEAM_ACCESS_GRANTED=$steam_access
LAST_WORK_SESSION_START=$session_start
CURRENT_SESSION_SECONDS=$current_session
EOF

	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
}

# Check if a process is running
is_process_running() {
	local process_name="$1"
	pgrep -x "$process_name" >/dev/null 2>&1
}

# 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
	window_name=$(xdotool getwindowname "$active_window_id" 2>/dev/null || echo "")

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

	local process_name=""
	if [[ -n $window_pid ]]; then
		process_name=$(ps -p "$window_pid" -o comm= 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 /etc/hosts 2>/dev/null || true

	# Add blocking entries if not already present
	local hosts_modified=0

	for domain in "${STEAM_DOMAINS[@]}" "${DISTRACTION_DOMAINS[@]}"; do
		if ! grep -q "^0.0.0.0[[:space:]]*$domain" /etc/hosts 2>/dev/null; then
			echo "0.0.0.0 $domain" | sudo tee -a /etc/hosts >/dev/null
			hosts_modified=1
		fi
	done

	# Re-apply immutable flag
	sudo chattr +i /etc/hosts 2>/dev/null || true

	if [[ $hosts_modified -eq 1 ]]; then
		log_info "Added distraction blocks to /etc/hosts"
	fi
}

# 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 /etc/hosts 2>/dev/null || true

	# Remove blocking entries using mktemp for security
	local temp_hosts
	temp_hosts=$(mktemp) || {
		log_error "Failed to create temporary file"
		return 1
	}

	sudo cp /etc/hosts "$temp_hosts"

	for domain in "${STEAM_DOMAINS[@]}" "${DISTRACTION_DOMAINS[@]}"; do
		sudo sed -i "/^0.0.0.0[[:space:]]*$domain/d" "$temp_hosts"
	done

	sudo mv "$temp_hosts" /etc/hosts
	sudo chmod 644 /etc/hosts

	# Re-apply immutable flag
	sudo chattr +i /etc/hosts 2>/dev/null || true

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

# Check if Steam is currently running (to track decay)
is_steam_running() {
	pgrep -x "steam" >/dev/null 2>&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
	last_status_log=$(date +%s)
	local last_decay_check
	last_decay_check=$(date +%s)

	while true; do
		local current_time
		current_time=$(date +%s)

		# 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

		sleep "$CHECK_INTERVAL"
	done
}

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

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