#!/bin/bash
# BlackRoad Memory System Controller
# Provides continuous memory for Claude Code sessions via PS-SHA∞ journals

set -e

VERSION="1.0.0"

# Configuration
MEMORY_DIR="$HOME/.blackroad/memory"
SESSION_DIR="$MEMORY_DIR/sessions"
JOURNAL_DIR="$MEMORY_DIR/journals"
LEDGER_DIR="$MEMORY_DIR/ledger"
CONTEXT_DIR="$MEMORY_DIR/context"

# Colors
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
PURPLE='\033[0;35m'
NC='\033[0m'

# Helper functions
log_info() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

log_success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1"
}

log_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

# Initialize memory system
init_memory() {
    log_info "Initializing BlackRoad Memory System..."

    mkdir -p "$SESSION_DIR" "$JOURNAL_DIR" "$LEDGER_DIR" "$CONTEXT_DIR"

    # Create genesis entry if not exists
    if [ ! -f "$JOURNAL_DIR/master-journal.jsonl" ]; then
        local timestamp="$(date -u +%Y-%m-%dT%H:%M:%SZ)"
        local genesis_hash="$(echo -n "genesis" | shasum -a 256 | cut -d' ' -f1)"

        echo "{\"timestamp\":\"${timestamp}\",\"action\":\"genesis\",\"entity\":\"memory-system\",\"details\":\"BlackRoad Memory System initialized\",\"sha256\":\"${genesis_hash}\",\"parent_hash\":\"0000000000000000\"}" > "$JOURNAL_DIR/master-journal.jsonl"

        echo "{\"hash\":\"${genesis_hash}\",\"timestamp\":\"${timestamp}\",\"parent\":\"0000000000000000\",\"action_count\":0}" > "$LEDGER_DIR/memory-ledger.jsonl"

        log_success "Memory system initialized with genesis hash: ${genesis_hash:0:8}..."
    else
        log_warning "Memory system already initialized"
    fi

    # Create config
    cat > "$MEMORY_DIR/config.json" <<EOF
{
  "version": "${VERSION}",
  "initialized": "$(date -u +%Y-%m-%dT%H:%M:%SZ)",
  "hash_function": "sha256",
  "encoding": "utf-8"
}
EOF

    log_success "Memory system ready at: $MEMORY_DIR"
}

# Create new session
new_session() {
    local session_name="${1:-default}"
    local session_id="$(date +%Y-%m-%d-%H%M)-${session_name}"
    local session_file="$SESSION_DIR/${session_id}.json"
    local timestamp="$(date -u +%Y-%m-%dT%H:%M:%SZ)"

    # Get last action
    local last_action="N/A"
    if [ -f "$JOURNAL_DIR/master-journal.jsonl" ]; then
        last_action="$(tail -1 "$JOURNAL_DIR/master-journal.jsonl" | jq -r '.action + ": " + .entity' 2>/dev/null || echo "N/A")"
    fi

    # Count git repos
    local git_repos=$(find ~/projects -name .git -type d 2>/dev/null | wc -l | tr -d ' ')

    cat > "$session_file" <<EOF
{
  "session_id": "${session_id}",
  "timestamp": "${timestamp}",
  "type": "session_start",
  "context": {
    "working_directory": "$(pwd)",
    "git_repos": ${git_repos},
    "last_action": "${last_action}",
    "hostname": "$(hostname)",
    "user": "$(whoami)"
  }
}
EOF

    # Create symlink to current session
    ln -sf "$session_file" "$SESSION_DIR/current-session.json"

    # Log session start
    log_action "session_start" "$session_id" "Working directory: $(pwd)"

    log_success "Session created: $session_id"
}

