"""Script section implementations with clear separation of concerns.

Each section is responsible for generating a specific part of the startup script.
This separation allows for easy testing, extension, and modification.
"""

import textwrap
from abc import ABC, abstractmethod
from dataclasses import dataclass
from typing import Any, Dict, List, Optional, Protocol

from ...core.constants import FCP_LOG_DIR, FCP_STARTUP_LOG, FLOW_LEGACY_LOG_DIR, FLOW_LOG_DIR


@dataclass
class ScriptContext:
    """Context passed to script sections for generation.

    Contains all configuration needed to generate startup script sections.
    Each script section can access this context to determine what to generate
    and how to configure the instance.

    Attributes:
        ports: List of ports to expose through the bastion.
        volumes: List of volume configurations to mount.
        docker_image: Docker image to run, if any.
        docker_command: Command to run in the Docker container.
        user_script: Custom shell script to execute during startup.
        environment: Environment variables to set.
        upload_code: Whether to extract uploaded code archive.
        code_archive: Base64-encoded gzipped tar archive of user code.
    """

    ports: List[int] = None
    volumes: List[Dict[str, Any]] = None
    docker_image: Optional[str] = None
    docker_command: Optional[List[str]] = None
    user_script: Optional[str] = None
    environment: Dict[str, str] = None
    upload_code: bool = False
    code_archive: Optional[str] = None  # Base64 encoded tar.gz

    def __post_init__(self):
        """Initialize default values for mutable attributes."""
        self.ports = self.ports or []
        self.volumes = self.volumes or []
        self.environment = self.environment or {}


class IScriptSection(Protocol):
    """Protocol for script sections."""

    @property
    def name(self) -> str:
        """Section name for identification."""
        ...

    @property
    def priority(self) -> int:
        """Execution priority (lower numbers run first)."""
        ...

    def should_include(self, context: ScriptContext) -> bool:
        """Determine if this section should be included."""
        ...

    def generate(self, context: ScriptContext) -> str:
        """Generate the script section content."""
        ...

    def validate(self, context: ScriptContext) -> List[str]:
        """Validate the context for this section. Return list of errors."""
        ...


class ScriptSection(ABC):
    """Base class for script sections."""

    @property
    @abstractmethod
    def name(self) -> str:
        pass

    @property
    def priority(self) -> int:
        return 50  # Default middle priority

    def should_include(self, context: ScriptContext) -> bool:
        """By default, include if generate returns non-empty content."""
        return bool(self.generate(context).strip())

    @abstractmethod
    def generate(self, context: ScriptContext) -> str:
        pass

    def validate(self, context: ScriptContext) -> List[str]:
        """Default validation (none)."""
        return []


class HeaderSection(ScriptSection):
    """Generate script header with safety checks."""

    @property
    def name(self) -> str:
        return "header"

    @property
    def priority(self) -> int:
        return 10  # Always first

    def should_include(self, context: ScriptContext) -> bool:
        return True  # Always include header

    def generate(self, context: ScriptContext) -> str:
        return textwrap.dedent(
            f"""
            #!/bin/bash
            set -euxo pipefail
            
            # FCP Startup Script with SLURM-like logging
            echo "Starting FCP instance initialization at $(date)"
            
            # Ensure we're running as root
            if [ "$EUID" -ne 0 ]; then
                echo "Error: Script must run as root" >&2
                exit 1
            fi
            
            # Get task ID from environment (FCP sets this)
            TASK_ID="${{FCP_BID_ID:-unknown}}"
            
            # Create log directories
            LOG_DIR="{FLOW_LOG_DIR}"
            mkdir -p "$LOG_DIR"
            mkdir -p {FLOW_LEGACY_LOG_DIR}  # Legacy directory for compatibility
            
            # Create SLURM-like log files
            STDOUT_LOG="$LOG_DIR/$TASK_ID.out"
            STDERR_LOG="$LOG_DIR/$TASK_ID.err"
            COMBINED_LOG="$LOG_DIR/$TASK_ID.log"
            
            # Create symlink for legacy compatibility
            ln -sf "$COMBINED_LOG" {FLOW_LEGACY_LOG_DIR}/startup.log
            
            # Redirect outputs to separate files (like SLURM)
            exec 1> >(tee -a "$STDOUT_LOG" "$COMBINED_LOG")
            exec 2> >(tee -a "$STDERR_LOG" "$COMBINED_LOG" >&2)
            
            # Log startup information
            echo "[$(date)] Task $TASK_ID starting on $(hostname)"
            echo "Logs: stdout=$STDOUT_LOG stderr=$STDERR_LOG combined=$COMBINED_LOG"
            
            # Create FCP-compatible log locations
            echo "Creating FCP log symlinks..."
            mkdir -p {FCP_LOG_DIR}
            ln -sf "$COMBINED_LOG" {FCP_STARTUP_LOG}
        """
        ).strip()


