#!/usr/bin/env python3

import os
import json
import yaml
import argparse
from typing import Dict, List, Set, Tuple
from pathlib import Path

class DockerConfigGenerator:
    def __init__(self, project_path: str):
        self.project_path = Path(project_path).resolve()
        self.detected_tech = set()
        self.node_version = None
        self.python_version = None
        self.ports = set()
        self.env_vars = set()
        self.dependencies = set()

    def analyze_existing_configs(self) -> None:
        """Analyze existing Docker configurations"""
        # Analyze existing Dockerfile
        dockerfile_path = self.project_path / 'Dockerfile'
        if dockerfile_path.exists():
            with open(dockerfile_path) as f:
                content = f.read()
                # Detect base images and versions
                if 'FROM node:' in content:
                    self.detected_tech.add('nodejs')
                    # Extract node version from FROM statement
                    import re
                    node_match = re.search(r'FROM node:(\d+)', content)
                    if node_match:
                        self.node_version = node_match.group(1)
                if 'python' in content.lower():
                    self.detected_tech.add('python')
                
                # Detect common patterns
                if 'npm ci' in content:
                    self.detected_tech.add('npm')
                if 'pip install' in content:
                    self.detected_tech.add('pip')

        # Analyze existing docker-compose.yml
        compose_path = self.project_path / 'docker-compose.yml'
        if compose_path.exists():
            with open(compose_path) as f:
                try:
                    compose_data = yaml.safe_load(f)
                    services = compose_data.get('services', {})
                    
                    # Analyze each service
                    for service_name, service_config in services.items():
                        # Extract ports
                        ports = service_config.get('ports', [])
                        for port in ports:
                            if isinstance(port, str):
                                port = port.split(':')[0].replace('${', '').split('-')[-1].replace('}', '')
                            self.ports.add(int(port))
                        
                        # Extract environment variables
                        env = service_config.get('environment', [])
                        if isinstance(env, list):
                            for var in env:
                                if isinstance(var, str):
                                    self.env_vars.add(var.split('=')[0])
                        elif isinstance(env, dict):
                            self.env_vars.update(env.keys())
                        
                        # Extract volumes
                        volumes = service_config.get('volumes', [])
                        for volume in volumes:
                            if isinstance(volume, str):
                                if 'node_modules' in volume:
                                    self.detected_tech.add('nodejs')
                                if '.pip' in volume:
                                    self.detected_tech.add('python')
                        
                        # Extract dependencies
                        depends_on = service_config.get('depends_on', [])
                        self.dependencies.update(depends_on)
                    
                    # Extract volume definitions
                    volumes = compose_data.get('volumes', {})
                    for volume_name in volumes:
                        if 'node_modules' in volume_name:
                            self.detected_tech.add('nodejs')
                        if 'pip' in volume_name:
                            self.detected_tech.add('python')
                
                except yaml.YAMLError as e:
                    print(f"Error parsing docker-compose.yml: {e}")

    def detect_technologies(self) -> None:
        """Detect technologies used in the project."""
        # First analyze existing Docker configs
        self.analyze_existing_configs()
        
        # Then analyze project files as before
        if self._find_file('package.json'):
            self.detected_tech.add('nodejs')
            with open(self.project_path / 'package.json') as f:
                package_data = json.load(f)
                deps = {**package_data.get('dependencies', {}), 
                       **package_data.get('devDependencies', {})}
                if 'react' in deps:
                    self.detected_tech.add('react')
                if 'express' in deps:
                    self.detected_tech.add('express')
                if '@tensorflow/tfjs' in deps or 'pytorch' in deps:
                    self.detected_tech.add('ml')
                if not self.node_version:  # Only detect if not already found
                    self._detect_node_version(package_data)

        # Check for Python
        if self._find_file('requirements.txt') or self._find_file('setup.py'):
            self.detected_tech.add('python')
            if self._find_file('requirements.txt'):
                with open(self.project_path / 'requirements.txt') as f:
                    reqs = f.read()
                    if 'torch' in reqs or 'tensorflow' in reqs:
                        self.detected_tech.add('ml')
                    if 'flask' in reqs:
                        self.ports.add(5000)
                    if 'fastapi' in reqs:
                        self.ports.add(8000)

        # Check for AI/ML specific files
        if self._find_file('*.onnx', glob=True) or self._find_file('*.pt', glob=True):
            self.detected_tech.add('ml')

        # Check for environment variables
        env_file = self._find_file('.env')
        if env_file:
            with open(env_file) as f:
                for line in f:
                    if line.strip() and not line.startswith('#'):
                        key = line.split('=')[0].strip()
                        self.env_vars.add(key)

    def _find_file(self, filename: str, glob: bool = False) -> str:
        """Find a file in the project directory."""
        if glob:
            for file in self.project_path.glob(filename):
                return str(file)
        else:
            file_path = self.project_path / filename
            if file_path.exists():
                return str(file_path)
        return None

    def _detect_node_version(self, package_data: Dict) -> None:
        """Detect Node.js version from package.json."""
        engines = package_data.get('engines', {})
        if 'node' in engines:
            version = engines['node'].replace('^', '').replace('~', '').split('.')[0]
            self.node_version = version
        else:
            self.node_version = '20'  # Default to latest LTS

    def generate_dockerfile(self) -> str:
        """Generate Dockerfile content based on detected technologies."""
        dockerfile = [
            "# syntax=docker/dockerfile:1.4",
            f"ARG NODE_VERSION={self.node_version}",
            "",
            "# Base development image",
            "FROM node:${NODE_VERSION}-slim AS base",
            "",
            "# Install Python and basic build dependencies",
            "RUN apt-get update && apt-get install -y \\",
            "    python3 \\",
            "    python3-pip \\",
            "    git \\",
            "    curl \\",
            "    build-essential \\",
            "    procps \\",
            "    && rm -rf /var/lib/apt/lists/*",
            "",
            "# Create cache directories",
            "RUN mkdir -p /root/.npm",
            "RUN mkdir -p /root/.pip",
            "",
            "# Set working directory",
            "WORKDIR /app",
            "",
            "# Development stage",
            "FROM base AS dev",
            "",
            "# Install development tools",
            "RUN apt-get update && apt-get install -y \\",
            "    vim \\",
            "    ssh \\",
            "    && rm -rf /var/lib/apt/lists/*",
            "",
            "# Create a non-root user for development",
            "ARG USERNAME=node",
            "ARG USER_UID=1000",
            "ARG USER_GID=$USER_UID",
            "",
            "# Create the user",
            "RUN groupadd --gid $USER_GID $USERNAME \\",
            "    && useradd --uid $USER_UID --gid $USER_GID -m $USERNAME \\",
            "    && apt-get update \\",
            "    && apt-get install -y sudo \\",
            "    && echo $USERNAME ALL=\\(root\\) NOPASSWD:ALL > /etc/sudoers.d/$USERNAME \\",
            "    && chmod 0440 /etc/sudoers.d/$USERNAME",
            "",
            "# Set npm config",
            "RUN npm config set cache /root/.npm \\",
            "    && npm config set prefer-offline true \\",
            "    && npm config set package-lock true",
            "",
            "# Copy package files",
            "COPY package*.json ./",
            "COPY .npmrc ./",
            "COPY requirements.txt ./",
            "",
            "# Install Node.js dependencies with cache",
            "RUN --mount=type=cache,target=/root/.npm \\",
            "    npm ci",
            "",
            "# Install Python dependencies with cache",
            "RUN --mount=type=cache,target=/root/.cache/pip \\",
            "    pip3 install -r requirements.txt",
            "",
            "# Switch to non-root user",
            "USER $USERNAME",
            "",
            "# Production stage",
            "FROM base AS prod",
            "",
            "# Copy package files",
            "COPY package*.json ./",
            "COPY .npmrc ./",
            "COPY requirements.txt ./",
            "",
            "# Install production dependencies",
            "RUN --mount=type=cache,target=/root/.npm \\",
            "    npm ci --only=production",
            "",
            "# Install Python production dependencies",
            "RUN --mount=type=cache,target=/root/.cache/pip \\",
            "    pip3 install -r requirements.txt",
            "",
            "# Copy application code",
            "COPY . .",
            "",
            "# Build the application",
            "RUN npm run build",
            "",
            "# Production command",
            'CMD ["npm", "start"]'
        ]

        return '\n'.join(dockerfile)

    def generate_compose(self) -> str:
        """Generate docker-compose.yml content."""
        compose_config = {
            'version': '3.8',
            'services': {
                'dev': {
                    'build': {
                        'context': '.',
                        'target': 'dev',
                        'args': {
                            'NODE_VERSION': self.node_version
                        }
                    },
                    'volumes': [
                        '../:/app:cached',
                        'node_modules:/app/node_modules' if 'nodejs' in self.detected_tech else None,
                        'npm-cache:/root/.npm' if 'nodejs' in self.detected_tech else None,
                        'pip-cache:/root/.cache/pip' if 'python' in self.detected_tech else None,
                        '~/.gitconfig:/root/.gitconfig',
                        '~/.ssh:/root/.ssh'
                    ],
                    'environment': list(self.env_vars) if self.env_vars else [
                        'NODE_ENV=development',
                        'PORT=${DEV_PORT_1:-3000}',
                        'BACKEND_PORT=${DEV_PORT_2:-5000}',
                    ],
                    'ports': [f"{port}:{port}" for port in self.ports] or [
                        '${DEV_PORT_1:-3000}:3000',
                        '${DEV_PORT_2:-5000}:5000',
                        '9229:9229'
                    ],
                }
            },
            'volumes': {
                'node_modules': {},
                'npm-cache': {},
                'pip-cache': {},
            }
        }

        # Add special services based on dependencies
        compose_config = self.add_special_services(compose_config)

        # Remove None values
        compose_config['services']['dev']['volumes'] = [v for v in compose_config['services']['dev']['volumes'] if v]

        return yaml.dump(compose_config, default_flow_style=False)

    def generate_devcontainer(self) -> str:
        """Generate devcontainer.json content."""
        config = {
            "name": "Development Environment",
            "dockerComposeFile": [
                "../docker-compose.yml",
                "docker-compose.extend.yml"
            ],
            "service": "app",
            "workspaceFolder": "/app",
            "customizations": {
                "vscode": {
                    "extensions": [
                        "ms-azuretools.vscode-docker",
                        "editorconfig.editorconfig"
                    ]
                }
            },
            "forwardPorts": list(self.ports)
        }

        # Add language-specific extensions
        if 'nodejs' in self.detected_tech:
            config["customizations"]["vscode"]["extensions"].extend([
                "dbaeumer.vscode-eslint",
                "esbenp.prettier-vscode",
                "christian-kohler.npm-intellisense"
            ])

        if 'python' in self.detected_tech:
            config["customizations"]["vscode"]["extensions"].extend([
                "ms-python.python",
                "ms-python.vscode-pylance"
            ])

        return json.dumps(config, indent=2)

    def generate_devcontainer_compose(self) -> str:
        """Generate docker-compose.extend.yml content for devcontainer."""
        config = {
            'version': '3.8',
            'services': {
                'app': {
                    'init': True,
                    'security_opt': ['seccomp:unconfined'],
                    'cap_add': ['SYS_PTRACE'],
                    'command': 'sleep infinity'
                }
            }
        }
        return yaml.dump(config, default_flow_style=False)

    def generate_configs(self) -> None:
        """Generate all configuration files."""
        # Create backup directory
        backup_dir = self.project_path / 'docker_config_backups'
        backup_dir.mkdir(exist_ok=True)
        
        # Backup existing files
        for filename in ['Dockerfile', 'docker-compose.yml']:
            filepath = self.project_path / filename
            if filepath.exists():
                backup_path = backup_dir / f"{filename}.backup"
                import shutil
                shutil.copy2(filepath, backup_path)
                print(f"Backed up {filename} to {backup_path}")

        # Create .devcontainer directory if it doesn't exist
        devcontainer_dir = self.project_path / '.devcontainer'
        devcontainer_dir.mkdir(exist_ok=True)

        # Generate and write new files
        with open(self.project_path / 'Dockerfile', 'w') as f:
            f.write(self.generate_dockerfile())

        with open(self.project_path / 'docker-compose.yml', 'w') as f:
            f.write(self.generate_compose())

        with open(devcontainer_dir / 'devcontainer.json', 'w') as f:
            f.write(self.generate_devcontainer())

        with open(devcontainer_dir / 'docker-compose.extend.yml', 'w') as f:
            f.write(self.generate_devcontainer_compose())

    def add_special_services(self, compose_config: dict) -> dict:
        """Add special services based on detected dependencies."""
        for dependency in self.dependencies:
            if dependency == 'ollama':
                compose_config['services']['ollama'] = {
                    'image': 'ollama/ollama',
                    'volumes': ['ollama-models:/root/.ollama'],
                    'ports': ['${OLLAMA_PORT_1:-11434}:11434']
                }
                compose_config['volumes']['ollama-models'] = {}
                
            elif dependency == 'sdwebui':
                compose_config['services']['sdwebui'] = {
                    'image': 'stable-diffusion-webui/stable-diffusion-webui',
                    'volumes': ['sd-models:/models'],
                    'ports': ['${SDWEBUI_PORT_1:-7860}:7860']
                }
                compose_config['volumes']['sd-models'] = {}
                
        return compose_config

def main():
    parser = argparse.ArgumentParser(description='Generate Docker configurations for a project')
    parser.add_argument('project_path', help='Path to the project directory')
    args = parser.parse_args()

    generator = DockerConfigGenerator(args.project_path)
    print("🔍 Analyzing project...")
    generator.detect_technologies()
    
    print("\n📦 Detected technologies:")
    for tech in generator.detected_tech:
        print(f"  - {tech}")
    
    print("\n🔧 Generating configuration files...")
    generator.generate_configs()
    
    print("\n✅ Generated files:")
    print("  - Dockerfile")
    print("  - docker-compose.yml")
    print("  - .devcontainer/devcontainer.json")
    print("  - .devcontainer/docker-compose.extend.yml")

if __name__ == "__main__":
    main()