"""
Auto-generate Dockerfiles from environment snapshots.
Users never need to write or see Dockerfiles - they're generated in temp directories.
"""

import logging
from pathlib import Path
from typing import Optional


class DockerGenerator:
    """Generate production-ready Dockerfiles automatically from environment snapshots."""
    
    def generate_dockerfile(
        self,
        python_version: str,
        requirements: list[str],
        entry_command: str,
        working_directory: Path,
        custom_base_image: str | None = None,
    ) -> tuple[str, str]:
        """
        Generate Dockerfile and requirements.txt content.
        
        Args:
            python_version: Python version (e.g., "3.10")
            requirements: List of pip requirements
            entry_command: Command to run (e.g., "python train.py")
            working_directory: Project directory path
            custom_base_image: Optional pre-built base image (for air-gapped environments)
        
        Returns:
            Tuple of (dockerfile_content, requirements_content)
        """
        # Generate requirements.txt content
        requirements_content = self._generate_requirements_txt(requirements)
        
        # Generate Dockerfile content
        dockerfile_content = self._generate_dockerfile_content(
            python_version=python_version,
            entry_command=entry_command,
            custom_base_image=custom_base_image,
        )
        
        logging.info("Generated Dockerfile automatically")
        return dockerfile_content, requirements_content
    
    def _generate_requirements_txt(self, requirements: list[str]) -> str:
        """
        Generate requirements.txt content from package list.
        
        Args:
            requirements: List of package specifications
        
        Returns:
            requirements.txt file content
        """
        if not requirements:
            return "# No additional requirements detected\n"
        
        # Clean and deduplicate
        requirements = list(dict.fromkeys(requirements))  # Remove duplicates, preserve order
        
        content = "# Auto-generated by Co-DataScientist\n"
        content += "# From your current Python environment\n\n"
        content += "\n".join(requirements) + "\n"
        
        return content
    
    def _generate_dockerfile_content(
        self,
        python_version: str,
        entry_command: str,
        custom_base_image: str | None = None,
    ) -> str:
        """
        Generate Dockerfile content optimized for ML workloads.
        
        Args:
            python_version: Python version (e.g., "3.10")
            entry_command: Command to run
            custom_base_image: Optional pre-built base image with all dependencies
                              (for air-gapped/offline environments)
        
        Returns:
            Dockerfile content
        """
        # Parse entry command into proper CMD format
        cmd_parts = self._parse_entry_command(entry_command)
        
        # If custom base image provided, assume it's complete (air-gap mode)
        # Generate minimal Dockerfile that only copies code
        # Dependencies come from either:
        #   - Baked mode: Installed in custom_base_image
        #   - Venv mode: Mounted at runtime via airgap_venv config
        if custom_base_image:
            dockerfile = f"""# Auto-generated by Co-DataScientist (Air-Gap Mode)
# Using pre-built base image
# Dependencies: from base image OR mounted venv (see config)
FROM {custom_base_image}

# Set working directory
WORKDIR /app

# Copy all project files
COPY . .

# Set environment variable for data path (can be overridden at runtime)
ENV INPUT_URI=/data

# Run the entry command
CMD {cmd_parts}
"""
            return dockerfile
        
        # Standard mode: full Dockerfile with apt-get and pip
        dockerfile = f"""# Auto-generated by Co-DataScientist
# Optimized for machine learning workloads
FROM python:{python_version}-slim

# Install system dependencies for common ML packages
# (pandas, numpy, scipy often need build tools)
RUN apt-get update && apt-get install -y \\
    build-essential \\
    && rm -rf /var/lib/apt/lists/*

# Set working directory
WORKDIR /app

# Copy and install requirements first (Docker layer caching)
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

# Copy all project files
COPY . .

# Set environment variable for data path (can be overridden at runtime)
ENV INPUT_URI=/data

# Run the entry command
CMD {cmd_parts}
"""
        
        return dockerfile
    
    def _parse_entry_command(self, entry_command: str) -> str:
        """
        Parse entry command into Docker CMD format.
        
        Args:
            entry_command: Command string (e.g., "python train.py")
        
        Returns:
            JSON array format for Docker CMD (e.g., '["python", "train.py"]')
        """
        import shlex
        import json
        
        # Split command respecting quotes
        parts = shlex.split(entry_command)
        
        # Return as proper JSON array format (with double quotes)
        return json.dumps(parts)
    
    def _generate_dockerignore(self) -> str:
        """
        Generate .dockerignore content to exclude output folders from Docker images.
        
        This is a safety mechanism (defense in depth) - output folders should already
        be excluded during copy_to_tmp, but this prevents them from accidentally
        being included in Docker images if they somehow make it to the build context.
        
        Returns:
            .dockerignore file content
        """
        return """# Auto-generated by Co-DataScientist
# Exclude all output/results folders from Docker images

# Main results folder (all outputs consolidated here)
results/

# Legacy output folders (backward compatibility)
co_datascientist_runs/
co_datascientist_checkpoints/
current_runs/
co_datascientist_output/

# Cache and temp files
__pycache__/
*.pyc
*.pyo
*.pyd
.pytest_cache/
.qa_cache/
*.log

# Git and IDE
.git/
.gitignore
.vscode/
.idea/

# Documentation
*.md
README*
"""
    
    def generate_hypothesis_dockerfile(
        self,
        base_image_tag: str,
        entry_command: str,
    ) -> str:
        """
        Generate a lightweight Dockerfile for hypothesis testing.
        This builds on top of the base image (which already has all dependencies).
        
        SPEED OPTIMIZATION: This Dockerfile only copies code changes, 
        resulting in ~5-10 second builds instead of 2-20 minutes!
        
        Args:
            base_image_tag: The base image with all dependencies installed
            entry_command: Command to run (e.g., "python train.py")
        
        Returns:
            Lightweight Dockerfile content
        """
        cmd_parts = self._parse_entry_command(entry_command)
        
        dockerfile = f"""# Auto-generated by Co-DataScientist (Fast Hypothesis Mode)
# This builds on the base image with pre-installed dependencies
FROM {base_image_tag}

# Set working directory (should already exist from base)
WORKDIR /app

# Copy ONLY the code files (dependencies already installed in base)
COPY . .

# Set environment variable for data path (can be overridden at runtime)
ENV INPUT_URI=/data

# Run the entry command
CMD {cmd_parts}
"""
        
        return dockerfile
    
    def write_docker_files(
        self,
        temp_directory: Path,
        dockerfile_content: str,
        requirements_content: str,
    ) -> Path:
        """
        Write Dockerfile, requirements.txt, and .dockerignore to temp directory.
        
        Args:
            temp_directory: Directory to write files to
            dockerfile_content: Content of Dockerfile
            requirements_content: Content of requirements.txt
        
        Returns:
            Path to the generated Dockerfile
        """
        # Write Dockerfile
        dockerfile_path = temp_directory / "Dockerfile"
        dockerfile_path.write_text(dockerfile_content)
        logging.info(f"Wrote Dockerfile to {dockerfile_path}")
        
        # Write requirements.txt
        requirements_path = temp_directory / "requirements.txt"
        requirements_path.write_text(requirements_content)
        logging.info(f"Wrote requirements.txt to {requirements_path}")
        
        # Write .dockerignore for safety (defense in depth)
        dockerignore_content = self._generate_dockerignore()
        dockerignore_path = temp_directory / ".dockerignore"
        dockerignore_path.write_text(dockerignore_content)
        logging.debug(f"Wrote .dockerignore to {dockerignore_path}")
        
        return dockerfile_path
    
    def write_hypothesis_dockerfile(
        self,
        temp_directory: Path,
        dockerfile_content: str,
    ) -> Path:
        """
        Write a lightweight hypothesis Dockerfile and .dockerignore.
        
        Args:
            temp_directory: Directory to write Dockerfile to
            dockerfile_content: Content of the lightweight Dockerfile
        
        Returns:
            Path to the generated Dockerfile
        """
        dockerfile_path = temp_directory / "Dockerfile"
        dockerfile_path.write_text(dockerfile_content)
        logging.debug(f"Wrote lightweight hypothesis Dockerfile to {dockerfile_path}")
        
        # Write .dockerignore for safety
        dockerignore_content = self._generate_dockerignore()
        dockerignore_path = temp_directory / ".dockerignore"
        dockerignore_path.write_text(dockerignore_content)
        logging.debug(f"Wrote .dockerignore to {dockerignore_path}")
        
        return dockerfile_path


def generate_docker_setup(
    working_directory: str,
    entry_command: str,
    python_version: Optional[str] = None,
    requirements: Optional[list[str]] = None,
) -> tuple[str, str]:
    """
    Convenience function to generate Docker setup from scratch.
    
    Args:
        working_directory: Project directory
        entry_command: Command to run the project
        python_version: Optional Python version (auto-detected if None)
        requirements: Optional requirements list (auto-detected if None)
    
    Returns:
        Tuple of (dockerfile_content, requirements_content)
    """
    from .environment_detector import EnvironmentDetector
    
    # Auto-detect environment if not provided
    if python_version is None or requirements is None:
        detector = EnvironmentDetector(working_directory)
        snapshot = detector.get_environment_snapshot()
        
        python_version = python_version or snapshot['python_version']
        requirements = requirements or snapshot['requirements']
    
    # Generate Docker files
    generator = DockerGenerator()
    return generator.generate_dockerfile(
        python_version=python_version,
        requirements=requirements,
        entry_command=entry_command,
        working_directory=Path(working_directory),
    )

