#!/usr/bin/env python3
"""
Generate custom MaxFrame DPE runtime Dockerfile with user-specified packages.

This script creates:
- Dockerfile with user packages
- README.md with build instructions
- requirements.txt for reproducibility
- .dockerignore for optimized builds

Supports Python 3.7-3.12 with interactive selection.
"""

import argparse
import json
import sys
from pathlib import Path
from typing import Dict, List, Optional


# Valid Python versions
VALID_PYTHON_VERSIONS = ["3.7", "3.8", "3.9", "3.10", "3.11", "3.12"]


def parse_package_list(packages: List[str]) -> List[str]:
    """Parse and validate package list."""
    validated = []
    for pkg in packages:
        pkg = pkg.strip()
        if not pkg:
            continue
        # Basic validation
        if any(char in pkg for char in ['<>', '|', '&', ';', '$']):
            raise ValueError(f"Invalid package specification: {pkg}")
        validated.append(pkg)
    return validated


def parse_python_versions(versions_str: Optional[str]) -> List[str]:
    """Parse comma-separated Python versions string."""
    if not versions_str:
        return VALID_PYTHON_VERSIONS  # Default to all

    versions = [v.strip() for v in versions_str.split(',')]

    # Validate
    for v in versions:
        if v not in VALID_PYTHON_VERSIONS:
            raise ValueError(
                f"Invalid Python version: {v}. "
                f"Valid versions: {', '.join(VALID_PYTHON_VERSIONS)}"
            )

    return versions


def prompt_python_versions() -> List[str]:
    """
    Interactively prompt user for Python versions.

    For Claude Code CLI integration, this function should be replaced by
    the AskUserQuestion tool when invoked as a skill.

    When running the script directly, provides command-line prompts.
    """
    print("\n" + "="*70)
    print("Select Python Versions for Custom MaxFrame DPE Runtime")
    print("="*70)
    print(f"\nAvailable Python versions: {', '.join(VALID_PYTHON_VERSIONS)}")
    print("\nOptions:")
    print("  [1] All versions (3.7-3.12) - Maximum compatibility")
    print("  [2] Recent versions (3.10-3.12) - Good balance, smaller size")
    print("  [3] Python 3.11 only - Recommended for production")
    print("  [4] Custom selection")
    print()

    while True:
        choice = input("Enter choice [1-4] (default: 1): ").strip()

        if not choice or choice == "1":
            print("\n✓ Selected: All Python versions (3.7-3.12)")
            return VALID_PYTHON_VERSIONS
        elif choice == "2":
            versions = ["3.10", "3.11", "3.12"]
            print(f"\n✓ Selected: Python {', '.join(versions)}")
            return versions
        elif choice == "3":
            versions = ["3.11"]
            print(f"\n✓ Selected: Python {', '.join(versions)} (Recommended)")
            return versions
        elif choice == "4":
            custom = input("\nEnter versions (comma-separated, e.g., '3.10,3.11'): ").strip()
            try:
                versions = parse_python_versions(custom)
                print(f"\n✓ Selected: Python {', '.join(versions)}")
                return versions
            except ValueError as e:
                print(f"\nError: {e}")
                print("Please try again.\n")
        else:
            print("Invalid choice. Please enter 1, 2, 3, or 4.\n")


def generate_conda_create_commands(python_versions: List[str]) -> str:
    """Generate conda create commands for selected Python versions."""
    commands = []
    for version in python_versions:
        env_name = f"py{version.replace('.', '')}"
        commands.append(f"conda create -y -n {env_name} python={version} --override-channels -c conda-forge")

    return " && \\\n    ".join(commands)


def generate_maxframe_install_loop(python_versions: List[str]) -> str:
    """Generate loop to install MaxFrame SDK in all environments.

    Note: MaxFrame SDK and pyodps are already available in the DPE environment,
    so we don't need to install them in the custom runtime image.
    """
    # Return empty string - MaxFrame SDK and pyodps are provided by DPE environment
    return ""


