"""Enhanced template generator for Modal deployments using Jinja2.

This module provides advanced template generation capabilities that avoid
nested f-string issues by using Jinja2 templating engine.
"""

import os
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Dict, List, Optional

from jinja2 import Environment, FileSystemLoader, select_autoescape
from loguru import logger

from modal_for_noobs.config_loader import config_loader


@dataclass
class RemoteFunctionConfig:
    """Configuration for a Modal remote function."""
    
    name: str
    keep_warm: int = 0
    gpu: str | None = None
    num_gpus: int | None = None
    timeout: int = 300
    memory: int = 8192
    cpu: float = 2.0
    secret: str | None = None
    volume: dict[str, str] | None = None
    schedule: str | None = None  # Cron schedule
    
    def to_dict(self) -> dict[str, Any]:
        """Convert to dictionary for template rendering."""
        return {
            "name": self.name,
            "keep_warm": self.keep_warm,
            "gpu": self.gpu,
            "num_gpus": self.num_gpus,
            "timeout": self.timeout,
            "memory": self.memory,
            "cpu": self.cpu,
            "secret": self.secret,
            "volume": self.volume,
            "schedule": self.schedule,
        }


@dataclass 
class TemplateConfig:
    """Configuration for template generation."""
    
    # Basic settings
    app_name: str
    deployment_mode: str = "minimum"
    description: str = "Modal deployment generated by modal-for-noobs"
    
    # Container settings
    timeout_seconds: int = 3600
    max_containers: int = 10
    min_containers: int = 1
    scaledown_window: int = 1200
    concurrent_inputs: int = 100
    
    # Dependencies
    python_dependencies: list[str] = field(default_factory=list)
    system_dependencies: list[str] = field(default_factory=list)
    requirements_file: Path | None = None
    
    # Infrastructure
    gpu_type: str | None = None
    num_gpus: int = 0
    memory_mb: int = 8192
    cpu: float = 2.0
    
    # Features
    provision_nfs: bool = False
    provision_logging: bool = True
    enable_dashboard: bool = True
    enable_monitoring: bool = True
    
    # Remote functions
    remote_functions: list[RemoteFunctionConfig] = field(default_factory=list)
    
    # Environment
    environment_variables: dict[str, str] = field(default_factory=dict)
    secrets: list[str] = field(default_factory=list)
    
    # Original code
    original_code: str = ""
    
    def to_dict(self) -> dict[str, Any]:
        """Convert to dictionary for template rendering."""
        return {
            "app_name": self.app_name,
            "deployment_mode": self.deployment_mode,
            "description": self.description,
            "timeout_seconds": self.timeout_seconds,
            "max_containers": self.max_containers,
            "min_containers": self.min_containers,
            "scaledown_window": self.scaledown_window,
            "concurrent_inputs": self.concurrent_inputs,
            "python_dependencies": self.python_dependencies,
            "system_dependencies": self.system_dependencies,
            "gpu_type": self.gpu_type,
            "num_gpus": self.num_gpus,
            "memory_mb": self.memory_mb,
            "cpu": self.cpu,
            "provision_nfs": self.provision_nfs,
            "provision_logging": self.provision_logging,
            "enable_dashboard": self.enable_dashboard,
            "enable_monitoring": self.enable_monitoring,
            "remote_functions": [rf.to_dict() for rf in self.remote_functions],
            "environment_variables": self.environment_variables,
            "secrets": self.secrets,
            "original_code": self.original_code,
        }


