#!/bin/bash
# MCP-ZERO LLM Container Setup
# Configures LLM support for MCP-ZERO agents in containerized environments
# Implements offline-first resilience pattern

# Generate LLM configuration
generate_llm_config() {
  local config_dir="$1"
  local offline_mode="$2"
  
  echo "Generating LLM configuration..."
  
  # Convert offline_mode to JSON boolean
  if [ "$offline_mode" -eq 1 ]; then
    offline_json="true"
  else
    offline_json="false"
  fi
  
  cat > "${config_dir}/llm.json" << EOF
{
  "llm": {
    "provider": "zeta",
    "models": {
      "default": "zeta-llm-small",
      "enhanced": "zeta-llm-medium",
      "offline": "zeta-llm-tiny"
    },
    "parameters": {
      "temperature": 0.7,
      "top_p": 0.9,
      "max_tokens": 1024
    },
    "offline_mode": {
      "enabled": ${offline_json},
      "fallback_model": "zeta-llm-tiny",
      "connection_attempts": 1,
      "permanent_fallback": true
    }
  },
  "embedding": {
    "model": "zeta-embedding",
    "dimensions": 384,
    "normalize": true,
    "offline_capable": true
  },
  "container": {
    "resource_limits": {
      "cpu": 2,
      "memory": "2Gi",
      "gpu_enabled": true,
      "gpu_fallback_cpu": true
    }
  }
}
EOF

  echo "LLM configuration generated successfully at ${config_dir}/llm.json"
}