def generate_package_install_loop(python_versions: List[str], packages: List[str]) -> str:
    """Generate loop to install user packages in all environments."""
    env_names = [f"py{v.replace('.', '')}" for v in python_versions]

    pkg_str = " \\\n            ".join(packages)

    # Check if torch or pytorch is in the package list, if so add pytorch channel
    channels = "-c pytorch \\\n            -c conda-forge" if any('torch' in pkg for pkg in packages) else "-c conda-forge"

    return f"""RUN for env in {' '.join(env_names)}; do \\
        echo "Installing user packages in $env..." && \\
        conda install -n $env --override-channels \\
            {channels} \\
            {pkg_str}; \\
    done && \\
    conda clean -afy"""


def generate_dockerfile(
    packages: List[str],
    python_versions: List[str],
    image_tag: str = "maxframe-dpe-custom:latest",
    additional_env: Optional[Dict[str, str]] = None
) -> str:
    """Generate Dockerfile content."""

    # Generate conda create commands
    conda_create_cmd = generate_conda_create_commands(python_versions)

    # Generate package install loop (MaxFrame SDK is provided by DPE, not needed here)
    package_install = generate_package_install_loop(python_versions, packages)

    # Default environment (use Python 3.11 if available, otherwise latest)
    if "3.11" in python_versions:
        default_env = "py311"
    else:
        default_env = f"py{python_versions[-1].replace('.', '')}"

    # Build environment variables
    env_vars = additional_env or {}

    # CRITICAL: MF_PYTHON_EXECUTABLE must point to the default environment's Python executable
    # This is enforced by MaxFrame to detect Python at runtime
    mf_python_executable = f"/py-runtime/envs/{default_env}/bin/python"

    dockerfile = f"""# Custom MaxFrame DPE Runtime Image
# Generated by alibabacloud-maxframe-runtime-image-creator skill
# Image tag: {image_tag}
# Python versions: {', '.join(python_versions)}

FROM ubuntu:22.04

# Metadata
LABEL maintainer="generated-by-maxframe-skill"
LABEL description="Custom MaxFrame DPE runtime with user packages"
LABEL version="1.0"

# Use Aliyun mirror for Ubuntu packages (faster in China)
RUN echo "deb http://mirrors.aliyun.com/ubuntu/ jammy main restricted universe multiverse\\n\
deb http://mirrors.aliyun.com/ubuntu/ jammy-updates main restricted universe multiverse\\n\
deb http://mirrors.aliyun.com/ubuntu/ jammy-backports main restricted universe multiverse\\n\
deb http://mirrors.aliyun.com/ubuntu/ jammy-security main restricted universe multiverse" \\
    > /etc/apt/sources.list

# Set environment variables
ENV DEBIAN_FRONTEND=noninteractive
ENV MINIFORGE_HOME="/py-runtime"
ENV PATH="${{MINIFORGE_HOME}}/bin:${{PATH}}"
ENV LANG=en_US.UTF-8
ENV LC_ALL=en_US.UTF-8
ENV LC_CTYPE=en_US.UTF-8
ENV TZ="Asia/Shanghai"
ENV TERM=xterm-256color

# Install system dependencies
RUN apt-get update && \\
    DEBIAN_FRONTEND=noninteractive apt-get install -y \\
    wget curl vim ca-certificates locales \\
    build-essential jq dnsutils ffmpeg tzdata && \\
    locale-gen en_US.UTF-8 && \\
    update-locale LANG=en_US.UTF-8 && \\
    rm -rf /var/lib/apt/lists/*

# Install Miniforge (conda-forge distribution)
RUN wget -q https://github.com/conda-forge/miniforge/releases/latest/download/Miniforge3-Linux-x86_64.sh -O miniforge.sh && \\
    bash miniforge.sh -b -p ${{MINIFORGE_HOME}} && \\
    rm -rf miniforge.sh

# Configure conda to use conda-forge only
RUN conda config --remove channels defaults || true && \\
    conda config --add channels conda-forge && \\
    conda config --set channel_priority strict && \\
    conda config --set show_channel_urls true

# Create conda environments for selected Python versions
RUN {conda_create_cmd} && \\
    conda clean -afy

# Install user packages in all Python environments
{package_install}
"""

    # Add environment variables if specified
    if env_vars:
        env_lines = "\n".join([f"ENV {k}={v}" for k, v in env_vars.items()])
        dockerfile += f"\n# Additional environment variables\n{env_lines}\n"

    dockerfile += f"""
# Set default conda environment
# CRITICAL: MF_PYTHON_EXECUTABLE is enforced by MaxFrame to detect Python executable at runtime
ENV CONDA_DEFAULT_ENV={default_env}
ENV MF_PYTHON_EXECUTABLE={mf_python_executable}

# Verify installation
RUN echo "=== Python Environments ===" && \\
    conda env list && \\
    echo "=== Default Python ({default_env}) ===" && \\
    conda run -n {default_env} python --version

# Default command (runs in default environment)
CMD ["conda", "run", "-n", "{default_env}", "python"]
"""

    return dockerfile


