#!/bin/bash
set -euo pipefail

# Source required modules for session management functions
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
source "$SCRIPT_DIR/kugetsu-config.sh"
source "$SCRIPT_DIR/kugetsu-index.sh"
source "$SCRIPT_DIR/kugetsu-worktree.sh"
source "$SCRIPT_DIR/kugetsu-log.sh"

count_active_dev_sessions() {
    local count=0
    if [ -d "$SESSIONS_DIR" ]; then
        for session_file in "$SESSIONS_DIR"/*.json; do
            if [ -f "$session_file" ]; then
                local filename=$(basename "$session_file")
                if [ "$filename" != "base.json" ] && [ "$filename" != "pm-agent.json" ]; then
                    count=$((count + 1))
                fi
            fi
        done
    fi
    echo "$count"
}

cmd_init() {
    local force=false

    while [ $# -gt 0 ]; do
        case "$1" in
            --force)
                force=true
                ;;
            *)
                ;;
        esac
        shift
    done

    ensure_dirs

    if [ ! -f "$KUGETSU_DIR/config" ]; then
        cat > "$KUGETSU_DIR/config" << 'EOF'
# User configuration overrides
# Values set here take precedence over defaults
# Changes take effect immediately (no re-init needed)

# Max concurrent dev agents (default: 3)
# MAX_CONCURRENT_AGENTS=5

# Git server configurations
# Format: GIT_SERVERS["hostname"]="https://hostname"
declare -A GIT_SERVERS
GIT_SERVERS["github.com"]="https://github.com"
GIT_SERVERS["git.fbrns.co"]="https://git.fbrns.co"
DEFAULT_GIT_SERVER="github.com"
EOF
        echo "Created config file: $KUGETSU_DIR/config"
    fi

    mkdir -p "$ENV_DIR"
    if [ ! -f "$ENV_DIR/default.env" ]; then
        cat > "$ENV_DIR/default.env" << 'EOF'
# Environment variables for agents
# Copy this file to <agent-type>.env (e.g., pm-agent.env, dev.env)
# and set your tokens and configuration

# Required: Gitea token for API access
# GITEA_TOKEN=your_gitea_token_here

# Optional: GitHub token (if using GitHub)
# GITHUB_TOKEN=your_github_token_here

# Optional: GitLab token (if using GitLab)
# GITLAB_TOKEN=your_gitlab_token_here
EOF
        echo "Created env template: $ENV_DIR/default.env"
    fi

    local existing_base=$(get_base_session_id)
    local existing_pm=$(get_pm_agent_session_id)
    
    if [ -n "$existing_base" ] && [ "$existing_base" != "null" ]; then
        if [ "$force" = true ]; then
            echo "Warning: Reinitializing sessions (force mode)" >&2
        else
            echo "Error: Base session already exists: $existing_base" >&2
            echo "Use --force to reinitialize" >&2
            exit 1
        fi
    fi

    if ! test -t 0; then
        echo "Error: init requires a terminal (TTY)" >&2
        echo "Please run this command in an interactive shell" >&2
        exit 1
    fi

    echo "Starting TUI to create base session..."
    echo "Press Ctrl+C to cancel or wait for session to be created"
    sleep 2

    local before_sessions=$(opencode session list 2>/dev/null | grep -E '^ses_' | awk '{print $1}' || true)

    opencode

    local after_sessions=$(opencode session list 2>/dev/null | grep -E '^ses_' | awk '{print $1}' || true)
    local session_ids=""
    while IFS= read -r line; do
        local sid=$(echo "$line" | awk '{print $1}')
        if [ -n "$sid" ] && ! echo "$before_sessions" | grep -q "^${sid}$"; then
            session_ids="$sid"
            break
        fi
    done <<< "$after_sessions"

    if [ -z "$session_ids" ]; then
        echo "Error: Could not find newly created session" >&2
        exit 1
    fi

    echo "$session_ids" > "$SESSIONS_DIR/base.json"
    set_base_in_index "$session_ids"

    echo "Base session created: $session_ids"
    echo "Starting PM agent..."

    before_sessions="$after_sessions"

    opencode

    after_sessions=$(opencode session list 2>/dev/null | grep -E '^ses_' | awk '{print $1}' || true)
    local pm_session_ids=""
    while IFS= read -r line; do
        local sid=$(echo "$line" | awk '{print $1}')
        if [ -n "$sid" ] && ! echo "$before_sessions" | grep -q "^${sid}$"; then
            pm_session_ids="$sid"
            break
        fi
    done <<< "$after_sessions"

    if [ -z "$pm_session_ids" ]; then
        echo "Warning: Could not find separate PM agent session" >&2
        pm_session_ids="$session_ids"
    fi

    echo "$pm_session_ids" > "$SESSIONS_DIR/pm-agent.json"
    set_pm_agent_in_index "$pm_session_ids"

    load_agent_env "pm-agent"

    local pm_system_prompt=""
    if [ -f "$KUGETSU_DIR/pm-agent.md" ]; then
        pm_system_prompt=$(cat "$KUGETSU_DIR/pm-agent.md")
        echo "Injecting PM agent system prompt from $KUGETSU_DIR/pm-agent.md"
    fi

    echo "PM agent session created: $pm_session_ids"
    echo ""
    echo "kugetsu initialized successfully!"
    echo "  Base session: $session_ids"
    echo "  PM agent: $pm_session_ids"
}

extract_issue_ref_from_message() {
    local message="$1"
    
    if [ -z "$message" ]; then
        echo ""
        return
    fi
    
    if [[ "$message" =~ ^([a-zA-Z0-9.-]+/[a-zA-Z0-9._-]+/[a-zA-Z0-9._-]+#[0-9]+) ]]; then
        echo "${BASH_REMATCH[1]}"
        return
    fi
    
    if [[ "$message" =~ (https?://)?([a-zA-Z0-9.-]+)/([a-zA-Z0-9._-]+)/([a-zA-Z0-9._-]+)/(issues|pull)/([0-9]+) ]]; then
        local instance="${BASH_REMATCH[2]}"
        local owner="${BASH_REMATCH[3]}"
        local repo="${BASH_REMATCH[4]}"
        local num="${BASH_REMATCH[6]}"
        echo "${instance}/${owner}/${repo}#${num}"
        return
    fi
    
    echo ""
}

cmd_delegate() {
    local message="${1:-}"
    
    if [ -z "$message" ]; then
        echo "Error: message is required" >&2
        echo "Usage: kugetsu delegate <message>" >&2
        exit 1
    fi
    
    local issue_ref=$(extract_issue_ref_from_message "$message")
    
    if [ -n "$issue_ref" ] && [[ "$issue_ref" =~ \#[0-9]+$ ]]; then
        # Enqueue for daemon to process via cmd_start/cmd_continue
        enqueue_task "$issue_ref" "$message"
        return
    fi
    
    # No issue ref detected — fork a new session from base session
    local base_session=$(get_base_session_id)
    if [ -z "$base_session" ] || [ "$base_session" = "null" ]; then
        echo "Error: Base session not found. Run 'kugetsu init' first." >&2
        exit 1
    fi
    
    mkdir -p "$LOGS_DIR"
    local log_file="$LOGS_DIR/delegate-$(date +%s).log"
    load_agent_env "pm-agent"
    
    local new_session=$(create_session "$base_session")
    if [ -z "$new_session" ]; then
        echo "Error: Failed to create session" >&2
        exit 1
    fi
    
    nohup sh -c "GITEA_TOKEN='${GITEA_TOKEN:-}' opencode run '$message' --session '$new_session'" >> "$log_file" 2>&1 &
    echo "Delegated to new session (logged to $(basename "$log_file"))"
}

create_session() {
    local base_session="${1:-$base_session_id}"
    
    if [ -z "$base_session" ] || [ "$base_session" = "null" ]; then
        echo "Error: base session not found. Run 'kugetsu init' first." >&2
        return 1
    fi
    
    local before_json=$(opencode session list --format=json 2>/dev/null)
    local before_ids=$(echo "$before_json" | python3 -c "import sys,json; sessions=json.load(sys.stdin); print(' '.join(s['id'] for s in sessions))" 2>/dev/null || echo "")
    
    opencode run --fork --session "$base_session" "new session" 2>/dev/null
    
    local after_json=$(opencode session list --format=json 2>/dev/null)
    local after_ids=$(echo "$after_json" | python3 -c "import sys,json; sessions=json.load(sys.stdin); print(' '.join(s['id'] for s in sessions))" 2>/dev/null || echo "")
    
    local new_session_id=""
    for sess in $after_ids; do
        if [[ ! " $before_ids " =~ " $sess " ]] && [[ "$sess" != "$base_session" ]]; then
            new_session_id="$sess"
            break
        fi
    done
    
    echo "$new_session_id"
}

build_dev_agent_message() {
    local issue_ref="$1"
    local user_message="${2:-}"
    
    local instance=$(echo "$issue_ref" | cut -d'/' -f1 | cut -d'#' -f1)
    local owner=$(echo "$issue_ref" | cut -d'/' -f2)
    local repo=$(echo "$issue_ref" | cut -d'/' -f3 | cut -d'#' -f1)
    local number=$(echo "$issue_ref" | grep -oE '#[0-9]+$' | tr -d '#')
    local worktree_path=$(issue_ref_to_worktree_path "$issue_ref")
    
    local base_message="You are assigned to work on $issue_ref.

Workflow:
1. Read the issue at $instance/$owner/$repo/issues/$number AND all comments on that issue
2. Check if a PR already exists for this issue
   - If PR exists and is open, review it and learn from it
   - If PR makes sense to continue, work on it instead
   - If PR is not worth continuing, create a new branch/PR but explain in PR description why you're creating a new one instead of continuing the existing PR
3. Read README.md (if exists) to understand the general concept of this repository
4. Read CONTRIBUTING.md (if exists) to understand how to contribute
   - If CONTRIBUTING.md doesn't exist, follow steps 5-9 as your guideline
5. Explore the repository to understand the codebase
6. If anything is unclear, post a comment on the issue asking for clarification before implementing
7. Implement the solution
8. Create a branch named fix/issue-$number and implement the fix
9. Create a PR when the implementation is complete

Work directory: $worktree_path"
    
    if [ -n "$user_message" ]; then
        echo "$base_message

Additional instructions from delegator:
$user_message"
    else
        echo "$base_message"
    fi
}

cmd_start() {
    local issue_ref="${1:-}"
    local message="${2:-}"
    
    if [ -z "$issue_ref" ]; then
        echo "Error: issue ref is required" >&2
        echo "Usage: kugetsu start <issue-ref> [message]" >&2
        exit 1
    fi
    
    validate_issue_ref "$issue_ref"
    
    local base_session_id=$(get_base_session_id)
    if [ -z "$base_session_id" ] || [ "$base_session_id" = "null" ]; then
        echo "Error: Base session not found. Run 'kugetsu init' first." >&2
        exit 1
    fi
    
    if worktree_exists "$issue_ref"; then
        echo "Issue '$issue_ref' already has a worktree. Use 'kugetsu continue' instead."
        exit 1
    fi
    
    local active_count=$(count_active_dev_sessions)
    if [ "$active_count" -ge "${MAX_CONCURRENT_AGENTS:-3}" ]; then
        echo "Error: Max concurrent agents (${MAX_CONCURRENT_AGENTS:-3}) reached. Use 'kugetsu continue' or wait for an agent to finish." >&2
        exit 1
    fi
    
    local session_file=$(issue_ref_to_filename "$issue_ref")
    local session_path="$SESSIONS_DIR/$session_file"
    
    if [ -f "$session_path" ]; then
        echo "Session file already exists: $session_file"
        echo "Use 'kugetsu continue $issue_ref' to continue work."
        exit 1
    fi
    
    create_worktree "$issue_ref" "$WORKTREES_DIR"
    
    local new_session_id=$(create_session "$base_session_id")
    
    if [ -z "$new_session_id" ]; then
        echo "Error: Could not create session" >&2
        remove_worktree_for_issue "$issue_ref"
        exit 1
    fi
    
    local worktree_path=$(issue_ref_to_worktree_path "$issue_ref")
    
    printf '{"type": "forked", "issue_ref": "%s", "opencode_session_id": "%s", "worktree_path": "%s", "created_at": "%s", "state": "idle"}\n' \
        "$issue_ref" "$new_session_id" "$worktree_path" "$(date -Iseconds)" > "$SESSIONS_DIR/$session_file"
    
    add_issue_to_index "$issue_ref" "$session_file"
    
    local dev_message=$(build_dev_agent_message "$issue_ref" "$message")
    
    load_agent_env "dev"
    
    cd "$worktree_path"
    nohup sh -c "GITEA_TOKEN='${GITEA_TOKEN:-}' opencode run '$dev_message' --session '$new_session_id'" >> "$LOGS_DIR/dev-$new_session_id.log" 2>&1 &
    
    echo "Session started for '$issue_ref': $new_session_id"
    echo "Worktree: $worktree_path"
}

cmd_continue() {
    local session_name=""
    local message=""
    local args=("$@")

    args=$(set_debug_mode "${args[@]}")

    for arg in $args; do
        if [ -z "$session_name" ]; then
            session_name="$arg"
        else
            message="$arg"
        fi
    done
    
    if [ -z "$session_name" ]; then
        echo "Error: issue ref is required" >&2
        echo "Usage: kugetsu continue <issue-ref> [message]" >&2
        exit 1
    fi
    
    validate_issue_ref "$session_name"
    
    local session_file=$(get_session_for_issue "$session_name")
    if [ -z "$session_file" ] || [ "$session_file" = "null" ]; then
        echo "Error: No session found for '$session_name'" >&2
        echo "Use 'kugetsu start $session_name' to create a new session." >&2
        exit 1
    fi
    
    local session_path="$SESSIONS_DIR/$session_file"
    
    if [ ! -f "$session_path" ]; then
        echo "Error: Session file not found: $session_path" >&2
        exit 1
    fi
    
    load_agent_env "dev"
    
    local opencode_session_id=$(python3 -c "import json; print(json.load(open('$session_path')).get('opencode_session_id', ''))" 2>/dev/null || echo "")
    local worktree_path=$(python3 -c "import json; print(json.load(open('$session_path')).get('worktree_path', ''))" 2>/dev/null || echo "")
    local issue_ref=$(python3 -c "import json; print(json.load(open('$session_path')).get('issue_ref', ''))" 2>/dev/null || echo "")
    
    if [ -z "$message" ]; then
        message=$(build_dev_agent_message "$issue_ref" "")
    fi
    
    if [ -n "$worktree_path" ] && [ -d "$worktree_path" ]; then
        cd "$worktree_path"
        nohup sh -c "GITEA_TOKEN='${GITEA_TOKEN:-}' opencode run '$message' --session '$opencode_session_id'" >> "$LOGS_DIR/dev-$opencode_session_id.log" 2>&1 &
    else
        nohup sh -c "GITEA_TOKEN='${GITEA_TOKEN:-}' opencode run '$message' --session '$opencode_session_id'" >> "$LOGS_DIR/dev-$opencode_session_id.log" 2>&1 &
    fi
}

cmd_list() {
    echo "=== kugetsu sessions ==="
    echo ""
    
    local base_id=$(get_base_session_id)
    if [ -n "$base_id" ] && [ "$base_id" != "null" ]; then
        echo "Base session: $base_id"
    else
        echo "Base session: not initialized"
    fi
    
    local pm_id=$(get_pm_agent_session_id)
    if [ -n "$pm_id" ] && [ "$pm_id" != "null" ]; then
        echo "PM agent: $pm_id"
    else
        echo "PM agent: not initialized"
    fi
    
    echo ""
    echo "Issue sessions:"
    
    if [ -d "$SESSIONS_DIR" ]; then
        for session_file in "$SESSIONS_DIR"/*.json; do
            if [ -f "$session_file" ]; then
                local filename=$(basename "$session_file")
                if [ "$filename" != "base.json" ] && [ "$filename" != "pm-agent.json" ]; then
                    local issue_ref=$(filename_to_issue_ref "$filename")
                    local opencode_sid=$(python3 -c "import json; print(json.load(open('$session_file')).get('opencode_session_id', 'unknown'))" 2>/dev/null || echo "unknown")
                    local state=$(python3 -c "import json; print(json.load(open('$session_file')).get('state', 'unknown'))" 2>/dev/null || echo "unknown")
                    local worktree_path=$(python3 -c "import json; print(json.load(open('$session_file')).get('worktree_path', ''))" 2>/dev/null || echo "")
                    local worktree_status=""
                    if [ -n "$worktree_path" ]; then
                        if [ -d "$worktree_path" ]; then
                            worktree_status="(worktree exists)"
                        else
                            worktree_status="(worktree MISSING)"
                        fi
                    fi
                    echo "  $filename"
                    echo "    Issue: $issue_ref"
                    echo "    Session: $opencode_sid"
                    echo "    State: $state"
                    echo "    $worktree_status"
                fi
            fi
        done
    fi
    
    if [ -d "$WORKTREES_DIR" ]; then
        echo ""
        echo "Worktrees without sessions:"
        for worktree in $(find "$WORKTREES_DIR" -mindepth 1 -maxdepth 1 -type d 2>/dev/null); do
            local worktree_name=$(basename "$worktree")
            local has_session=false
            for session_file in "$SESSIONS_DIR"/*.json; do
                if [ -f "$session_file" ]; then
                    local wt_path=$(python3 -c "import json; print(json.load(open('$session_file')).get('worktree_path', ''))" 2>/dev/null || echo "")
                    if [ "$wt_path" = "$worktree" ]; then
                        has_session=true
                        break
                    fi
                fi
            done
            if [ "$has_session" = false ]; then
                echo "  $worktree_name (no session)"
            fi
        done
    fi
}

cmd_prune() {
    local force=false
    
    if [ "$1" = "--force" ]; then
        force=true
    fi
    
    echo "=== kugetsu prune ==="
    echo ""
    
    local orphaned=()
    
    if [ -d "$SESSIONS_DIR" ]; then
        for session_file in "$SESSIONS_DIR"/*.json; do
            [ -f "$session_file" ] || continue
            local filename=$(basename "$session_file")
            if [ "$filename" = "base.json" ] || [ "$filename" = "pm-agent.json" ]; then
                continue
            fi
            
            local opencode_sid=$(python3 -c "import json; print(json.load(open('$session_file')).get('opencode_session_id', ''))" 2>/dev/null || echo "")
            
            if [ -n "$opencode_sid" ]; then
                local exists=$(opencode session list 2>/dev/null | grep -c "^$opencode_sid" || echo "0")
                if [ "$exists" -eq 0 ]; then
                    orphaned+=("$session_file")
                fi
            else
                orphaned+=("$session_file")
            fi
        done
    fi
    
    if [ ${#orphaned[@]} -eq 0 ]; then
        echo "No orphaned sessions found."
        return
    fi
    
    echo "Found ${#orphaned[@]} orphaned session(s):"
    for session in "${orphaned[@]}"; do
        echo "  $session"
    done
    echo ""
    
    if [ "$force" = false ]; then
        read -p "Remove these sessions? [y/N] " -r answer
        if [[ ! "$answer" =~ ^[Yy]$ ]]; then
            echo "Aborted."
            return
        fi
    fi
    
    for session in "${orphaned[@]}"; do
        local issue_ref=$(python3 -c "import json; print(json.load(open('$session')).get('issue_ref', ''))" 2>/dev/null || echo "")
        local worktree_path=$(python3 -c "import json; print(json.load(open('$session')).get('worktree_path', ''))" 2>/dev/null || echo "")
        
        if [ -n "$worktree_path" ] && [ -d "$worktree_path" ]; then
            echo "Removing worktree: $worktree_path"
            git worktree remove "$worktree_path" 2>/dev/null || rm -rf "$worktree_path"
        fi
        
        rm -f "$session"
        
        if [ -n "$issue_ref" ]; then
            remove_issue_from_index "$issue_ref"
        fi
        
        echo "Removed: $session"
    done
    
    echo ""
    echo "Pruned ${#orphaned[@]} orphaned session(s)."
}

cmd_destroy() {
    local target="${1:-}"
    local force=false
    
    if [ "${2:-}" = "-y" ]; then
        force=true
    fi
    
    if [ -z "$target" ]; then
        echo "Error: target is required" >&2
        echo "Usage: kugetsu destroy <issue-ref> [-y]" >&2
        echo "       kugetsu destroy --pm-agent [-y]" >&2
        echo "       kugetsu destroy --base [-y]" >&2
        exit 1
    fi
    
    if [ "$target" = "--pm-agent" ]; then
        if [ "$force" = false ]; then
            echo "Warning: Destroying PM agent session is not recommended." >&2
            read -p "Continue? [y/N] " -r answer
            if [[ ! "$answer" =~ ^[Yy]$ ]]; then
                echo "Aborted."
                return
            fi
        fi
        
        local pm_session=$(get_pm_agent_session_id)
        if [ -n "$pm_session" ] && [ "$pm_session" != "null" ]; then
            echo "Stopping PM agent session: $pm_session"
            opencode session stop "$pm_session" 2>/dev/null || true
        fi
        
        rm -f "$SESSIONS_DIR/pm-agent.json"
        set_pm_agent_in_index "null"
        echo "PM agent session destroyed"
    elif [ "$target" = "--base" ]; then
        if [ "$force" = false ]; then
            echo "Warning: Destroying base session will remove ALL sessions." >&2
            read -p "Continue? [y/N] " -r answer
            if [[ ! "$answer" =~ ^[Yy]$ ]]; then
                echo "Aborted."
                return
            fi
        fi
        
        for session_file in "$SESSIONS_DIR"/*.json; do
            [ -f "$session_file" ] || continue
            rm -f "$session_file"
        done
        
        for worktree in "$WORKTREES_DIR"/.kugetsu-worktrees/*; do
            if [ -d "$worktree" ]; then
                git worktree remove "$worktree" 2>/dev/null || rm -rf "$worktree"
            fi
        done
        
        write_index "null" "null" "{}"
        echo "Base session and all worktrees destroyed"
    else
        validate_issue_ref "$target"
        
        local session_file=$(get_session_for_issue "$target")
        if [ -z "$session_file" ] || [ "$session_file" = "null" ]; then
            echo "Error: No session found for '$target'" >&2
            exit 1
        fi
        
        local session_path="$SESSIONS_DIR/$session_file"
        
        if [ "$force" = true ]; then
            remove_worktree_for_issue "$target"
            rm -f "$session_path"
            remove_issue_from_index "$target"
            echo "Session for '$target' destroyed"
        else
            echo "Warning: This will delete session and worktree for '$target'" >&2
            read -p "Continue? [y/N] " -r answer
            if [[ ! "$answer" =~ ^[Yy]$ ]]; then
                echo "Aborted."
                return
            fi
            
            remove_worktree_for_issue "$target"
            rm -f "$session_path"
            remove_issue_from_index "$target"
            echo "Session for '$target' destroyed"
        fi
    fi
}

cmd_status() {
    echo "=== kugetsu status ==="
    
    local base_id=$(get_base_session_id)
    local pm_id=$(get_pm_agent_session_id)
    
    if [ -z "$base_id" ] || [ "$base_id" = "null" ]; then
        echo "Status: Not initialized"
        echo "Run 'kugetsu init' to initialize."
        return
    fi
    
    echo "Base session: $base_id"
    
    if [ -n "$pm_id" ] && [ "$pm_id" != "null" ]; then
        echo "PM agent: $pm_id"
    else
        echo "PM agent: not running"
    fi
    
    local active_count=$(count_active_dev_sessions)
    echo "Active issue sessions: $active_count / ${MAX_CONCURRENT_AGENTS:-3}"
    
    echo ""
    echo "OpenCode sessions:"
    opencode session list 2>/dev/null || echo "  (unable to list sessions)"
}