class PortForwardingSection(ScriptSection):
    """Configure port forwarding through FCP bastion."""

    @property
    def name(self) -> str:
        return "port_forwarding"

    @property
    def priority(self) -> int:
        return 20

    def should_include(self, context: ScriptContext) -> bool:
        return bool(context.ports)

    def generate(self, context: ScriptContext) -> str:
        if not context.ports:
            return ""

        # Generate nginx configs
        nginx_configs = []
        for port in context.ports:
            nginx_configs.append(self._generate_nginx_config(port))

        return textwrap.dedent(
            f"""
            # Port forwarding setup
            echo "Configuring port forwarding for ports: {', '.join(map(str, context.ports))}"
            
            # Install nginx for local reverse proxy
            export DEBIAN_FRONTEND=noninteractive
            apt-get update -qq
            apt-get install -y -qq nginx
            
            # Remove default site
            rm -f /etc/nginx/sites-enabled/default
            
            {chr(10).join(nginx_configs)}
            
            # Test and reload nginx
            nginx -t
            systemctl enable nginx
            systemctl restart nginx
            
            {self._generate_foundrypf_service()}
        """
        ).strip()

    def _generate_nginx_config(self, port: int) -> str:
        """Generate nginx configuration for a single port."""
        return textwrap.dedent(
            f"""
            # Configure port {port}
            cat > /etc/nginx/sites-available/port{port} <<'NGINX_EOF'
            server {{
                listen {port};
                server_name _;
                
                location / {{
                    proxy_pass http://127.0.0.1:{port};
                    proxy_http_version 1.1;
                    proxy_set_header Upgrade $http_upgrade;
                    proxy_set_header Connection 'upgrade';
                    proxy_set_header Host $host;
                    proxy_set_header X-Real-IP $remote_addr;
                    proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for;
                    proxy_set_header X-Forwarded-Proto $scheme;
                    proxy_cache_bypass $http_upgrade;
                    proxy_read_timeout 86400;
                }}
            }}
            NGINX_EOF
            ln -sf /etc/nginx/sites-available/port{port} /etc/nginx/sites-enabled/
        """
        ).strip()

    def _generate_foundrypf_service(self) -> str:
        """Generate systemd service for foundrypf."""
        return textwrap.dedent(
            """
            # Create systemd service for foundrypf
            cat > /etc/systemd/system/foundrypf.service <<'SYSTEMD_EOF'
            [Unit]
            Description=Foundry Port Forwarding
            After=network-online.target
            Wants=network-online.target
            
            [Service]
            Type=simple
            ExecStart=/usr/local/bin/foundrypf
            Restart=always
            RestartSec=10
            StandardOutput=journal
            StandardError=journal
            SyslogIdentifier=foundrypf
            
            [Install]
            WantedBy=multi-user.target
            SYSTEMD_EOF
            
            systemctl daemon-reload
            systemctl enable foundrypf
            systemctl start foundrypf
        """
        ).strip()

    def validate(self, context: ScriptContext) -> List[str]:
        errors = []
        for port in context.ports:
            if not (1 <= port <= 65535):
                errors.append(f"Invalid port number: {port}")
        return errors


