#!/bin/bash
# compare.sh — A/B test: Qwen vanilla vs Qwen with Claude Style
#
# Strategy:
# 1. Clean up any previous test artifacts
# 2. Uninstall extension
# 3. Run vanilla Qwen (YOLO mode)
# 4. Clean up the file it created
# 5. Reinstall extension
# 6. Run Qwen with Claude Style (YOLO mode)
# 7. Compare

set -e

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

PROMPT="Create a Python file 'user_database.py' with a class that manages a user database. Support adding users, finding by email, updating fields, and deleting users. Store in a dict."
OUTPUT_DIR="/tmp/claude-style-compare"
SCRIPT_PATH="$(readlink -f "${BASH_SOURCE[0]}")"
SCRIPT_DIR="$(cd "$(dirname "$SCRIPT_PATH")" && pwd)"
PROJECT_DIR="$(cd "$SCRIPT_DIR/.." && pwd)"
WORK_DIR=""
mkdir -p "$OUTPUT_DIR"

cleanup_work_dir() {
    if [ -n "$WORK_DIR" ] && [ -d "$WORK_DIR" ]; then
        rm -rf "$WORK_DIR"
    fi
}
trap cleanup_work_dir EXIT

echo -e "${BLUE}╔══════════════════════════════════════════╗${NC}"
echo -e "${BLUE}║    Claude Style — A/B Comparison         ║${NC}"
echo -e "${BLUE}╚══════════════════════════════════════════╝${NC}"
echo ""
echo -e "${YELLOW}Prompt:${NC}"
echo "  $PROMPT"
echo ""

# ============================================================
# Step 1: Uninstall
# ============================================================
echo -e "${YELLOW}[1/4] Uninstalling Claude Style...${NC}"
"$PROJECT_DIR/scripts/uninstall.sh" 2>&1 | grep -E "^(✓|✗|⚠|Removing|Extension)" || true
echo ""

# ============================================================
# Step 2: Vanilla Qwen
# ============================================================
echo -e "${YELLOW}[2/4] Running Qwen WITHOUT Claude Style...${NC}"
WORK_DIR=$(mktemp -d)
pushd "$WORK_DIR" > /dev/null
qwen --prompt "$PROMPT" --output-format text -y 2>&1 | tee "$OUTPUT_DIR/vanilla.txt" > /dev/null 2>&1
if [ -f "$WORK_DIR/user_database.py" ]; then
    cp "$WORK_DIR/user_database.py" "$OUTPUT_DIR/vanilla_code.py"
    VANILLA_CODE_LINES=$(wc -l < "$OUTPUT_DIR/vanilla_code.py")
    echo -e "${GREEN}✓ Generated $VANILLA_CODE_LINES lines of code${NC}"
else
    echo -e "${YELLOW}⚠ No file generated${NC}"
fi
popd > /dev/null
WORK_DIR=""
echo -e "${GREEN}✓ Vanilla output saved${NC}"
echo ""

# ============================================================
# Step 3: Clean up + Reinstall
# ============================================================
echo -e "${YELLOW}[3/4] Cleaning up + reinstalling Claude Style...${NC}"
"$PROJECT_DIR/scripts/setup.sh" 2>&1 | grep -E "^(✓|✗|⚠|Installing|Extension|Configuring)" || true
echo ""

# ============================================================
# Step 4: Qwen with Claude Style
# ============================================================
echo -e "${YELLOW}[4/4] Running Qwen WITH Claude Style...${NC}"
WORK_DIR=$(mktemp -d)
pushd "$WORK_DIR" > /dev/null
qwen --prompt "$PROMPT" --output-format text -y 2>&1 | tee "$OUTPUT_DIR/claude-style.txt" > /dev/null 2>&1
if [ -f "$WORK_DIR/user_database.py" ]; then
    cp "$WORK_DIR/user_database.py" "$OUTPUT_DIR/claude-style_code.py"
    STYLE_CODE_LINES=$(wc -l < "$OUTPUT_DIR/claude-style_code.py")
    echo -e "${GREEN}✓ Generated $STYLE_CODE_LINES lines of code${NC}"
else
    echo -e "${YELLOW}⚠ No file generated${NC}"
fi
popd > /dev/null
WORK_DIR=""
echo -e "${GREEN}✓ Claude Style output saved${NC}"
echo ""

# ============================================================
# Comparison
# ============================================================
echo -e "${BLUE}────────────────────────────────────────────${NC}"
echo -e "${YELLOW}─── Vanilla Qwen ───${NC}"
cat "$OUTPUT_DIR/vanilla.txt"
echo ""

echo -e "${YELLOW}─── Claude Style ───${NC}"
cat "$OUTPUT_DIR/claude-style.txt"
echo ""

# Metrics
echo -e "${BLUE}─── Metrics ───${NC}"
VANILLA_LINES=$(wc -l < "$OUTPUT_DIR/vanilla.txt")
STYLE_LINES=$(wc -l < "$OUTPUT_DIR/claude-style.txt")
VANILLA_WORDS=$(wc -w < "$OUTPUT_DIR/vanilla.txt")
STYLE_WORDS=$(wc -w < "$OUTPUT_DIR/claude-style.txt")

echo "  Vanilla:       $VANILLA_LINES lines, $VANILLA_WORDS words"
echo "  Claude Style:  $STYLE_LINES lines, $STYLE_WORDS words"
echo ""

# Code extraction using awk (extracts ALL code blocks)
VANILLA_CODE=$(awk '/```/{p=!p; next} p{print}' "$OUTPUT_DIR/vanilla.txt" || true)
STYLE_CODE=$(awk '/```/{p=!p; next} p{print}' "$OUTPUT_DIR/claude-style.txt" || true)

VANILLA_CODE_LINES=$(printf '%s\n' "$VANILLA_CODE" | grep -c . 2>/dev/null || echo "0")
STYLE_CODE_LINES=$(printf '%s\n' "$STYLE_CODE" | grep -c . 2>/dev/null || echo "0")

echo -e "${BLUE}─── Code Only ───${NC}"
echo "  Vanilla:       $VANILLA_CODE_LINES lines of code"
echo "  Claude Style:  $STYLE_CODE_LINES lines of code"
echo ""

# Text comparison (excluding code blocks)
VANILLA_TEXT_LINES=$(sed '/```/,/```/d' "$OUTPUT_DIR/vanilla.txt" | grep -c . || echo "0")
STYLE_TEXT_LINES=$(sed '/```/,/```/d' "$OUTPUT_DIR/claude-style.txt" | grep -c . || echo "0")

echo -e "${BLUE}─── Explanatory Text ───${NC}"
echo "  Vanilla:       $VANILLA_TEXT_LINES lines"
echo "  Claude Style:  $STYLE_TEXT_LINES lines"
echo ""

if [ "$STYLE_TEXT_LINES" -lt "$VANILLA_TEXT_LINES" ]; then
    echo -e "${GREEN}✓ Claude Style is $(( VANILLA_TEXT_LINES - STYLE_TEXT_LINES )) fewer lines of explanation${NC}"
elif [ "$STYLE_TEXT_LINES" -eq "$VANILLA_TEXT_LINES" ]; then
    echo -e "${YELLOW}⚠ Similar explanation length${NC}"
else
    echo -e "${YELLOW}⚠ Claude Style used more explanation${NC}"
fi

echo ""
echo -e "${YELLOW}Full outputs:${NC}"
echo "  diff $OUTPUT_DIR/vanilla.txt $OUTPUT_DIR/claude-style.txt"
echo ""
echo -e "${BLUE}────────────────────────────────────────────${NC}"