#!/bin/bash
# =============================================================================
# MPC System - Deployment Script
# =============================================================================
# This script manages the MPC System Docker services
#
# Deployment Modes:
#   1. Development (default): All services on one machine (docker-compose.yml)
#   2. Production Central:    Central services only (docker-compose.prod.yml)
#   3. Production Party:      Standalone party (docker-compose.party.yml)
#
# External Ports (Development):
#   4000  - Account Service HTTP API
#   8081  - Session Coordinator API
#   8082  - Message Router HTTP
#   8083  - Server Party API
#
# External Ports (Production Central):
#   50051 - Message Router gRPC (for party connections)
#   50052 - Session Coordinator gRPC (for party connections)
#   4000  - Account Service HTTP API
#   8081  - Session Coordinator HTTP API
#   8082  - Message Router HTTP API
# =============================================================================

set -e

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

log_info() { echo -e "${BLUE}[INFO]${NC} $1"; }
log_success() { echo -e "${GREEN}[OK]${NC} $1"; }
log_warn() { echo -e "${YELLOW}[WARN]${NC} $1"; }
log_error() { echo -e "${RED}[ERROR]${NC} $1"; }
log_header() { echo -e "${CYAN}=== $1 ===${NC}"; }

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
cd "$SCRIPT_DIR"

# Determine which environment file to load
load_env() {
    local env_file="$1"
    if [ -f "$env_file" ]; then
        log_info "Loading environment from $env_file"
        set -a
        source "$env_file"
        set +a
        return 0
    fi
    return 1
}

# Load environment based on mode
load_environment() {
    local mode="$1"
    case "$mode" in
        prod)
            load_env ".env.prod" || load_env ".env" || {
                log_error "No .env.prod or .env file found"
                exit 1
            }
            ;;
        party)
            load_env ".env.party" || {
                log_error "No .env.party file found. Create from .env.party.example"
                exit 1
            }
            ;;
        *)
            load_env ".env" || {
                if [ -f ".env.example" ]; then
                    log_warn ".env file not found. Creating from .env.example"
                    cp .env.example .env
                    log_error "Please configure .env file and run again"
                fi
                exit 1
            }
            ;;
    esac
}

# Service lists
CORE_SERVICES="postgres"
DEV_MPC_SERVICES="session-coordinator message-router server-party-1 server-party-2 server-party-3 server-party-api account-service"
PROD_CENTRAL_SERVICES="postgres message-router session-coordinator account-service server-party-api"

# ============================================
# Development Mode Commands (docker-compose.yml)
# ============================================
dev_commands() {
    load_environment "dev"

    case "$1" in
        build)
            log_info "Building MPC System services..."
            docker compose build
            log_success "MPC System built successfully"
            ;;

        build-no-cache)
            log_info "Building MPC System (no cache)..."
            docker compose build --no-cache
            log_success "MPC System built successfully"
            ;;

        up|start)
            log_info "Starting MPC System (Development)..."
            docker compose up -d
            log_success "MPC System started"
            echo ""
            log_info "Services status:"
            docker compose ps
            ;;

        down|stop)
            log_info "Stopping MPC System..."
            docker compose down
            log_success "MPC System stopped"
            ;;

        restart)
            log_info "Restarting MPC System..."
            docker compose down
            docker compose up -d
            log_success "MPC System restarted"
            ;;

        logs)
            if [ -n "$2" ]; then
                docker compose logs -f "$2"
            else
                docker compose logs -f
            fi
            ;;

        logs-tail)
            if [ -n "$2" ]; then
                docker compose logs --tail 100 "$2"
            else
                docker compose logs --tail 100
            fi
            ;;

        status|ps)
            log_info "MPC System status:"
            docker compose ps
            ;;

        health)
            log_info "Checking MPC System health..."
            echo ""
            log_header "Infrastructure"
            for svc in $CORE_SERVICES; do
                if docker compose ps "$svc" --format json 2>/dev/null | grep -q '"Health":"healthy"'; then
                    log_success "$svc is healthy"
                else
                    log_warn "$svc is not healthy"
                fi
            done

            echo ""
            log_header "MPC Services"
            for svc in $DEV_MPC_SERVICES; do
                if docker compose ps "$svc" --format json 2>/dev/null | grep -q '"Health":"healthy"'; then
                    log_success "$svc is healthy"
                else
                    log_warn "$svc is not healthy"
                fi
            done

            echo ""
            log_header "External API"
            if curl -sf "http://localhost:4000/health" > /dev/null 2>&1; then
                log_success "Account Service API (port 4000) is accessible"
            else
                log_error "Account Service API (port 4000) is not accessible"
            fi
            ;;

        clean)
            log_warn "This will remove all containers and volumes!"
            read -p "Are you sure? (y/N) " -n 1 -r
            echo
            if [[ $REPLY =~ ^[Yy]$ ]]; then
                docker compose down -v
                log_success "MPC System cleaned"
            else
                log_info "Cancelled"
            fi
            ;;

        shell)
            if [ -n "$2" ]; then
                log_info "Opening shell in $2..."
                docker compose exec "$2" sh
            else
                log_info "Opening shell in account-service..."
                docker compose exec account-service sh
            fi
            ;;

        test-api)
            log_info "Testing Account Service API..."
            echo ""
            echo "Health check:"
            curl -s "http://localhost:4000/health" | jq . 2>/dev/null || curl -s "http://localhost:4000/health"
            echo ""
            ;;

        *)
            return 1
            ;;
    esac
}

