#!/usr/bin/env python3

"""
Claude Code Stack - TrueNAS MCP Deployment Script

Deploys the Claude Code + Agents UI stack to TrueNAS SCALE using the MCP protocol.
Requires mcp-server to be running on TrueNAS and accessible via MCP.

Usage:
    python deploy-mcp.py --mcp-url https://mcp.wilddragon.net/mcp --pool tank
"""

import asyncio
import json
import argparse
import os
from pathlib import Path
from typing import Optional
import sys

try:
    import httpx
except ImportError:
    print("Error: httpx is required. Install with: pip install httpx")
    sys.exit(1)

# ANSI Color codes
class Colors:
    GREEN = '\033[0;32m'
    YELLOW = '\033[1;33m'
    RED = '\033[0;31m'
    BLUE = '\033[0;34m'
    RESET = '\033[0m'

def colored(text: str, color: str) -> str:
    """Return colored text for terminal output"""
    return f"{color}{text}{Colors.RESET}"

def log_info(msg: str):
    print(colored(f"[INFO] {msg}", Colors.GREEN))

def log_warn(msg: str):
    print(colored(f"[WARN] {msg}", Colors.YELLOW))

def log_error(msg: str):
    print(colored(f"[ERROR] {msg}", Colors.RED))

def log_step(msg: str):
    print(colored(f"\n>>> {msg}", Colors.BLUE))

# Docker Compose YAML template
DOCKER_COMPOSE_TEMPLATE = """version: '3.8'

services:
  agents-ui:
    image: node:20-alpine
    container_name: claude-agents-ui
    restart: unless-stopped
    ports:
      - "3000:3000"
    environment:
      ANTHROPIC_API_KEY: ${ANTHROPIC_API_KEY}
      NODE_ENV: production
      CLAUDE_DIR: /root/.claude
    volumes:
      - claude-config:/root/.claude
      - workspace:/workspace
    depends_on:
      - claude-code-backend
    networks:
      - claude-stack
    healthcheck:
      test: ["CMD", "wget", "--quiet", "--tries=1", "--spider", "http://localhost:3000"]
      interval: 30s
      timeout: 10s
      retries: 3

  claude-code-backend:
    image: ghcr.io/anthropics/claude-code:latest
    container_name: claude-code-runtime
    restart: unless-stopped
    ports:
      - "5000:5000"
    environment:
      ANTHROPIC_API_KEY: ${ANTHROPIC_API_KEY}
      CLAUDE_MODEL: claude-opus-4-1
      WORKSPACE_DIR: /workspace
    volumes:
      - claude-config:/root/.claude
      - workspace:/workspace
      - ${HOME}/.ssh:/root/.ssh:ro
    networks:
      - claude-stack
    healthcheck:
      test: ["CMD", "test", "-d", "/workspace"]
      interval: 30s
      timeout: 10s
      retries: 3

  postgres:
    image: postgres:16-alpine
    container_name: claude-postgres
    restart: unless-stopped
    ports:
      - "5432:5432"
    environment:
      POSTGRES_USER: claude
      POSTGRES_PASSWORD: ${POSTGRES_PASSWORD}
      POSTGRES_DB: claude_agents
    volumes:
      - postgres-data:/var/lib/postgresql/data
    networks:
      - claude-stack
    healthcheck:
      test: ["CMD-SHELL", "pg_isready -U claude"]
      interval: 10s
      timeout: 5s
      retries: 5

  redis:
    image: redis:7-alpine
    container_name: claude-redis
    restart: unless-stopped
    ports:
      - "6379:6379"
    volumes:
      - redis-data:/data
    networks:
      - claude-stack
    command: redis-server --appendonly yes
    healthcheck:
      test: ["CMD", "redis-cli", "ping"]
      interval: 10s
      timeout: 5s
      retries: 5

volumes:
  claude-config:
    driver: local
  workspace:
    driver: local
  postgres-data:
    driver: local
  redis-data:
    driver: local

networks:
  claude-stack:
    driver: bridge
"""

ENV_TEMPLATE = """# Claude Code Stack Environment Configuration
ANTHROPIC_API_KEY={api_key}
POSTGRES_PASSWORD={postgres_password}
SESSION_SECRET={session_secret}
NODE_ENV=production
CLAUDE_MODEL=claude-opus-4-1
RUN_CLAUDE_CODE=true
LOG_LEVEL=info
"""