# Log action to journal (with lock-free concurrent write support)
log_action() {
    local action="$1"
    local entity="$2"
    local details="${3:-}"

    if [ ! -f "$JOURNAL_DIR/master-journal.jsonl" ]; then
        log_error "Memory system not initialized. Run: $0 init"
        return 1
    fi

    # Use high-precision timestamp with nanoseconds for uniqueness
    local timestamp="$(date -u +%Y-%m-%dT%H:%M:%S.%3NZ)"

    # Lock-free read: get parent hash (atomic tail operation)
    local parent_hash="$(tail -1 "$JOURNAL_DIR/master-journal.jsonl" | jq -r '.sha256' 2>/dev/null || echo '0000000000000000')"

    # Add process ID and random nonce for uniqueness in concurrent writes
    local nonce="$$-$(od -An -N4 -tu4 < /dev/urandom | tr -d ' ')"
    local hash_input="${timestamp}${action}${entity}${details}${parent_hash}${nonce}"
    local sha256="$(echo -n "$hash_input" | shasum -a 256 | cut -d' ' -f1)"
    local action_count=$(wc -l < "$JOURNAL_DIR/master-journal.jsonl")

    # Get agent SHA-2048 identity if available
    local agent_sha2048="${CLAUDE_SHA2048_SHORT:-}"
    local agent_id="${MY_CLAUDE:-}"

    # Prepare JSON entry (use jq to properly escape multiline strings, -c for compact)
    local temp_file="$JOURNAL_DIR/.temp-${sha256:0:16}.jsonl"
    jq -nc \
        --arg timestamp "$timestamp" \
        --arg action "$action" \
        --arg entity "$entity" \
        --arg details "$details" \
        --arg sha256 "$sha256" \
        --arg parent_hash "$parent_hash" \
        --arg nonce "$nonce" \
        --arg agent_id "$agent_id" \
        --arg agent_sha2048 "$agent_sha2048" \
        '{timestamp: $timestamp, action: $action, entity: $entity, details: $details, sha256: $sha256, parent_hash: $parent_hash, nonce: $nonce, agent_id: $agent_id, agent_sha2048: $agent_sha2048}' \
        > "$temp_file"

    # Lock-free atomic append (>> is atomic at filesystem level for small writes)
    cat "$temp_file" >> "$JOURNAL_DIR/master-journal.jsonl"
    rm -f "$temp_file"

    # Update ledger (also atomic append)
    local ledger_entry="{\"hash\":\"${sha256}\",\"timestamp\":\"${timestamp}\",\"parent\":\"${parent_hash}\",\"action_count\":${action_count},\"nonce\":\"${nonce}\"}"
    echo "$ledger_entry" >> "$LEDGER_DIR/memory-ledger.jsonl"

    echo -e "${PURPLE}[MEMORY]${NC} Logged: ${action} → ${entity} (hash: ${sha256:0:8}...)"
}

# Synthesize context from journals
synthesize_context() {
    log_info "Synthesizing context from memory journals..."

    if [ ! -f "$JOURNAL_DIR/master-journal.jsonl" ]; then
        log_error "No journal found. Run: $0 init"
        return 1
    fi

    # Recent actions
    echo "# Recent Actions" > "$CONTEXT_DIR/recent-actions.md"
    echo "" >> "$CONTEXT_DIR/recent-actions.md"
    tail -50 "$JOURNAL_DIR/master-journal.jsonl" | \
        jq -r '"- [" + .timestamp[0:19] + "] **" + .action + "**: " + .entity + " — " + .details' \
        >> "$CONTEXT_DIR/recent-actions.md"

    # Infrastructure state
    echo "# Infrastructure State" > "$CONTEXT_DIR/infrastructure-state.md"
    echo "" >> "$CONTEXT_DIR/infrastructure-state.md"
    grep -E '"action":"(deployed|configured|allocated)"' "$JOURNAL_DIR/master-journal.jsonl" 2>/dev/null | \
        tail -30 | \
        jq -r '"- **" + .entity + "**: " + .details + " (" + .timestamp[0:10] + ")"' \
        >> "$CONTEXT_DIR/infrastructure-state.md" || echo "No infrastructure changes yet" >> "$CONTEXT_DIR/infrastructure-state.md"

    # Decisions made
    echo "# Decisions Made" > "$CONTEXT_DIR/decisions.md"
    echo "" >> "$CONTEXT_DIR/decisions.md"
    grep -E '"action":"decided"' "$JOURNAL_DIR/master-journal.jsonl" 2>/dev/null | \
        jq -r '"- **" + .entity + "**: " + .details + " (" + .timestamp[0:10] + ")"' \
        >> "$CONTEXT_DIR/decisions.md" || echo "No decisions logged yet" >> "$CONTEXT_DIR/decisions.md"

    log_success "Context synthesized to: $CONTEXT_DIR/"
}

