#!/usr/bin/env bash
# ═══════════════════════════════════════════════════════════════════════════
# start.sh — Inference launcher for the Dyslexia Academic Writing Corrector
# ═══════════════════════════════════════════════════════════════════════════
#
# Usage:
#   bash start.sh --cli                  # Interactive REPL mode
#   bash start.sh --api                  # FastAPI server mode
#   bash start.sh --cli --text "..."     # Single text correction
#   bash start.sh --api --port 8080      # Custom port
#
set -euo pipefail

# ── Defaults ────────────────────────────────────────────────────────────────
MODE=""
CONFIG="configs/inference_config.yaml"
TEXT=""
MASTER_COPY=""
STYLE_ALPHA="0.6"
PORT="8000"
WORKERS="1"

# ── Colors ──────────────────────────────────────────────────────────────────
GREEN='\033[0;32m'
CYAN='\033[0;36m'
YELLOW='\033[1;33m'
RED='\033[0;31m'
BOLD='\033[1m'
NC='\033[0m'

# ── Parse arguments ────────────────────────────────────────────────────────
while [[ $# -gt 0 ]]; do
    case $1 in
        --cli)         MODE="cli"; shift ;;
        --api)         MODE="api"; shift ;;
        --config)      CONFIG="$2"; shift 2 ;;
        --config=*)    CONFIG="${1#*=}"; shift ;;
        --text)        TEXT="$2"; shift 2 ;;
        --text=*)      TEXT="${1#*=}"; shift ;;
        --master-copy) MASTER_COPY="$2"; shift 2 ;;
        --port)        PORT="$2"; shift 2 ;;
        --port=*)      PORT="${1#*=}"; shift ;;
        --workers)     WORKERS="$2"; shift 2 ;;
        --alpha)       STYLE_ALPHA="$2"; shift 2 ;;
        -h|--help)
            echo "Usage: bash start.sh [--cli|--api] [OPTIONS]"
            echo ""
            echo "Modes:"
            echo "  --cli          Interactive REPL or single-text correction"
            echo "  --api          Start FastAPI server"
            echo ""
            echo "Options:"
            echo "  --config PATH  Config file (default: configs/inference_config.yaml)"
            echo "  --text TEXT    Text to correct (CLI mode, skip interactive)"
            echo "  --master-copy  Optional master copy for style matching"
            echo "  --alpha FLOAT  Style blend weight 0-1 (default: 0.6)"
            echo "  --port PORT    API server port (default: 8000)"
            echo "  --workers N    API server workers (default: 1)"
            exit 0
            ;;
        *) echo -e "${RED}Unknown option: $1${NC}"; exit 1 ;;
    esac
done

# ── Python detection ───────────────────────────────────────────────────────
if command -v python3 &>/dev/null; then
    PYTHON=python3
elif command -v python &>/dev/null; then
    PYTHON=python
else
    echo -e "${RED}Python not found!${NC}"
    exit 1
fi

# ── Mode selection ─────────────────────────────────────────────────────────
if [ -z "$MODE" ]; then
    echo ""
    echo -e "${BOLD}╔══════════════════════════════════════════════════════════╗${NC}"
    echo -e "${BOLD}║  Dyslexia Academic Writing Corrector — Inference        ║${NC}"
    echo -e "${BOLD}╚══════════════════════════════════════════════════════════╝${NC}"
    echo ""
    echo -e "  ${CYAN}1)${NC} Interactive CLI (REPL)"
    echo -e "  ${CYAN}2)${NC} API Server (FastAPI)"
    echo ""
    read -rp "  Select mode [1/2]: " choice
    case "$choice" in
        1) MODE="cli" ;;
        2) MODE="api" ;;
        *) MODE="cli" ;;
    esac
fi

# ── Check model exists ────────────────────────────────────────────────────
if [ ! -d "checkpoints/best_model" ]; then
    echo -e "${YELLOW}[WARN] No trained model found at checkpoints/best_model${NC}"
    echo -e "${YELLOW}       Will use base model. Run train.sh first for best results.${NC}"
fi

# ── Launch ─────────────────────────────────────────────────────────────────
case "$MODE" in
    cli)
        echo -e "${GREEN}Starting CLI inference...${NC}"
        CLI_ARGS="--config $CONFIG --style-alpha $STYLE_ALPHA"
        if [ -n "$TEXT" ]; then
            CLI_ARGS="$CLI_ARGS --text \"$TEXT\""
        fi
        if [ -n "$MASTER_COPY" ]; then
            CLI_ARGS="$CLI_ARGS --master-copy \"$MASTER_COPY\""
        fi
        eval $PYTHON scripts/run_inference.py $CLI_ARGS
        ;;
    api)
        echo -e "${GREEN}Starting API server on port $PORT...${NC}"
        echo -e "  Docs:   ${CYAN}http://localhost:$PORT/docs${NC}"
        echo -e "  Health: ${CYAN}http://localhost:$PORT/health${NC}"
        echo ""
        $PYTHON -m uvicorn src.api.main:app \
            --host 0.0.0.0 \
            --port "$PORT" \
            --workers "$WORKERS" \
            --log-level info
        ;;
esac