class TrueNASMCPDeployer:
    def __init__(self, mcp_url: str, pool_name: str, api_key: str):
        self.mcp_url = mcp_url
        self.pool_name = pool_name
        self.api_key = api_key
        self.project_dir = f"/mnt/{pool_name}/docker/claude-code-stack"
        self.client = httpx.AsyncClient(timeout=30.0)

    async def call_mcp(self, method: str, params: dict) -> dict:
        """Call MCP method via HTTP"""
        payload = {
            "jsonrpc": "2.0",
            "method": method,
            "params": params,
            "id": 1
        }
        
        try:
            response = await self.client.post(
                self.mcp_url,
                json=payload,
                headers={"Content-Type": "application/json"}
            )
            
            if response.status_code != 200:
                log_error(f"MCP call failed: {response.status_code}")
                log_error(f"Response: {response.text}")
                return None
            
            result = response.json()
            
            if "error" in result:
                log_error(f"MCP error: {result['error']}")
                return None
            
            return result.get("result")
        
        except Exception as e:
            log_error(f"MCP connection error: {str(e)}")
            return None

    async def create_directory(self) -> bool:
        """Create project directory on TrueNAS"""
        log_step(f"Creating project directory: {self.project_dir}")
        
        result = await self.call_mcp(
            "truenas:start_process",
            {
                "command": f"mkdir -p {self.project_dir} && chmod 755 {self.project_dir}",
                "shell": True
            }
        )
        
        if result:
            log_info(f"Directory created: {self.project_dir}")
            return True
        return False

    async def write_file(self, filename: str, content: str) -> bool:
        """Write file to TrueNAS via MCP"""
        filepath = f"{self.project_dir}/{filename}"
        
        try:
            # Write file using exec command
            escaped_content = content.replace('"', '\\"').replace('\n', '\\n')
            
            result = await self.call_mcp(
                "truenas:start_process",
                {
                    "command": f'cat > {filepath} << \'EOF\'\n{content}\nEOF',
                    "shell": True
                }
            )
            
            if result:
                log_info(f"Created: {filename}")
                return True
            
        except Exception as e:
            log_error(f"Failed to write {filename}: {str(e)}")
        
        return False

    async def generate_env(self) -> str:
        """Generate .env file with random secrets"""
        import secrets
        import string
        
        postgres_password = secrets.token_urlsafe(16)
        session_secret = secrets.token_urlsafe(32)
        
        return ENV_TEMPLATE.format(
            api_key=self.api_key,
            postgres_password=postgres_password,
            session_secret=session_secret
        )

    async def deploy_docker_compose(self) -> bool:
        """Deploy docker-compose.yml"""
        log_step("Deploying docker-compose.yml")
        return await self.write_file("docker-compose.yml", DOCKER_COMPOSE_TEMPLATE)

    async def deploy_env_file(self) -> bool:
        """Deploy .env file"""
        log_step("Generating and deploying .env file")
        env_content = await self.generate_env()
        return await self.write_file(".env", env_content)

    async def pull_images(self) -> bool:
        """Pull Docker images"""
        log_step("Pulling Docker images")
        
        result = await self.call_mcp(
            "truenas:start_process",
            {
                "command": f"cd {self.project_dir} && docker-compose pull",
                "shell": True
            }
        )
        
        return result is not None

    async def start_services(self) -> bool:
        """Start Docker services"""
        log_step("Starting Docker services")
        
        result = await self.call_mcp(
            "truenas:start_process",
            {
                "command": f"cd {self.project_dir} && docker-compose up -d",
                "shell": True
            }
        )
        
        return result is not None

    async def get_service_status(self) -> Optional[str]:
        """Get service status"""
        result = await self.call_mcp(
            "truenas:start_process",
            {
                "command": f"cd {self.project_dir} && docker-compose ps",
                "shell": True
            }
        )
        
        return result

    async def deploy(self) -> bool:
        """Execute full deployment"""
        try:
            print(colored("="*50, Colors.BLUE))
            print(colored("Claude Code Stack - TrueNAS MCP Deployment", Colors.BLUE))
            print(colored("="*50, Colors.BLUE))
            print()
            
            # Check API key
            if not self.api_key:
                log_error("ANTHROPIC_API_KEY is required")
                return False
            
            log_info(f"MCP Server: {self.mcp_url}")
            log_info(f"Pool: {self.pool_name}")
            log_info(f"Project Directory: {self.project_dir}")
            print()
            
            # Execute deployment steps
            if not await self.create_directory():
                return False
            
            if not await self.deploy_docker_compose():
                return False
            
            if not await self.deploy_env_file():
                return False
            
            if not await self.pull_images():
                log_warn("Failed to pull images, continuing...")
            
            if not await self.start_services():
                return False
            
            # Wait for services
            log_step("Waiting for services to stabilize...")
            await asyncio.sleep(5)
            
            # Get status
            log_step("Checking service status")
            status = await self.get_service_status()
            if status:
                print(status)
            
            return True
        
        finally:
            await self.client.aclose()

async def main():
    parser = argparse.ArgumentParser(
        description="Deploy Claude Code Stack to TrueNAS via MCP"
    )
    
    parser.add_argument(
        "--mcp-url",
        required=True,
        help="MCP server URL (e.g., https://mcp.wilddragon.net/mcp)"
    )
    
    parser.add_argument(
        "--pool",
        default="tank",
        help="TrueNAS pool name (default: tank)"
    )
    
    parser.add_argument(
        "--api-key",
        help="Anthropic API Key (or set ANTHROPIC_API_KEY env var)"
    )
    
    args = parser.parse_args()
    
    # Get API key
    api_key = args.api_key or os.environ.get("ANTHROPIC_API_KEY")
    
    if not api_key:
        log_error("ANTHROPIC_API_KEY not provided")
        log_info("Set it via: export ANTHROPIC_API_KEY=your_key")
        log_info("Or pass: --api-key your_key")
        sys.exit(1)
    
    # Create deployer
    deployer = TrueNASMCPDeployer(
        mcp_url=args.mcp_url,
        pool_name=args.pool,
        api_key=api_key
    )
    
    # Execute deployment
    success = await deployer.deploy()
    
    if success:
        print()
        print(colored("="*50, Colors.GREEN))
        print(colored("Deployment Successful!", Colors.GREEN))
        print(colored("="*50, Colors.GREEN))
        print()
        print("Access Agents UI at: http://your-truenas-ip:3000")
        print(f"Project directory: {deployer.project_dir}")
        sys.exit(0)
    else:
        print()
        print(colored("="*50, Colors.RED))
        print(colored("Deployment Failed!", Colors.RED))
        print(colored("="*50, Colors.RED))
        sys.exit(1)

if __name__ == "__main__":
    asyncio.run(main())