# ============================================
# Production Central Commands (docker-compose.prod.yml)
# ============================================
prod_commands() {
    load_environment "prod"

    case "$1" in
        build)
            log_info "Building Production Central services..."
            docker compose -f docker-compose.prod.yml build
            log_success "Production services built"
            ;;

        up|start)
            log_info "Starting Production Central services..."
            docker compose -f docker-compose.prod.yml up -d
            log_success "Production Central services started"
            echo ""
            log_header "Services Status"
            docker compose -f docker-compose.prod.yml ps
            echo ""
            log_header "Public Endpoints"
            echo "  Message Router gRPC:     ${MESSAGE_ROUTER_GRPC_PORT:-50051}"
            echo "  Session Coordinator gRPC: ${SESSION_COORDINATOR_GRPC_PORT:-50052}"
            echo "  Account Service HTTP:     ${ACCOUNT_SERVICE_PORT:-4000}"
            ;;

        down|stop)
            log_info "Stopping Production Central services..."
            docker compose -f docker-compose.prod.yml down
            log_success "Production Central services stopped"
            ;;

        restart)
            log_info "Restarting Production Central services..."
            docker compose -f docker-compose.prod.yml down
            docker compose -f docker-compose.prod.yml up -d
            log_success "Production Central services restarted"
            ;;

        logs)
            if [ -n "$2" ]; then
                docker compose -f docker-compose.prod.yml logs -f "$2"
            else
                docker compose -f docker-compose.prod.yml logs -f
            fi
            ;;

        status|ps)
            log_info "Production Central status:"
            docker compose -f docker-compose.prod.yml ps
            ;;

        health)
            log_info "Checking Production Central health..."
            echo ""
            log_header "Central Services"
            for svc in $PROD_CENTRAL_SERVICES; do
                if docker compose -f docker-compose.prod.yml ps "$svc" --format json 2>/dev/null | grep -q '"Health":"healthy"'; then
                    log_success "$svc is healthy"
                else
                    log_warn "$svc is not healthy"
                fi
            done

            echo ""
            log_header "Public Endpoints"
            if curl -sf "http://localhost:${MESSAGE_ROUTER_HTTP_PORT:-8082}/health" > /dev/null 2>&1; then
                log_success "Message Router (port ${MESSAGE_ROUTER_HTTP_PORT:-8082}) is accessible"
            else
                log_error "Message Router is not accessible"
            fi
            if curl -sf "http://localhost:${SESSION_COORDINATOR_HTTP_PORT:-8081}/health" > /dev/null 2>&1; then
                log_success "Session Coordinator (port ${SESSION_COORDINATOR_HTTP_PORT:-8081}) is accessible"
            else
                log_error "Session Coordinator is not accessible"
            fi
            if curl -sf "http://localhost:${ACCOUNT_SERVICE_PORT:-4000}/health" > /dev/null 2>&1; then
                log_success "Account Service (port ${ACCOUNT_SERVICE_PORT:-4000}) is accessible"
            else
                log_error "Account Service is not accessible"
            fi
            ;;

        clean)
            log_warn "This will remove all Production Central containers and volumes!"
            read -p "Are you sure? (y/N) " -n 1 -r
            echo
            if [[ $REPLY =~ ^[Yy]$ ]]; then
                docker compose -f docker-compose.prod.yml down -v
                log_success "Production Central cleaned"
            else
                log_info "Cancelled"
            fi
            ;;

        *)
            return 1
            ;;
    esac
}