class TemplateGenerator:
    """Advanced template generator using Jinja2."""
    
    def __init__(self):
        """Initialize template generator."""
        # Set up Jinja2 environment
        template_dir = Path(__file__).parent / "templates" / "jinja2"
        self.env = Environment(
            loader=FileSystemLoader(template_dir),
            autoescape=select_autoescape(['py']),
            trim_blocks=True,
            lstrip_blocks=True,
        )
        
        # Add custom filters
        self.env.filters['classify'] = self._classify_filter
        self.env.filters['format_gpu'] = self._format_gpu_filter
        
    def _classify_filter(self, value: str) -> str:
        """Convert string to class name format."""
        import re
        sanitized = re.sub(r'[^a-zA-Z0-9\s]', '', value)
        title = sanitized.title()
        cleaned = title.strip(' _')
        if cleaned and not cleaned[0].isalpha():
            cleaned = 'Class' + cleaned
        return cleaned
    
    def _format_gpu_filter(self, value: str) -> str:
        """Format GPU type for Modal."""
        gpu_mapping = {
            "any": "gpu.Any()",
            "t4": "gpu.T4()",
            "l4": "gpu.L4()",
            "a10g": "gpu.A10G()",
            "a100": "gpu.A100()",
            "h100": "gpu.H100()",
        }
        return gpu_mapping.get(value.lower(), "gpu.Any()")
    
    def generate_deployment(self, config: TemplateConfig) -> str:
        """Generate a complete Modal deployment script.
        
        Args:
            config: Template configuration
            
        Returns:
            str: Generated deployment script
        """
        # Select template based on deployment mode
        template_name = f"{config.deployment_mode}_deployment.j2"
        
        # Check if specific template exists, otherwise use base template
        try:
            template = self.env.get_template(template_name)
        except Exception:
            logger.info(f"Specific template {template_name} not found, using base template")
            template = self.env.get_template("base_deployment.j2")
        
        # Add image configuration
        image_config = self._generate_image_config(config)
        
        # Load dashboard module content
        dashboard_module_content = self._load_dashboard_module()
        
        # Render template
        rendered = template.render(
            **config.to_dict(),
            image_config=image_config,
            dashboard_module=dashboard_module_content,
            has_gpu=bool(config.gpu_type),
            has_remote_functions=bool(config.remote_functions),
            has_secrets=bool(config.secrets),
            has_env_vars=bool(config.environment_variables),
        )
        
        return rendered
    
    def _generate_image_config(self, config: TemplateConfig) -> str:
        """Generate Modal image configuration."""
        # Load package configuration
        packages = config.python_dependencies
        if not packages:
            package_config = config_loader.load_base_packages()
            packages = package_config.get(config.deployment_mode, [])
        
        # Add packages from requirements file if provided
        if config.requirements_file and config.requirements_file.exists():
            try:
                req_content = config.requirements_file.read_text().strip()
                if req_content:
                    req_packages = [line.strip() for line in req_content.split('\n') if line.strip() and not line.startswith('#')]
                    packages.extend(req_packages)
            except Exception as e:
                logger.warning(f"Failed to read requirements file: {e}")
        
        # Build image configuration based on mode
        if config.deployment_mode in ["optimized", "marimo"] or config.gpu_type:
            base_image = 'modal.Image.from_registry("nvidia/cuda:12.1-devel-ubuntu22.04", add_python="3.11")'
        else:
            base_image = 'modal.Image.debian_slim(python_version="3.11")'
        
        # Format packages
        packages_formatted = ',\n        '.join([f'"{pkg}"' for pkg in packages])
        
        # Build image config
        image_parts = [base_image]
        
        if packages:
            image_parts.append(f"    .pip_install(\n        {packages_formatted}\n    )")
        
        if config.system_dependencies:
            sys_deps = ', '.join([f'"{dep}"' for dep in config.system_dependencies])
            image_parts.append(f"    .apt_install({sys_deps})")
        
        # Add GPU verification for GPU-enabled modes
        if config.deployment_mode in ["optimized", "marimo"] or config.gpu_type:
            image_parts.append('''    .run_commands(
        "apt-get update && apt-get install -y git build-essential",
        "nvidia-smi",
        "python -c 'import torch; print(f\\"PyTorch {torch.__version__} - CUDA: {torch.cuda.is_available()}\\")'",
    )''')
        
        return "image = (\n    " + "\n".join(image_parts) + "\n)"
    
    def _load_dashboard_module(self) -> str:
        """Load the dashboard module content from the templates directory."""
        dashboard_file = Path(__file__).parent / "templates" / "dashboard.py"
        if not dashboard_file.exists():
            logger.warning("Dashboard module not found, using minimal placeholder")
            return '"""Minimal dashboard placeholder."""\nimport gradio as gr\ndef create_dashboard_interface(app): return app'
        
        return dashboard_file.read_text()
    
    def create_gradio_wrapper(self, config: TemplateConfig) -> str:
        """Create a Gradio-specific wrapper for the deployment."""
        template = self.env.get_template("gradio_wrapper.j2")
        return template.render(**config.to_dict())
    
    def create_fastapi_wrapper(self, config: TemplateConfig) -> str:
        """Create a FastAPI-specific wrapper for the deployment."""
        template = self.env.get_template("fastapi_wrapper.j2")
        return template.render(**config.to_dict())


def generate_from_wizard_input(
    app_name: str,
    deployment_mode: str,
    original_code: str,
    provision_nfs: bool = False,
    provision_logging: bool = True,
    system_dependencies: list[str] | None = None,
    python_dependencies: list[str] | None = None,
    remote_functions: list[dict[str, Any]] | None = None,
    gpu_type: str | None = None,
    secrets: list[str] | None = None,
    environment_variables: dict[str, str] | None = None,
    requirements_file: Path | None = None,
) -> str:
    """Generate deployment from wizard input.
    
    This function provides a convenient interface for the CLI wizard.
    """
    # Create remote function configs
    remote_func_configs = []
    if remote_functions:
        for rf in remote_functions:
            remote_func_configs.append(RemoteFunctionConfig(
                name=rf.get("name", "remote_function"),
                keep_warm=rf.get("keep_warm", 0),
                gpu=rf.get("gpu"),
                num_gpus=rf.get("num_gpus"),
                timeout=rf.get("timeout", 300),
                memory=rf.get("memory", 8192),
                cpu=rf.get("cpu", 2.0),
                secret=rf.get("secret"),
                volume=rf.get("volume"),
                schedule=rf.get("schedule"),
            ))
    
    # Create template config
    config = TemplateConfig(
        app_name=app_name,
        deployment_mode=deployment_mode,
        original_code=original_code,
        provision_nfs=provision_nfs,
        provision_logging=provision_logging,
        system_dependencies=system_dependencies or [],
        python_dependencies=python_dependencies or [],
        remote_functions=remote_func_configs,
        gpu_type=gpu_type,
        secrets=secrets or [],
        environment_variables=environment_variables or {},
        requirements_file=requirements_file,
    )
    
    # Generate deployment using the working old template system
    from modal_for_noobs.templates.deployment import generate_modal_deployment
    
    # Create a fake app file path for the old system
    fake_app_file = Path(f"{app_name}.py")
    
    logger.debug(f"Calling generate_modal_deployment with app_name={app_name}, mode={deployment_mode}")
    
    try:
        result = generate_modal_deployment(
            app_file=fake_app_file,
            original_code=original_code,
            deployment_mode=deployment_mode,
            timeout_seconds=3600,  # Default timeout
            scaledown_window=1200,  # Default scaledown
        )
        logger.debug("generate_modal_deployment completed successfully")
        return result
    except Exception as e:
        logger.error(f"generate_modal_deployment failed: {e}")
        raise