def generate_readme(
    packages: List[str],
    python_versions: List[str],
    image_tag: str
) -> str:
    """Generate README with build instructions."""

    package_list = "\n".join([f"- {pkg}" for pkg in packages])
    version_list = ", ".join(python_versions)

    # Use Python 3.11 as default if available, otherwise latest
    if "3.11" in python_versions:
        default_env = "py311"
    else:
        default_env = f"py{python_versions[-1].replace('.', '')}"

    readme = f"""# Custom MaxFrame DPE Runtime

This directory contains a custom Docker image for MaxFrame DPE runtime with additional Python libraries.

## Image Details

- **Base Image**: `ubuntu:22.04` (public)
- **Conda Distribution**: Miniforge (conda-forge)
- **Custom Image Tag**: `{image_tag}`
- **Python Versions**: {version_list}
- **Default Environment**: `{default_env}`
- **Packages Added**: {len(packages)}

## Installed Packages

{package_list}

## Python Environments

This image includes the following Python environments:

"""

    for version in python_versions:
        env_name = f"py{version.replace('.', '')}"
        readme += f"- `{env_name}` - Python {version}\n"

    readme += f"""
**Default environment**: `{default_env}`

## Build Instructions

### 1. Build the Docker Image

```bash
docker build -t {image_tag} .
```

**Build with no cache (recommended for production):**

```bash
docker build --no-cache -t {image_tag} .
```

### 2. Test the Image Locally

**Basic health check:**

```bash
docker run --rm {image_tag} conda run -n {default_env} python --version
```

**Verify installed packages:**

```bash
docker run --rm {image_tag} conda list -n {default_env}
```

**Test specific package imports:**

```bash
docker run --rm {image_tag} conda run -n {default_env} python -c "
import sys
packages = [{', '.join([f'"{pkg.split("=")[0].split(">")[0].split("<")[0]}"' for pkg in packages[:3]])}]
for pkg in packages:
    try:
        mod = __import__(pkg)
        version = getattr(mod, '__version__', 'unknown')
        print(f'{{pkg}}: {{version}}')
    except ImportError as e:
        print(f'{{pkg}}: FAILED - {{e}}')
"
```

**Test in different Python environment:**

```bash
# Replace py310 with any available environment (py37, py38, py39, py310, py311, py312)
docker run --rm {image_tag} conda run -n py39 python --version
```

### 3. Push to Registry

**Tag for your registry:**

```bash
docker tag {image_tag} your-registry.example.com/{image_tag}
```

**Push to registry:**

```bash
docker push your-registry.example.com/{image_tag}
```

### 4. Use in MaxFrame

Configure your MaxFrame session to use the custom image:

```python
import os
from maxframe.session import new_session
from odps import ODPS
import dotenv

dotenv.load_dotenv()

# Create ODPS connection
o = ODPS(
    access_id=os.getenv("ODPS_ACCESS_ID"),
    secret_access_key=os.getenv("ODPS_ACCESS_KEY"),
    project=os.getenv("ODPS_PROJECT"),
    endpoint=os.getenv("ODPS_ENDPOINT"),
)

# Create session with custom image
session = new_session(
    o,
    image="your-registry.example.com/{image_tag}"
)

try:
    # Your MaxFrame operations
    import maxframe.dataframe as md
    df = md.read_odps_table("your_table")
    # ... processing code ...
    result.execute()
finally:
    session.destroy()
```

## Image Size

Check the image size after building:

```bash
docker images {image_tag}
```

**Note:** Including all Python versions (3.7-3.12) results in larger image size (~3-5 GB).
To reduce size, consider including only needed Python versions.

## Troubleshooting

### Issue: Package not found

If a package is not available in conda-forge:

1. Check available channels:
   ```bash
   docker run --rm continuumio/miniconda3 conda search -c conda-forge package-name
   ```

2. Try alternative package name

3. Use pip fallback (modify Dockerfile):
   ```dockerfile
   RUN for env in py37 py38 py39 py310 py311 py312; do \\
           conda run -n $env pip install package-name; \\
       done
   ```

### Issue: Large image size

1. Verify package cache is cleaned:
   ```dockerfile
   RUN conda clean -afy
   ```

2. Check for unnecessary dependencies

3. **Use fewer Python versions** to reduce size significantly

### Issue: Import errors

1. Verify package installation:
   ```bash
   docker run --rm {image_tag} conda list -n {default_env} | grep package-name
   ```

2. Check Python path:
   ```bash
   docker run --rm {image_tag} conda run -n {default_env} python -c "import sys; print(sys.path)"
   ```

3. Try different Python environment

## Security Considerations

- **Do NOT** include credentials in the image
- **Do NOT** include private SSH keys
- **Do NOT** include proprietary data

Pass secrets via environment variables when running containers.

## Maintenance

### Update Packages

To update packages, rebuild the image:

```bash
docker build --no-cache -t {image_tag} .
```

### Update Base Image

When Ubuntu or Miniforge updates:

1. Pull the new base image
2. Rebuild your custom image
3. Test thoroughly
4. Update version tag

## Files in This Directory

- `Dockerfile` - Docker image definition
- `README.md` - This file
- `requirements.txt` - Package list for reference
- `.dockerignore` - Build optimization

## Next Steps

1. Review the generated `Dockerfile`
2. Build and test the image
3. Push to your registry
4. Configure MaxFrame sessions to use this image
5. Monitor performance and resource usage

## Support

For issues with:
- **MaxFrame**: Contact MaxFrame support
- **Conda packages**: Check conda-forge documentation
- **Docker**: Consult Docker documentation

---

Generated by `alibabacloud-maxframe-runtime-image-creator` skill
"""

    return readme