# Show session summary
show_summary() {
    echo -e "${BLUE}╔════════════════════════════════════════╗${NC}"
    echo -e "${BLUE}║  🧠 BlackRoad Memory System Status    ║${NC}"
    echo -e "${BLUE}╚════════════════════════════════════════╝${NC}"
    echo ""

    if [ ! -f "$SESSION_DIR/current-session.json" ]; then
        log_warning "No active session. Run: $0 new [name]"
        return 0
    fi

    # Current session
    local session_id=$(jq -r '.session_id' "$SESSION_DIR/current-session.json" 2>/dev/null || echo "unknown")
    local session_time=$(jq -r '.timestamp' "$SESSION_DIR/current-session.json" 2>/dev/null || echo "unknown")
    local working_dir=$(jq -r '.context.working_directory' "$SESSION_DIR/current-session.json" 2>/dev/null || echo "unknown")

    echo -e "  ${GREEN}Session:${NC} $session_id"
    echo -e "  ${GREEN}Started:${NC} $session_time"
    echo -e "  ${GREEN}Directory:${NC} $working_dir"
    echo ""

    # Journal stats
    if [ -f "$JOURNAL_DIR/master-journal.jsonl" ]; then
        local total_entries=$(wc -l < "$JOURNAL_DIR/master-journal.jsonl")
        local last_hash=$(tail -1 "$JOURNAL_DIR/master-journal.jsonl" | jq -r '.sha256' | cut -c1-16)
        local last_action=$(tail -1 "$JOURNAL_DIR/master-journal.jsonl" | jq -r '.action + ": " + .entity')

        echo -e "  ${GREEN}Total entries:${NC} $total_entries"
        echo -e "  ${GREEN}Last hash:${NC} ${last_hash}..."
        echo -e "  ${GREEN}Last action:${NC} $last_action"
        echo ""

        # Recent changes
        echo -e "${BLUE}Recent changes:${NC}"
        tail -5 "$JOURNAL_DIR/master-journal.jsonl" | \
            jq -r '"  [" + .timestamp[11:19] + "] " + .action + ": " + .entity' | \
            sed "s/^/  /"
    else
        log_warning "No journal entries yet"
    fi

    echo ""
}