class VolumeSection(ScriptSection):
    """Handle volume mounting with proper error handling."""

    @property
    def name(self) -> str:
        return "volumes"

    @property
    def priority(self) -> int:
        return 30

    def should_include(self, context: ScriptContext) -> bool:
        return bool(context.volumes)

    def generate(self, context: ScriptContext) -> str:
        if not context.volumes:
            return ""

        mount_commands = []
        for i, volume in enumerate(context.volumes):
            # Check if this is a file share
            if volume.get("interface") == "file":
                mount_commands.append(self._generate_file_mount(i, volume))
            else:
                mount_commands.append(self._generate_block_mount(i, volume))

        return textwrap.dedent(
            f"""
            # Volume mounting
            echo "Mounting {len(context.volumes)} volume(s)"
            
            {chr(10).join(mount_commands)}
            
            # Verify all mounts
            echo "Mounted volumes:"
            mount | grep -E "(^/dev/xvd[f-z]|type nfs)"
        """
        ).strip()

    def _generate_file_mount(self, index: int, volume: Dict[str, Any]) -> str:
        """Mount file share using NFS."""
        mount_path = volume.get("mount_path", f"/data{index}")
        # File shares are exposed as NFS endpoints in FCP
        nfs_endpoint = volume.get("nfs_endpoint", f"fileshare-{volume['volume_id']}.fcp.internal")

        return textwrap.dedent(
            f"""
            # File share {index}: {nfs_endpoint} -> {mount_path}
            echo "Mounting file share to {mount_path}"
            
            # Install NFS client if not present
            if ! command -v mount.nfs >/dev/null; then
                apt-get update -qq && apt-get install -y -qq nfs-common
            fi
            
            # Create mount point
            mkdir -p {mount_path}
            
            # Mount with standard NFS options
            mount -t nfs4 -o nfsvers=4.1,rsize=1048576,wsize=1048576,hard,timeo=600,retrans=2 \\
                {nfs_endpoint}:/ {mount_path}
            
            # Add to fstab for persistence
            if ! grep -q "{nfs_endpoint}" /etc/fstab; then
                echo "{nfs_endpoint}:/ {mount_path} nfs4 nfsvers=4.1,rsize=1048576,wsize=1048576,hard,timeo=600,retrans=2,_netdev 0 0" >> /etc/fstab
            fi
            
            # Verify mount
            if mountpoint -q {mount_path}; then
                echo "File share mounted successfully at {mount_path}"
                chmod 755 {mount_path}
            else
                echo "ERROR: Failed to mount file share at {mount_path}" >&2
                exit 1
            fi
        """
        ).strip()

    def _generate_block_mount(self, index: int, volume: Dict[str, Any]) -> str:
        """Mount block storage device."""
        mount_path = volume.get("mount_path", f"/data{index}")
        device = f"/dev/xvd{chr(102 + index)}"  # f, g, h, ...

        return textwrap.dedent(
            f"""
            # Volume {index}: {device} -> {mount_path}
            echo "Processing volume {index}: {device} -> {mount_path}"
            
            # Create mount point
            mkdir -p {mount_path}
            
            # Wait for device (max 60 seconds)
            DEVICE={device}
            TIMEOUT=60
            ELAPSED=0
            while [ ! -b "$DEVICE" ] && [ $ELAPSED -lt $TIMEOUT ]; do
                echo "Waiting for device $DEVICE... ($ELAPSED/$TIMEOUT)"
                sleep 5
                ELAPSED=$((ELAPSED + 5))
            done
            
            if [ ! -b "$DEVICE" ]; then
                echo "ERROR: Device $DEVICE not found after $TIMEOUT seconds" >&2
                exit 1
            fi
            
            # Check if already formatted
            if ! blkid "$DEVICE" >/dev/null 2>&1; then
                echo "Formatting new volume $DEVICE"
                mkfs.ext4 -F "$DEVICE"
            fi
            
            # Mount the volume
            mount "$DEVICE" {mount_path}
            
            # Add to fstab for persistence
            if ! grep -q "$DEVICE" /etc/fstab; then
                echo "$DEVICE {mount_path} ext4 defaults,nofail,x-systemd.device-timeout=10 0 2" >> /etc/fstab
            fi
            
            # Set permissions
            chmod 755 {mount_path}
        """
        ).strip()

    def validate(self, context: ScriptContext) -> List[str]:
        errors = []
        if len(context.volumes) > 20:  # AWS limit
            errors.append(f"Too many volumes: {len(context.volumes)} (max 20)")

        for i, volume in enumerate(context.volumes):
            mount_path = volume.get("mount_path", f"/data{i}")
            if not mount_path.startswith("/"):
                errors.append(f"Volume {i}: mount_path must be absolute: {mount_path}")

        return errors