# Generate container orchestration for LLMs
generate_llm_docker() {
  local docker_dir="$1"
  local offline_mode="$2"
  
  echo "Generating LLM container configuration..."
  
  cat > "${docker_dir}/docker-compose.llm.yaml" << EOF
version: '3.8'

services:
  # ZetaLLM Service
  mcp-zetallm:
    build:
      context: .
      dockerfile: docker/Dockerfile.zetallm
    volumes:
      - ./.mcp/config:/app/config
      - ./.mcp/locks:/app/locks
      - ./agents/models/language:/app/models
    environment:
      - MCP_CONFIG_PATH=/app/config/llm.json
      - MCP_OFFLINE_MODE=$([ "$offline_mode" -eq 1 ] && echo "true" || echo "false")
    restart: unless-stopped
    networks:
      - mcp-network
    deploy:
      resources:
        limits:
          cpus: '2'
          memory: 2G
    healthcheck:
      test: ["CMD", "curl", "-f", "http://localhost:8000/health"]
      interval: 30s
      timeout: 10s
      retries: 3
      start_period: 40s

  # ZetaTensor Service
  mcp-zetatensor:
    build:
      context: .
      dockerfile: docker/Dockerfile.zetatensor
    volumes:
      - ./.mcp/config:/app/config
      - ./.mcp/locks:/app/locks
      - ./agents/models:/app/models
    environment:
      - MCP_CONFIG_PATH=/app/config/llm.json
      - MCP_OFFLINE_MODE=$([ "$offline_mode" -eq 1 ] && echo "true" || echo "false")
    restart: unless-stopped
    networks:
      - mcp-network
    deploy:
      resources:
        limits:
          cpus: '1'
          memory: 1G

  # ZetaFace Service (for vision models)
  mcp-zetaface:
    build:
      context: .
      dockerfile: docker/Dockerfile.zetaface
    volumes:
      - ./.mcp/config:/app/config
      - ./.mcp/locks:/app/locks
      - ./agents/models/vision:/app/models
    environment:
      - MCP_CONFIG_PATH=/app/config/llm.json
      - MCP_OFFLINE_MODE=$([ "$offline_mode" -eq 1 ] && echo "true" || echo "false")
    restart: unless-stopped
    networks:
      - mcp-network
    deploy:
      resources:
        limits:
          cpus: '1'
          memory: 1G

  # ZetaCUDA Service (optional GPU acceleration)
  mcp-zetacuda:
    build:
      context: .
      dockerfile: docker/Dockerfile.zetacuda
    volumes:
      - ./.mcp/config:/app/config
      - ./.mcp/locks:/app/locks
    environment:
      - MCP_CONFIG_PATH=/app/config/llm.json
      - MCP_OFFLINE_MODE=$([ "$offline_mode" -eq 1 ] && echo "true" || echo "false")
      - NVIDIA_VISIBLE_DEVICES=all
    restart: unless-stopped
    networks:
      - mcp-network
    runtime: nvidia
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: all
              capabilities: [gpu]

networks:
  mcp-network:
    external: true
EOF

  echo "LLM container configuration generated successfully at ${docker_dir}/docker-compose.llm.yaml"

  # Generate Dockerfiles for ZETA components
  # ZetaLLM Dockerfile
  cat > "${docker_dir}/Dockerfile.zetallm" << EOF
FROM python:3.9-slim

WORKDIR /app

# Install system dependencies
RUN apt-get update && apt-get install -y --no-install-recommends \\
    build-essential \\
    curl \\
    && rm -rf /var/lib/apt/lists/*

# Copy requirements first for better caching
COPY ./requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

# Copy source code
COPY ./src/zeta_llm ./src/zeta_llm
COPY ./src/common ./src/common

# Create directories for volumes
RUN mkdir -p /app/config /app/locks /app/models

# Copy offline check script
COPY ./src/zeta_llm/scripts/offline_check.sh /app/offline_check.sh
RUN chmod +x /app/offline_check.sh

# Create offline-ready wrapper for the LLM
RUN echo '#!/usr/bin/env python3\\n\\
import os, sys, json\\n\\
sys.path.append("/app")\\n\\
\\n\\
# Implement single connection attempt pattern\\n\\
def check_offline():\\n\\
    """Check if system should run in offline mode"""\\n\\
    if os.path.exists("/app/locks/offline.lock"):\\n\\
        print("Offline lock file detected, running in offline mode")\\n\\
        os.environ["MCP_OFFLINE_MODE"] = "true"\\n\\
        return True\\n\\
    return False\\n\\
\\n\\
# Check offline status before importing rest of system\\n\\
is_offline = check_offline()\\n\\
\\n\\
from src.zeta_llm.server import main\\n\\
\\n\\
if __name__ == "__main__":\\n\\
    main(offline_mode=is_offline)' > /app/run.py\\n\\

# Set entrypoint with offline check
ENTRYPOINT ["/bin/sh", "-c", "/app/offline_check.sh && python /app/run.py"]

# Expose port for health checks
EXPOSE 8000
EOF

  # ZetaTensor Dockerfile
  cat > "${docker_dir}/Dockerfile.zetatensor" << EOF
FROM python:3.9-slim

WORKDIR /app

# Install system dependencies
RUN apt-get update && apt-get install -y --no-install-recommends \\
    build-essential \\
    curl \\
    && rm -rf /var/lib/apt/lists/*

# Copy requirements first for better caching
COPY ./requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

# Copy source code
COPY ./src/zeta_tensor ./src/zeta_tensor
COPY ./src/common ./src/common

# Create directories for volumes
RUN mkdir -p /app/config /app/locks /app/models

# Copy offline check script
COPY ./src/zeta_tensor/scripts/offline_check.sh /app/offline_check.sh
RUN chmod +x /app/offline_check.sh

# Create wrapper script with offline-first pattern
RUN echo '#!/usr/bin/env python3\\n\\
import os, sys\\n\\
sys.path.append("/app")\\n\\
\\n\\
# Implement single connection attempt pattern\\n\\
def check_offline():\\n\\
    """Check if system should run in offline mode"""\\n\\
    if os.path.exists("/app/locks/offline.lock"):\\n\\
        print("Offline lock file detected, running in offline mode")\\n\\
        os.environ["MCP_OFFLINE_MODE"] = "true"\\n\\
        return True\\n\\
    return False\\n\\
\\n\\
# Check offline status before importing rest of system\\n\\
is_offline = check_offline()\\n\\
\\n\\
from src.zeta_tensor.service import main\\n\\
\\n\\
if __name__ == "__main__":\\n\\
    main(offline_mode=is_offline)' > /app/run.py\\n\\

# Set entrypoint with offline check
ENTRYPOINT ["/bin/sh", "-c", "/app/offline_check.sh && python /app/run.py"]
EOF

  # ZetaFace Dockerfile
  cat > "${docker_dir}/Dockerfile.zetaface" << EOF
FROM python:3.9-slim

WORKDIR /app

# Install system dependencies
RUN apt-get update && apt-get install -y --no-install-recommends \\
    build-essential \\
    libgl1-mesa-glx \\
    libglib2.0-0 \\
    && rm -rf /var/lib/apt/lists/*

# Copy requirements first for better caching
COPY ./requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

# Copy source code
COPY ./src/zeta_face ./src/zeta_face
COPY ./src/common ./src/common

# Create directories for volumes
RUN mkdir -p /app/config /app/locks /app/models

# Copy offline check script
COPY ./src/zeta_face/scripts/offline_check.sh /app/offline_check.sh
RUN chmod +x /app/offline_check.sh

# Create wrapper script with offline-first pattern
RUN echo '#!/usr/bin/env python3\\n\\
import os, sys\\n\\
sys.path.append("/app")\\n\\
\\n\\
# Implement single connection attempt pattern\\n\\
def check_offline():\\n\\
    """Check if system should run in offline mode"""\\n\\
    if os.path.exists("/app/locks/offline.lock"):\\n\\
        print("Offline lock file detected, running in offline mode")\\n\\
        os.environ["MCP_OFFLINE_MODE"] = "true"\\n\\
        return True\\n\\
    return False\\n\\
\\n\\
# Check offline status before importing rest of system\\n\\
is_offline = check_offline()\\n\\
\\n\\
from src.zeta_face.service import main\\n\\
\\n\\
if __name__ == "__main__":\\n\\
    main(offline_mode=is_offline)' > /app/run.py\\n\\

# Set entrypoint with offline check
ENTRYPOINT ["/bin/sh", "-c", "/app/offline_check.sh && python /app/run.py"]
EOF

  # ZetaCUDA Dockerfile
  cat > "${docker_dir}/Dockerfile.zetacuda" << EOF
FROM nvidia/cuda:11.6.2-cudnn8-runtime-ubuntu20.04

WORKDIR /app

# Install system dependencies
RUN apt-get update && apt-get install -y --no-install-recommends \\
    python3.9 \\
    python3-pip \\
    python3-setuptools \\
    build-essential \\
    curl \\
    && rm -rf /var/lib/apt/lists/*

# Set Python 3.9 as default
RUN update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.9 1

# Copy requirements first for better caching
COPY ./requirements.txt .
RUN pip3 install --no-cache-dir -r requirements.txt

# Copy source code
COPY ./src/zeta_cuda ./src/zeta_cuda
COPY ./src/common ./src/common

# Create directories for volumes
RUN mkdir -p /app/config /app/locks

# Copy offline check script
COPY ./src/zeta_cuda/scripts/offline_check.sh /app/offline_check.sh
RUN chmod +x /app/offline_check.sh

# Create wrapper script with offline-first pattern
RUN echo '#!/usr/bin/env python3\\n\\
import os, sys\\n\\
sys.path.append("/app")\\n\\
\\n\\
# Implement single connection attempt pattern\\n\\
def check_offline():\\n\\
    """Check if system should run in offline mode"""\\n\\
    if os.path.exists("/app/locks/offline.lock"):\\n\\
        print("Offline lock file detected, running in offline mode")\\n\\
        os.environ["MCP_OFFLINE_MODE"] = "true"\\n\\
        return True\\n\\
    return False\\n\\
\\n\\
# Check offline status before importing rest of system\\n\\
is_offline = check_offline()\\n\\
\\n\\
from src.zeta_cuda.accelerator import main\\n\\
\\n\\
if __name__ == "__main__":\\n\\
    main(offline_mode=is_offline)' > /app/run.py\\n\\

# Set entrypoint with offline check
ENTRYPOINT ["/bin/sh", "-c", "/app/offline_check.sh && python3 /app/run.py"]
EOF

  echo "ZETA component Dockerfiles generated successfully"
}

# Generate container startup script for LLM services
generate_llm_startup_script() {
  local target_dir="$1"
  local scripts_dir="${target_dir}/scripts"
  
  echo "Generating LLM startup script..."
  
  mkdir -p "${scripts_dir}"
  
  cat > "${scripts_dir}/start_llm_services.sh" << EOF
#!/bin/bash
# Start MCP-ZERO LLM services with offline-first resilience

set -e

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

# Check if offline mode is active
check_offline_mode() {
  if [ -f "./.mcp/locks/offline.lock" ]; then
    echo -e "\${YELLOW}Offline lock file detected, starting in offline mode\${NC}"
    export MCP_OFFLINE_MODE=true
    return 0
  fi
  return 1
}

# Start LLM services
start_llm_services() {
  echo -e "\${BLUE}Starting MCP-ZERO LLM services...\${NC}"
  
  # Check if docker-compose is available
  if command -v docker-compose &> /dev/null; then
    # Start LLM services
    docker-compose -f docker/docker-compose.llm.yaml up -d
    
    echo -e "\${GREEN}LLM services started successfully\${NC}"
  else
    echo -e "\${RED}Error: docker-compose not found\${NC}"
    echo -e "\${YELLOW}Please install docker-compose and try again\${NC}"
    exit 1
  fi
}

# Check for GPU support
check_gpu_support() {
  echo -e "\${BLUE}Checking for GPU support...\${NC}"
  
  if command -v nvidia-smi &> /dev/null; then
    echo -e "\${GREEN}NVIDIA GPU detected\${NC}"
    # Check if nvidia-docker runtime is available
    if docker info | grep -q "nvidia"; then
      echo -e "\${GREEN}NVIDIA Docker runtime available\${NC}"
      export GPU_AVAILABLE=true
    else
      echo -e "\${YELLOW}NVIDIA Docker runtime not found, falling back to CPU\${NC}"
      export GPU_AVAILABLE=false
    fi
  else
    echo -e "\${YELLOW}No NVIDIA GPU detected, running in CPU-only mode\${NC}"
    export GPU_AVAILABLE=false
  fi
}

# Main execution
main() {
  echo -e "\${BLUE}MCP-ZERO LLM Services\${NC}"
  echo -e "\${BLUE}====================\${NC}"
  
  # Check offline mode
  check_offline_mode
  
  # Check GPU support
  check_gpu_support
  
  # Start LLM services
  start_llm_services
  
  echo -e "\${GREEN}MCP-ZERO LLM services are now running\${NC}"
  echo -e "\${YELLOW}To stop services:\${NC}"
  echo -e "docker-compose -f docker/docker-compose.llm.yaml down"
}

main "\$@"
EOF

  chmod +x "${scripts_dir}/start_llm_services.sh"
  
  echo "LLM startup script generated successfully at ${scripts_dir}/start_llm_services.sh"
}

# Setup LLM container environment
setup_llm_containers() {
  local config_dir="$1"
  local docker_dir="$2"
  local target_dir="$3"
  local offline_mode="$4"
  
  # Generate LLM configuration
  generate_llm_config "${config_dir}" "${offline_mode}"
  
  # Generate LLM docker configuration
  generate_llm_docker "${docker_dir}" "${offline_mode}"
  
  # Generate LLM startup script
  generate_llm_startup_script "${target_dir}"
  
  echo "LLM container setup completed successfully"
}

# If script is executed directly, show usage
if [[ "${BASH_SOURCE[0]}" == "${0}" ]]; then
  echo "This script is intended to be sourced by burn.sh"
  echo "Usage: source $(basename "$0")"
fi