# Show system status
show_status() {
    echo -e "${BLUE}╔════════════════════════════════════════╗${NC}"
    echo -e "${BLUE}║  BlackRoad Memory System Status        ║${NC}"
    echo -e "${BLUE}╚════════════════════════════════════════╝${NC}"
    echo ""

    # Check initialization
    if [ -f "$MEMORY_DIR/config.json" ] && [ -f "$JOURNAL_DIR/master-journal.jsonl" ]; then
        echo -e "  ${GREEN}Initialized:${NC} YES"
        local init_time=$(jq -r '.initialized // "unknown"' "$MEMORY_DIR/config.json" 2>/dev/null)
        echo -e "  ${GREEN}Since:${NC} $init_time"
    else
        echo -e "  ${RED}Initialized:${NC} NO"
        echo ""
        log_error "Memory system not initialized. Run: $0 init"
        return 1
    fi

    # Current session
    if [ -f "$SESSION_DIR/current-session.json" ]; then
        local session_name=$(jq -r '.session_id' "$SESSION_DIR/current-session.json" 2>/dev/null || echo "unknown")
        echo -e "  ${GREEN}Current session:${NC} $session_name"
    else
        echo -e "  ${YELLOW}Current session:${NC} none"
    fi

    echo ""

    # Journal stats
    if [ -f "$JOURNAL_DIR/master-journal.jsonl" ]; then
        local total_entries=$(wc -l < "$JOURNAL_DIR/master-journal.jsonl" | tr -d ' ')
        echo -e "  ${GREEN}Total journal entries:${NC} $total_entries"

        local last_line=$(tail -1 "$JOURNAL_DIR/master-journal.jsonl")
        local last_timestamp=$(echo "$last_line" | jq -r '.timestamp' 2>/dev/null || echo "unknown")
        local last_action=$(echo "$last_line" | jq -r '.action' 2>/dev/null || echo "unknown")
        local last_entity=$(echo "$last_line" | jq -r '.entity' 2>/dev/null || echo "unknown")
        echo -e "  ${GREEN}Last entry:${NC} $last_timestamp"
        echo -e "  ${GREEN}Last action:${NC} $last_action: $last_entity"
    else
        echo -e "  ${YELLOW}Total journal entries:${NC} 0"
    fi

    echo ""

    # Ledger size
    if [ -f "$LEDGER_DIR/memory-ledger.jsonl" ]; then
        local ledger_size=$(ls -lh "$LEDGER_DIR/memory-ledger.jsonl" | awk '{print $5}')
        echo -e "  ${GREEN}Ledger size:${NC} $ledger_size"
    else
        echo -e "  ${YELLOW}Ledger size:${NC} not found"
    fi

    echo ""

    # Chain integrity (quick check: last 5 entries)
    echo -e "  ${BLUE}Chain integrity (last 5 entries):${NC}"
    if [ -f "$JOURNAL_DIR/master-journal.jsonl" ]; then
        local total=$(wc -l < "$JOURNAL_DIR/master-journal.jsonl" | tr -d ' ')
        local start=$((total > 5 ? total - 4 : 1))
        local chain_valid=true
        local checked=0

        tail -5 "$JOURNAL_DIR/master-journal.jsonl" | while IFS= read -r line; do
            local parent_hash=$(echo "$line" | jq -r 'if .parent_hash then .parent_hash elif .prev_hash then .prev_hash else "0000000000000000" end' 2>/dev/null)
            local stored_hash=$(echo "$line" | jq -r 'if .sha256 then .sha256 elif .hash then .hash else "unknown" end' 2>/dev/null)
            local action=$(echo "$line" | jq -r '.action' 2>/dev/null)

            if [ "$parent_hash" = "0000000000000000" ] || [ "$parent_hash" = "null" ]; then
                echo -e "    ${GREEN}OK${NC} ${stored_hash:0:12}... ($action) [genesis/root]"
            elif grep -q "\"$parent_hash\"" "$JOURNAL_DIR/master-journal.jsonl" 2>/dev/null; then
                echo -e "    ${GREEN}OK${NC} ${stored_hash:0:12}... ($action)"
            else
                echo -e "    ${RED}BROKEN${NC} ${stored_hash:0:12}... ($action) - parent not found"
            fi
        done
    fi

    echo ""
}

# Verify integrity
verify_integrity() {
    log_info "Verifying memory integrity..."

    if [ ! -f "$JOURNAL_DIR/master-journal.jsonl" ]; then
        log_error "No journal found"
        return 1
    fi

    local entries=$(wc -l < "$JOURNAL_DIR/master-journal.jsonl")
    local valid=0
    local invalid=0
    local line_num=0

    while IFS= read -r line; do
        line_num=$((line_num + 1))

        # Support both old format (prev_hash/hash) and new format (parent_hash/sha256)
        local parent_hash=$(echo "$line" | jq -r 'if .parent_hash then .parent_hash elif .prev_hash then .prev_hash else "0000000000000000" end')
        local stored_hash=$(echo "$line" | jq -r 'if .sha256 then .sha256 elif .hash then .hash else "unknown" end')

        # Skip genesis entries and session roots (entries with no valid parent)
        if [ "$parent_hash" = "0000000000000000" ] || [ "$parent_hash" = "null" ] || [ -z "$parent_hash" ]; then
            valid=$((valid + 1))
            continue
        fi

        # Verify parent exists in previous entries (check both hash field names)
        if [ $line_num -le 1 ]; then
            invalid=$((invalid + 1))
            log_warning "Broken chain at entry $line_num (hash: ${stored_hash:0:8}...)"
        elif head -$((line_num - 1)) "$JOURNAL_DIR/master-journal.jsonl" | grep -qE "\"(sha256|hash)\":\"$parent_hash\""; then
            valid=$((valid + 1))
        else
            invalid=$((invalid + 1))
            log_warning "Broken chain at entry $line_num (hash: ${stored_hash:0:8}...)"
        fi
    done < "$JOURNAL_DIR/master-journal.jsonl"

    echo ""
    if [ $invalid -eq 0 ]; then
        log_success "Memory integrity verified ✅ ($valid entries, 0 broken)"
    else
        log_error "Found $invalid broken entries ❌ ($valid valid entries)"
        return 1
    fi
}