class S3Section(ScriptSection):
    """Handle S3 mounting via s3fs."""

    @property
    def name(self) -> str:
        return "s3_mounts"

    @property
    def priority(self) -> int:
        return 35  # After volumes, before docker

    def should_include(self, context: ScriptContext) -> bool:
        # Check if we have S3 mount environment variables
        return any(
            k.startswith("S3_MOUNT_") and k.endswith("_BUCKET") for k in context.environment.keys()
        )

    def generate(self, context: ScriptContext) -> str:
        # Count S3 mounts
        mount_count = int(context.environment.get("S3_MOUNTS_COUNT", "0"))
        if mount_count == 0:
            return ""

        mount_commands = []
        for i in range(mount_count):
            mount_key = f"S3_MOUNT_{i}"
            bucket = context.environment.get(f"{mount_key}_BUCKET")
            path = context.environment.get(f"{mount_key}_PATH", "")
            target = context.environment.get(f"{mount_key}_TARGET")

            if bucket and target:
                mount_commands.append(self._generate_s3_mount(bucket, path, target))

        return textwrap.dedent(
            f"""
            # S3 mounting via s3fs
            echo "Setting up S3 mounts"
            
            # Install s3fs
            export DEBIAN_FRONTEND=noninteractive
            apt-get update -qq
            apt-get install -y -qq s3fs
            
            # Create credential file for s3fs
            echo "${{AWS_ACCESS_KEY_ID}}:${{AWS_SECRET_ACCESS_KEY}}" > /tmp/s3fs_passwd
            chmod 600 /tmp/s3fs_passwd
            
            {chr(10).join(mount_commands)}
            
            # Verify S3 mounts
            echo "S3 mounts configured:"
            mount | grep s3fs
        """
        ).strip()

    def _generate_s3_mount(self, bucket: str, path: str, target: str) -> str:
        """Generate s3fs mount command for a single S3 location."""
        s3_path = f"{bucket}:/{path}" if path else bucket

        return textwrap.dedent(
            f"""
            # Mount S3: {bucket}/{path} -> {target}
            echo "Mounting S3 bucket {bucket} to {target}"
            
            # Create mount point
            mkdir -p {target}
            
            # Mount with s3fs
            s3fs {s3_path} {target} \\
                -o passwd_file=/tmp/s3fs_passwd \\
                -o allow_other \\
                -o use_cache=/tmp/s3fs_cache \\
                -o ro \\
                -o retries=5 \\
                -o connect_timeout=10 \\
                -o readwrite_timeout=30
            
            # Verify mount
            if mountpoint -q {target}; then
                echo "Successfully mounted S3 to {target}"
            else
                echo "ERROR: Failed to mount S3 to {target}" >&2
                exit 1
            fi
        """
        ).strip()

    def validate(self, context: ScriptContext) -> List[str]:
        errors = []

        # Check for AWS credentials
        if not context.environment.get("AWS_ACCESS_KEY_ID"):
            errors.append("AWS_ACCESS_KEY_ID not set in environment")
        if not context.environment.get("AWS_SECRET_ACCESS_KEY"):
            errors.append("AWS_SECRET_ACCESS_KEY not set in environment")

        return errors


class CodeUploadSection(ScriptSection):
    """Extract and prepare uploaded local code.

    This section handles extracting user code that was packaged and embedded
    in the startup script. The code archive is base64-encoded and gzipped
    to minimize startup script size while preserving directory structure.

    The extracted code is placed in /workspace, which becomes the working
    directory for Docker containers when upload_code=True.
    """

    @property
    def name(self) -> str:
        """Return section identifier.

        Returns:
            str: Section name used for identification and debugging.
        """
        return "code_upload"

    @property
    def priority(self) -> int:
        """Return execution priority.

        Returns:
            int: Priority value (35) - runs after volumes, before Docker.
        """
        return 35  # After volumes, before Docker

    def should_include(self, context: ScriptContext) -> bool:
        """Determine if code upload section should be included.

        Args:
            context: Script generation context containing upload settings.

        Returns:
            bool: True if upload_code is enabled and archive is provided.
        """
        return context.upload_code and context.code_archive is not None

    def generate(self, context: ScriptContext) -> str:
        """Generate shell script to extract uploaded code.

        Args:
            context: Script generation context with code archive.

        Returns:
            str: Shell script that extracts code to /workspace.
        """
        if not context.upload_code or not context.code_archive:
            return ""

        return textwrap.dedent(
            f"""
            # Extract uploaded code
            echo "Extracting uploaded code to /workspace..."
            mkdir -p /workspace
            cd /workspace
            
            # Decode and extract the code archive
            echo "{context.code_archive}" | base64 -d | tar -xzf -
            
            # List extracted files
            echo "Extracted files:"
            ls -la
            
            # Set proper permissions
            chmod -R 755 /workspace
        """
        ).strip()

    def validate(self, context: ScriptContext) -> List[str]:
        """Validate context for code upload requirements.

        Args:
            context: Script generation context to validate.

        Returns:
            List[str]: List of validation errors, empty if valid.
        """
        errors = []
        if context.upload_code and not context.code_archive:
            errors.append("upload_code is True but no code archive provided")
        return errors