# ============================================
# Production Party Commands (docker-compose.party.yml)
# ============================================
party_commands() {
    load_environment "party"

    # Validate required environment variables
    if [ -z "$PARTY_ID" ]; then
        log_error "PARTY_ID must be set (e.g., server-party-1)"
        exit 1
    fi
    if [ -z "$MESSAGE_ROUTER_ADDR" ]; then
        log_error "MESSAGE_ROUTER_ADDR must be set (e.g., grpc.mpc.example.com:50051)"
        exit 1
    fi

    case "$1" in
        build)
            log_info "Building Party ($PARTY_ID)..."
            docker compose -f docker-compose.party.yml build
            log_success "Party built"
            ;;

        up|start)
            log_info "Starting Party: $PARTY_ID"
            log_info "Connecting to Message Router: $MESSAGE_ROUTER_ADDR"
            docker compose -f docker-compose.party.yml up -d
            log_success "Party $PARTY_ID started"
            echo ""
            docker compose -f docker-compose.party.yml ps
            ;;

        down|stop)
            log_info "Stopping Party: $PARTY_ID..."
            docker compose -f docker-compose.party.yml down
            log_success "Party stopped"
            ;;

        restart)
            log_info "Restarting Party: $PARTY_ID..."
            docker compose -f docker-compose.party.yml down
            docker compose -f docker-compose.party.yml up -d
            log_success "Party restarted"
            ;;

        logs)
            docker compose -f docker-compose.party.yml logs -f server-party
            ;;

        status|ps)
            log_info "Party $PARTY_ID status:"
            docker compose -f docker-compose.party.yml ps
            ;;

        health)
            log_info "Checking Party $PARTY_ID health..."
            echo ""
            log_header "Local Services"
            if docker compose -f docker-compose.party.yml ps postgres --format json 2>/dev/null | grep -q '"Health":"healthy"'; then
                log_success "Local PostgreSQL is healthy"
            else
                log_warn "Local PostgreSQL is not healthy"
            fi
            if docker compose -f docker-compose.party.yml ps server-party --format json 2>/dev/null | grep -q '"Health":"healthy"'; then
                log_success "Server Party is healthy"
            else
                log_warn "Server Party is not healthy"
            fi

            echo ""
            log_header "Central Service Connectivity"
            # Extract host and port from MESSAGE_ROUTER_ADDR
            MR_HOST=$(echo "$MESSAGE_ROUTER_ADDR" | cut -d: -f1)
            MR_PORT=$(echo "$MESSAGE_ROUTER_ADDR" | cut -d: -f2)
            if timeout 5 bash -c "echo >/dev/tcp/$MR_HOST/$MR_PORT" 2>/dev/null; then
                log_success "Message Router ($MESSAGE_ROUTER_ADDR) is reachable"
            else
                log_error "Message Router ($MESSAGE_ROUTER_ADDR) is NOT reachable"
            fi
            ;;

        clean)
            log_warn "This will remove Party $PARTY_ID containers and LOCAL KEY STORAGE!"
            log_warn "Your encrypted key shares will be DELETED!"
            read -p "Are you absolutely sure? (yes/N) " confirm
            echo
            if [ "$confirm" = "yes" ]; then
                docker compose -f docker-compose.party.yml down -v
                log_success "Party $PARTY_ID cleaned"
            else
                log_info "Cancelled"
            fi
            ;;

        *)
            return 1
            ;;
    esac
}

# ============================================
# Main Command Router
# ============================================
show_help() {
    echo "MPC System Deployment Script"
    echo ""
    echo "Usage: $0 <mode> <command> [options]"
    echo ""
    echo "Deployment Modes:"
    echo "  (default)     Development mode - all services on one machine"
    echo "  prod          Production Central - Message Router, Session Coordinator, Account"
    echo "  party         Production Party - standalone server-party (distributed)"
    echo ""
    echo "Development Commands (default mode):"
    echo "  $0 build              Build all Docker images"
    echo "  $0 up|start           Start all services"
    echo "  $0 down|stop          Stop all services"
    echo "  $0 restart            Restart all services"
    echo "  $0 logs [service]     Follow logs"
    echo "  $0 status|ps          Show services status"
    echo "  $0 health             Check all services health"
    echo "  $0 clean              Remove containers and volumes"
    echo ""
    echo "Production Central Commands:"
    echo "  $0 prod build         Build central services"
    echo "  $0 prod up            Start central services"
    echo "  $0 prod down          Stop central services"
    echo "  $0 prod logs          Follow central logs"
    echo "  $0 prod health        Check central health"
    echo ""
    echo "Production Party Commands (run on each party machine):"
    echo "  $0 party build        Build party service"
    echo "  $0 party up           Start party (connects to central)"
    echo "  $0 party down         Stop party"
    echo "  $0 party logs         Follow party logs"
    echo "  $0 party health       Check party health and connectivity"
    echo ""
    echo "Environment Files:"
    echo "  .env              Development configuration"
    echo "  .env.prod         Production Central configuration"
    echo "  .env.party        Production Party configuration"
    echo ""
    echo "Examples:"
    echo "  # Development (all on one machine)"
    echo "  $0 up"
    echo ""
    echo "  # Production Central (on central server)"
    echo "  $0 prod up"
    echo ""
    echo "  # Production Party (on each party machine)"
    echo "  PARTY_ID=server-party-1 $0 party up"
    echo "  PARTY_ID=server-party-2 $0 party up"
    echo "  PARTY_ID=server-party-3 $0 party up"
}

# Route commands based on first argument
case "$1" in
    prod)
        shift
        prod_commands "$@" || {
            echo "Usage: $0 prod {build|up|down|restart|logs|status|health|clean}"
            exit 1
        }
        ;;

    party)
        shift
        party_commands "$@" || {
            echo "Usage: $0 party {build|up|down|restart|logs|status|health|clean}"
            echo ""
            echo "Required environment variables:"
            echo "  PARTY_ID                  Unique party identifier"
            echo "  MESSAGE_ROUTER_ADDR       Central Message Router address (only connection needed)"
            echo "  CRYPTO_MASTER_KEY         Encryption key for key shares"
            exit 1
        }
        ;;

    help|--help|-h)
        show_help
        ;;

    "")
        show_help
        exit 1
        ;;

    *)
        # Default to development mode
        dev_commands "$@" || {
            show_help
            exit 1
        }
        ;;
esac