# Export context for Claude Code
export_context() {
    log_info "Exporting context for Claude Code..."

    synthesize_context

    cat > "$CONTEXT_DIR/session-restore-context.md" <<EOF
# 🧠 BlackRoad Session Context

**Generated:** $(date -u +%Y-%m-%dT%H:%M:%SZ)

---

## Session Summary

$(show_summary 2>&1)

---

## Recent Infrastructure Changes

$(grep -E '"action":"(deployed|configured|allocated)"' "$JOURNAL_DIR/master-journal.jsonl" 2>/dev/null | tail -20 | jq -r '"- [" + .timestamp[0:19] + "] **" + .action + "**: " + .entity + " — " + .details' || echo "No infrastructure changes yet")

---

## Recent Decisions

$(grep -E '"action":"decided"' "$JOURNAL_DIR/master-journal.jsonl" 2>/dev/null | tail -10 | jq -r '"- [" + .timestamp[0:19] + "] **" + .entity + "**: " + .details' || echo "No decisions logged yet")

---

## Active Deployments

$(ssh pi@aria64 "docker ps --format '{{.Names}} → {{.Ports}}'" 2>/dev/null || echo "Cannot reach aria64")

---

## Current Working State

**Directory:** $(pwd)

**Git Status:**
\`\`\`
$(git status -s 2>/dev/null || echo "Not in git repository")
\`\`\`

---

**Memory integrity:** $(verify_integrity &>/dev/null && echo "✅ Verified" || echo "⚠️  Check needed")

EOF

    log_success "Context exported to: $CONTEXT_DIR/session-restore-context.md"
    echo ""
    echo "View with: cat $CONTEXT_DIR/session-restore-context.md"
}

# Show help
show_help() {
    cat <<EOF
BlackRoad Memory System v${VERSION}

USAGE:
    memory-system.sh <command> [options]

COMMANDS:
    init                          Initialize memory system
    new [session-name]            Create new session
    log <action> <entity> [details]  Log action to journal
    status                        Show system status and chain integrity
    summary                       Show current session summary
    synthesize                    Synthesize context from journals
    verify                        Verify memory integrity
    export                        Export context for Claude Code
    help                          Show this help

EXAMPLES:
    # Initialize
    memory-system.sh init

    # Start new session
    memory-system.sh new infrastructure-audit

    # Log actions
    memory-system.sh log deployed "api.blackroad.io" "Port 8080, FastAPI"
    memory-system.sh log created "new-script.sh" "Automated deployment"
    memory-system.sh log decided "architecture" "Using Headscale for mesh"

    # View state
    memory-system.sh summary
    memory-system.sh export

    # Verify
    memory-system.sh verify

MEMORY LOCATIONS:
    Sessions:  $SESSION_DIR/
    Journals:  $JOURNAL_DIR/
    Ledger:    $LEDGER_DIR/
    Context:   $CONTEXT_DIR/

EOF
}

# Main command handler
case "${1:-help}" in
    init)
        init_memory
        ;;
    new)
        new_session "${2:-default}"
        ;;
    log)
        if [ -z "$2" ] || [ -z "$3" ]; then
            log_error "Usage: $0 log <action> <entity> [details]"
            exit 1
        fi
        log_action "$2" "$3" "${4:-}"
        ;;
    synthesize)
        synthesize_context
        ;;
    summary)
        show_summary
        ;;
    status)
        show_status
        ;;
    verify)
        verify_integrity
        ;;
    export)
        export_context
        ;;
    help|--help|-h)
        show_help
        ;;
    *)
        log_error "Unknown command: $1"
        echo ""
        show_help
        exit 1
        ;;
esac