#!/bin/bash
#
# SessionEnd Hook: Save final state and log session end.
#

# Get Grid directory
PROJECT_DIR="${CLAUDE_PROJECT_DIR:-$(pwd)}"
GRID_DIR="$PROJECT_DIR/.grid"
SCRATCHPAD="$GRID_DIR/SCRATCHPAD.md"
STATE_FILE="$GRID_DIR/STATE.md"
BUDGET_FILE="$GRID_DIR/budget.json"

# Check if Grid is initialized
if [[ ! -d "$GRID_DIR" ]]; then
    exit 0
fi

TIMESTAMP=$(date +"%Y-%m-%dT%H:%M:%S")

# Log session end to scratchpad
if [[ -f "$SCRATCHPAD" ]]; then
    echo "" >> "$SCRATCHPAD"
    echo "--- SESSION END: $TIMESTAMP ---" >> "$SCRATCHPAD"

    # Count operations this session
    SESSION_START=$(grep -n "SESSION START" "$SCRATCHPAD" | tail -1 | cut -d: -f1)
    if [[ -n "$SESSION_START" ]]; then
        TOTAL_LINES=$(wc -l < "$SCRATCHPAD" | tr -d ' ')
        SESSION_LINES=$((TOTAL_LINES - SESSION_START))
        EDIT_COUNT=$(tail -n "$SESSION_LINES" "$SCRATCHPAD" | grep -c "VERIFY:" || echo "0")
        BUDGET_COUNT=$(tail -n "$SESSION_LINES" "$SCRATCHPAD" | grep -c "BUDGET:" || echo "0")

        echo "Session Summary: $EDIT_COUNT edits verified, $BUDGET_COUNT spawns tracked" >> "$SCRATCHPAD"
    fi
fi

# Archive session in budget history
if [[ -f "$BUDGET_FILE" ]]; then
    # Use Python for JSON manipulation (more reliable than sed)
    python3 << EOF
import json
from datetime import datetime
from pathlib import Path

budget_path = Path("$BUDGET_FILE")
try:
    with open(budget_path, 'r') as f:
        budget = json.load(f)

    current_session = budget.get("current_session", {})

    # Only archive if session had activity
    if current_session.get("spawns"):
        history = budget.setdefault("history", {
            "total_cost": 0,
            "total_spawns": 0,
            "total_input_tokens": 0,
            "total_output_tokens": 0,
            "sessions": []
        })

        # Update totals
        session_cost = current_session.get("estimated_cost", 0)
        session_spawns = len(current_session.get("spawns", []))

        history["total_cost"] = history.get("total_cost", 0) + session_cost
        history["total_spawns"] = history.get("total_spawns", 0) + session_spawns

        # Archive session
        session_record = {
            "id": current_session.get("id"),
            "started": current_session.get("started"),
            "ended": datetime.now().isoformat(),
            "cost": session_cost,
            "spawns": session_spawns
        }

        sessions = history.setdefault("sessions", [])
        sessions.append(session_record)

        # Keep only last 100 sessions
        if len(sessions) > 100:
            sessions = sessions[-100:]
            history["sessions"] = sessions

        budget["history"] = history

        # Reset current session
        budget["current_session"] = {
            "id": None,
            "started": None,
            "cluster": None,
            "model_tier": current_session.get("model_tier", "quality"),
            "estimated_cost": 0,
            "spawns": []
        }

        with open(budget_path, 'w') as f:
            json.dump(budget, f, indent=2)

        print(f"Session archived: \${session_cost:.2f}, {session_spawns} spawns")

except Exception as e:
    print(f"Warning: Could not archive session: {e}")
EOF
fi

exit 0