def generate_requirements(packages: List[str], python_versions: List[str]) -> str:
    """Generate requirements.txt for reference."""

    content = f"""# Requirements for custom MaxFrame DPE runtime
# This file is for reference only
# Packages are installed via conda in all Python environments

# Python versions: {', '.join(python_versions)}

# User packages:
"""

    content += "\n".join([f"# {pkg}" for pkg in packages])
    content += "\n\n# Note: These packages are installed via conda in the Dockerfile\n"
    content += "# The conda package manager handles dependencies differently than pip\n"

    return content


def generate_dockerignore() -> str:
    """Generate .dockerignore file."""

    return """# Git
.git
.gitignore

# Documentation
*.md
README.md
docs/

# Python
__pycache__
*.pyc
*.pyo
*.pyd
.Python
*.so
.pytest_cache/
.coverage
htmlcov/

# Virtual environments
venv/
env/
ENV/

# IDE
.vscode/
.idea/
*.swp
*.swo
*~

# OS
.DS_Store
Thumbs.db

# Build artifacts
dist/
build/
*.egg-info/

# Temporary files
*.tmp
*.bak
*.log
tmp/
temp/

# Credentials (should never be in context)
.env
*.pem
*.key
credentials.json
secrets.json
"""


def main():
    parser = argparse.ArgumentParser(
        description="Generate custom MaxFrame DPE runtime Dockerfile"
    )

    parser.add_argument(
        "--packages",
        nargs="+",
        required=True,
        help="List of conda packages to install (e.g., 'transformers' 'pandas=1.5' 'numpy>=1.23')"
    )

    parser.add_argument(
        "--python-versions",
        type=str,
        default=None,
        help="Comma-separated Python versions (e.g., '3.10,3.11,3.12'). Default: all (3.7-3.12)"
    )

    parser.add_argument(
        "--image-tag",
        default="maxframe-dpe-custom:latest",
        help="Docker image tag (default: maxframe-dpe-custom:latest)"
    )

    parser.add_argument(
        "--output-dir",
        default=".",
        help="Output directory for generated files (default: current directory)"
    )

    parser.add_argument(
        "--env-vars",
        help="JSON string of environment variables (e.g., '{\"VAR1\": \"value1\"}')"
    )

    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Print generated files without writing"
    )

    args = parser.parse_args()

    # Parse and validate packages
    try:
        packages = parse_package_list(args.packages)
        if not packages:
            print("Error: No valid packages specified", file=sys.stderr)
            sys.exit(1)
    except ValueError as e:
        print(f"Error: {e}", file=sys.stderr)
        sys.exit(1)

    # Parse Python versions (interactive if not specified)
    try:
        if args.python_versions:
            python_versions = parse_python_versions(args.python_versions)
        elif args.dry_run:
            python_versions = VALID_PYTHON_VERSIONS  # Default for dry-run
        else:
            # Interactive prompt
            python_versions = prompt_python_versions()
    except ValueError as e:
        print(f"Error: {e}", file=sys.stderr)
        sys.exit(1)

    # Parse environment variables
    env_vars = None
    if args.env_vars:
        try:
            env_vars = json.loads(args.env_vars)
        except json.JSONDecodeError as e:
            print(f"Error parsing env-vars: {e}", file=sys.stderr)
            sys.exit(1)

    # Generate files
    dockerfile = generate_dockerfile(
        packages=packages,
        python_versions=python_versions,
        image_tag=args.image_tag,
        additional_env=env_vars
    )

    readme = generate_readme(
        packages=packages,
        python_versions=python_versions,
        image_tag=args.image_tag
    )

    requirements = generate_requirements(packages, python_versions)
    dockerignore = generate_dockerignore()

    if args.dry_run:
        print("=== Dockerfile ===")
        print(dockerfile)
        print("\n=== README.md ===")
        print(readme)
        print("\n=== requirements.txt ===")
        print(requirements)
        print("\n=== .dockerignore ===")
        print(dockerignore)
        return

    # Write files
    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    (output_dir / "Dockerfile").write_text(dockerfile)
    (output_dir / "README.md").write_text(readme)
    (output_dir / "requirements.txt").write_text(requirements)
    (output_dir / ".dockerignore").write_text(dockerignore)

    print(f"\n✓ Generated files in {output_dir.absolute()}:")
    print(f"  - Dockerfile")
    print(f"  - README.md")
    print(f"  - requirements.txt")
    print(f"  - .dockerignore")
    print(f"\nPython versions: {', '.join(python_versions)}")
    print(f"\nPackages to install:")
    for pkg in packages:
        print(f"  - {pkg}")
    print(f"\nImage tag: {args.image_tag}")
    print(f"\nNext steps:")
    print(f"  1. Review the generated Dockerfile")
    print(f"  2. Build: docker build -t {args.image_tag} .")
    print(f"  3. Test: docker run --rm {args.image_tag} conda run -n py310 python --version")
    print(f"  4. See README.md for detailed instructions")


if __name__ == "__main__":
    main()