class DockerSection(ScriptSection):
    """Docker installation and container management."""

    @property
    def name(self) -> str:
        return "docker"

    @property
    def priority(self) -> int:
        return 40

    def should_include(self, context: ScriptContext) -> bool:
        return bool(context.docker_image)

    def generate(self, context: ScriptContext) -> str:
        if not context.docker_image:
            return ""

        docker_run_cmd = self._build_docker_run_command(context)

        return textwrap.dedent(
            f"""
            # Docker setup
            echo "Setting up Docker and running {context.docker_image}"
            
            # Install Docker
            if ! command -v docker >/dev/null 2>&1; then
                echo "Installing Docker..."
                curl -fsSL https://get.docker.com -o /tmp/get-docker.sh
                sh /tmp/get-docker.sh
                rm /tmp/get-docker.sh
            fi
            
            # Enable and start Docker
            systemctl enable docker
            systemctl start docker
            
            # Wait for Docker to be ready
            timeout 30 sh -c 'until docker info >/dev/null 2>&1; do sleep 1; done'
            
            # Pull the image if not already present
            echo "Checking Docker image: {context.docker_image}"
            docker image inspect {context.docker_image} >/dev/null 2>&1 || docker pull {context.docker_image}
            
            # Run the container
            echo "Starting container..."
            {docker_run_cmd}
            
            # Verify container is running
            sleep 5
            docker ps
            docker logs main --tail 50
        """
        ).strip()

    def _build_docker_run_command(self, context: ScriptContext) -> str:
        """Build the docker run command with all options."""
        cmd_parts = [
            "docker run",
            "-d",
            "--restart=unless-stopped",
            "--name=main",
            "--log-driver=json-file",
            "--log-opt max-size=100m",
            "--log-opt max-file=3",
        ]

        # Add port mappings
        for port in context.ports:
            cmd_parts.append(f"-p {port}:{port}")

        # Add volume mounts (excluding Docker's own storage)
        for i, volume in enumerate(context.volumes):
            mount_path = volume.get("mount_path", f"/data{i}")
            # Don't mount /var/lib/docker inside the container
            if mount_path != "/var/lib/docker":
                cmd_parts.append(f"-v {mount_path}:{mount_path}")

        # Add working directory if code was uploaded
        if context.upload_code:
            cmd_parts.append("-w /workspace")
            # Also mount /workspace from host to container
            cmd_parts.append("-v /workspace:/workspace")

        # Add environment variables
        for key, value in context.environment.items():
            cmd_parts.append(f'-e {key}="{value}"')

        # Add image
        cmd_parts.append(context.docker_image)

        # Add command if specified
        if context.docker_command:
            # docker_command is now a list from TaskConfig validation
            # Need to properly escape command arguments
            import shlex

            for arg in context.docker_command:
                cmd_parts.append(shlex.quote(arg))

        return " \\\n    ".join(cmd_parts)

    def validate(self, context: ScriptContext) -> List[str]:
        errors = []
        # Only warn about images that look like they need a registry
        if (
            context.docker_image
            and "/" not in context.docker_image
            and ":" not in context.docker_image
        ):
            # Skip validation for common official images
            official_images = {
                "ubuntu",
                "debian",
                "alpine",
                "centos",
                "fedora",
                "nginx",
                "redis",
                "postgres",
                "mysql",
                "python",
                "node",
                "golang",
            }
            image_name = context.docker_image.split(":")[0]
            if image_name not in official_images:
                errors.append(
                    f"Docker image should include registry/namespace: {context.docker_image}"
                )
        return errors


class UserScriptSection(ScriptSection):
    """Execute user-provided startup script."""

    @property
    def name(self) -> str:
        return "user_script"

    @property
    def priority(self) -> int:
        return 90  # Run near the end

    def should_include(self, context: ScriptContext) -> bool:
        return bool(context.user_script and context.user_script.strip())

    def generate(self, context: ScriptContext) -> str:
        if not context.user_script or not context.user_script.strip():
            return ""

        return textwrap.dedent(
            f"""
            # User startup script
            echo "Executing user startup script"
            
            cat > /tmp/user_startup.sh <<'USER_SCRIPT_EOF'
            {context.user_script}
            USER_SCRIPT_EOF
            
            chmod +x /tmp/user_startup.sh
            /tmp/user_startup.sh
            
            echo "User startup script completed"
        """
        ).strip()


class CompletionSection(ScriptSection):
    """Mark startup as complete and perform final tasks."""

    @property
    def name(self) -> str:
        return "completion"

    @property
    def priority(self) -> int:
        return 100  # Always last

    def should_include(self, context: ScriptContext) -> bool:
        return True

    def generate(self, context: ScriptContext) -> str:
        return textwrap.dedent(
            """
            # Startup complete
            echo "FCP startup script completed successfully at $(date)"
            
            # Create completion marker
            touch /var/run/fcp-startup-complete
            
            # Log system info
            echo "System information:"
            uname -a
            df -h
            free -h
            
            echo "Startup log saved to /var/log/fcp/startup.log"
        """
        ).strip()
