"""
HmsJython - Jython Script Generation and Execution for HEC-HMS

This module provides static methods for generating Jython scripts and executing
them via direct Java invocation. This approach provides:
- Support for HMS 3.3+ (32-bit) and HMS 4.4.1+ (64-bit)
- Reliable execution across all supported versions (fixes batch file bugs)
- Configurable JVM memory settings for large models
- Automatic version detection and appropriate execution path

Supported HEC-HMS Versions:
    - HMS 4.4.1+: Full support (64-bit, Python 3 compatible scripts)
    - HMS 3.3-3.5: Full support (32-bit, requires python2_compatible=True)
    - HMS 4.0-4.3: Not supported (legacy classpath structure)
    - HMS 3.0-3.2: Untested

HEC-HMS Jython API Reference:
    from hms.model import JythonHms
    JythonHms.OpenProject(name, path)
    JythonHms.OpenBasinModel(name)
    JythonHms.Compute(runName)
    JythonHms.SaveProject()  # or SaveAllProjectComponents() for HMS 3.x
    JythonHms.Exit(code)

Example (HMS 4.x):
    script = HmsJython.generate_compute_script(project_path, run_name)
    success, stdout, stderr = HmsJython.execute_script(script, hms_exe_path)

Example (HMS 3.x):
    script = HmsJython.generate_compute_script(
        project_path, run_name, python2_compatible=True
    )
    success, stdout, stderr = HmsJython.execute_script(script, hms_3x_path)

All methods are static and designed to be used without instantiation.
"""

import os
import subprocess
import tempfile
from pathlib import Path
from typing import Dict, List, Optional, Union, Tuple, Any
from datetime import datetime

from .LoggingConfig import get_logger
from .Decorators import log_call

logger = get_logger(__name__)


class HmsJython:
    """
    Jython script generation and execution for HEC-HMS automation.

    Uses direct Java invocation for reliable execution across all HMS 4.x versions
    with configurable memory settings for large models.

    All methods are static - no instantiation required.

    Example:
        >>> from hms_commander import HmsJython
        >>> script = HmsJython.generate_compute_script(
        ...     project_path="C:/Projects/MyProject",
        ...     run_name="Run 1"
        ... )
        >>> # Execute with 8GB heap for large model
        >>> success, stdout, stderr = HmsJython.execute_script(
        ...     script, hms_exe_path, max_memory="8G"
        ... )
    """

    # Default script header
    SCRIPT_HEADER = '''"""
HEC-HMS Jython Script
Generated by hms-commander
Generated at: {timestamp}
"""

from hms.model import JythonHms
import sys

'''

    # Script footer with proper exit
    SCRIPT_FOOTER = '''
# Exit cleanly
JythonHms.Exit(0)
'''

    # Default JVM memory settings
    DEFAULT_MAX_MEMORY = "4G"
    DEFAULT_INITIAL_MEMORY = "128M"
    # 32-bit JVM memory limits (can't exceed ~1.5GB)
    MAX_MEMORY_32BIT = "1280M"
    INITIAL_MEMORY_32BIT = "64M"

    # -------------------------------------------------------------------------
    # Version Detection and Validation
    # -------------------------------------------------------------------------

    @staticmethod
    def _get_hms_version(hms_path: Path) -> Tuple[int, ...]:
        """
        Extract HMS version from installation path.

        Args:
            hms_path: Path to HMS installation directory

        Returns:
            Tuple of version components (e.g., (4, 11) or (4, 7, 1))
        """
        version_str = hms_path.name  # e.g., "4.11" or "4.7.1"
        try:
            return tuple(int(p) for p in version_str.split('.'))
        except ValueError:
            # Try to read from config file or jar manifest
            logger.warning(f"Could not parse version from path: {hms_path}")
            return (0,)  # Unknown version

    @staticmethod
    def _format_version(version: Tuple[int, ...]) -> str:
        """Format version tuple as string."""
        return '.'.join(map(str, version))

    @staticmethod
    def _is_32bit_java(java_exe: Path) -> bool:
        """
        Detect if Java executable is 32-bit.

        Args:
            java_exe: Path to java.exe

        Returns:
            True if 32-bit Java, False if 64-bit
        """
        # Heuristic: If in "Program Files (x86)", it's 32-bit
        # Also check for "java" folder (old HMS) vs "jre" folder (new HMS)
        path_str = str(java_exe).lower()
        if "program files (x86)" in path_str:
            return True
        # Older HMS versions (4.0-4.3) use "java" folder and are 32-bit
        if "\\java\\bin\\" in path_str or "/java/bin/" in path_str:
            return True
        return False

    @staticmethod
    def _check_version_supported(version: Tuple[int, ...]) -> None:
        """
        Raise error if HMS version doesn't support Jython scripting.

        Args:
            version: Version tuple (e.g., (4, 11) or (3, 3))

        Raises:
            RuntimeError: If HMS version < 3.3 or 4.0-4.3 (unsupported legacy structure)
        """
        # HMS 3.3+ supports Jython scripting (discovered via decompilation)
        if version[0] == 3:
            if len(version) < 2 or version[1] < 3:
                version_str = HmsJython._format_version(version)
                raise RuntimeError(
                    f"HMS {version_str} may not support command-line Jython scripting. "
                    f"hms-commander supports HMS 3.3+ and HMS 4.4.1+. "
                    f"For older versions, use the HEC-HMS GUI."
                )
            # HMS 3.3+ is supported
            return

        if version[0] < 3:
            version_str = HmsJython._format_version(version)
            raise RuntimeError(
                f"HMS {version_str} does not support command-line Jython scripting. "
                f"Jython scripting requires HEC-HMS 3.3 or later."
            )

        # HMS 4.0-4.3 have a different internal structure that isn't currently supported
        # These versions use explicit jar lists instead of wildcards, different native lib paths, etc.
        if version[0] == 4 and (len(version) == 1 or version[1] < 4):
            version_str = HmsJython._format_version(version)
            raise RuntimeError(
                f"HMS {version_str} uses a legacy classpath structure not currently supported. "
                f"hms-commander supports HMS 3.3+ and HMS 4.4.1+. "
                f"For HMS 4.0-4.3, use the HEC-HMS GUI or the batch file directly."
            )

    # -------------------------------------------------------------------------
    # HMS 3.x Support
    # -------------------------------------------------------------------------

    @staticmethod
    def _is_hms_3x(version: Tuple[int, ...]) -> bool:
        """Check if this is an HMS 3.x version."""
        return version[0] == 3

    @staticmethod
    def _build_classpath_3x(hms_install_path: Path) -> str:
        """
        Build explicit classpath for HMS 3.x (no wildcard support in older Java).

        HMS 3.x classpath order (from HEC-HMS.cmd):
        - .
        - hms.jar
        - ui
        - lib/*.jar (explicit list)

        Args:
            hms_install_path: Path to HMS 3.x installation

        Returns:
            Semicolon-separated classpath string
        """
        path_sep = ";" if os.name == 'nt' else ":"

        # Build classpath in same order as HEC-HMS.cmd
        classpath_parts = [
            str(hms_install_path),                          # .
            str(hms_install_path / "hms.jar"),              # hms.jar
            str(hms_install_path / "ui"),                   # ui
        ]

        # Add all jars from lib directory
        lib_dir = hms_install_path / "lib"
        if lib_dir.exists():
            for jar_file in sorted(lib_dir.glob("*.jar")):
                classpath_parts.append(str(jar_file))

        return path_sep.join(classpath_parts)

    @staticmethod
    def _execute_via_java_3x(
        script_path: Path,
        hms_install_path: Path,
        timeout: int = 3600,
        max_memory: str = "512M",
        initial_memory: str = "32M",
    ) -> Tuple[bool, str, str]:
        """
        Execute HMS 3.x script via direct Java invocation.

        HMS 3.x differences from 4.x:
        - 32-bit Java in java/bin/ (not jre/bin/)
        - Explicit classpath (no wildcards)
        - Native DLLs in root directory (not bin/)
        - Must run from HMS installation directory

        Args:
            script_path: Path to the Jython script file
            hms_install_path: Path to HMS 3.x installation directory
            timeout: Maximum execution time in seconds
            max_memory: Maximum JVM heap size (capped for 32-bit)
            initial_memory: Initial JVM heap size

        Returns:
            Tuple of (success, stdout, stderr)
        """
        # Java executable path for HMS 3.x
        java_exe = hms_install_path / "java" / "bin" / "java.exe"
        if not java_exe.exists():
            java_exe = hms_install_path / "java" / "bin" / "java"

        if not java_exe.exists():
            return False, "", f"Java executable not found: {java_exe}"

        hms_jar = hms_install_path / "hms.jar"
        if not hms_jar.exists():
            return False, "", f"HMS jar not found: {hms_jar}"

        # Build explicit classpath for HMS 3.x
        classpath = HmsJython._build_classpath_3x(hms_install_path)

        # Cap memory for 32-bit JVM
        def parse_memory_mb(mem_str):
            mem_str = mem_str.upper()
            if mem_str.endswith('G'):
                return int(mem_str[:-1]) * 1024
            elif mem_str.endswith('M'):
                return int(mem_str[:-1])
            return int(mem_str) // (1024 * 1024)

        max_mem_mb = parse_memory_mb(max_memory)
        if max_mem_mb > 1280:  # 32-bit limit
            logger.warning(f"32-bit Java: capping max_memory from {max_memory} to 1280M")
            max_memory = "1280M"

        # Build Java command
        # From HMS 3.3 batch file: -Xcheck:jni -DMapPanel.NoVolatileImage=true
        cmd = [
            str(java_exe),
            "-Xcheck:jni",
            f"-Xms{initial_memory}",
            f"-Xmx{max_memory}",
            "-DMapPanel.NoVolatileImage=true",
            f"-Dpython.path={classpath}",
            "-Djava.library.path=.",  # HMS 3.x: native libs in root
            "-classpath", classpath,
            "hms.Hms",
            "-script", str(script_path.resolve())
        ]

        logger.debug(f"HMS 3.x Java command: {' '.join(cmd)}")

        # Execute from HMS installation directory (required for HMS 3.x)
        try:
            result = subprocess.run(
                cmd,
                capture_output=True,
                timeout=timeout,
                text=True,
                cwd=str(hms_install_path)  # Critical: run from HMS directory
            )

            stdout = result.stdout if result.stdout else ""
            stderr = result.stderr if result.stderr else ""

            # Check for success indicators
            computation_completed = "Finished computing" in stdout or "Computation completed" in stdout
            project_opened = "Finished opening project" in stdout
            script_ended = "End script" in stdout and "Exit code 0" in stdout
            has_error = "ERROR" in stdout and "Exit code 0" not in stdout

            success = (script_ended or computation_completed) and not has_error

            return success, stdout, stderr

        except subprocess.TimeoutExpired:
            logger.error(f"HMS 3.x execution timed out after {timeout} seconds")
            return False, "", f"Timeout after {timeout} seconds"

        except Exception as e:
            logger.error(f"Error executing HMS 3.x: {e}")
            return False, "", str(e)

    # -------------------------------------------------------------------------
    # Direct Java Execution (Core - HMS 4.x)
    # -------------------------------------------------------------------------

    @staticmethod
    def _execute_via_java(
        script_path: Path,
        hms_install_path: Path,
        timeout: int = 3600,
        max_memory: str = "4G",
        initial_memory: str = "128M",
        additional_java_opts: Optional[List[str]] = None
    ) -> Tuple[bool, str, str]:
        """
        Execute HMS script via direct Java invocation.

        This bypasses HEC-HMS.cmd entirely, providing:
        - Fix for path quoting bugs in HMS 4.0-4.11
        - Configurable JVM memory settings
        - Consistent behavior across all versions

        Args:
            script_path: Path to the Jython script file
            hms_install_path: Path to HMS installation directory
            timeout: Maximum execution time in seconds
            max_memory: Maximum JVM heap size (e.g., "4G", "8G", "16G")
            initial_memory: Initial JVM heap size (e.g., "128M", "256M")
            additional_java_opts: Extra JVM options

        Returns:
            Tuple of (success, stdout, stderr)
        """
        # Paths - handle different HMS version structures
        # HMS 4.4+ (64-bit): jre/bin/java.exe
        # HMS 4.0-4.3 (32-bit): java/bin/java.exe
        java_exe = hms_install_path / "jre" / "bin" / "java.exe"
        if not java_exe.exists():
            java_exe = hms_install_path / "jre" / "bin" / "java"
        if not java_exe.exists():
            # Try 32-bit location (HMS 4.0-4.3)
            java_exe = hms_install_path / "java" / "bin" / "java.exe"
        if not java_exe.exists():
            java_exe = hms_install_path / "java" / "bin" / "java"

        hms_jar = hms_install_path / "hms.jar"
        lib_dir = hms_install_path / "lib"

        # Verify critical files exist
        if not java_exe.exists():
            return False, "", f"Java executable not found: {java_exe}"
        if not hms_jar.exists():
            return False, "", f"HMS jar not found: {hms_jar}"

        # Check for 32-bit Java and adjust memory if needed
        is_32bit = HmsJython._is_32bit_java(java_exe)
        if is_32bit:
            # 32-bit JVM can't allocate more than ~1.5GB
            # Parse memory value and cap if necessary
            def parse_memory(mem_str):
                """Parse memory string like '4G' or '1280M' to bytes."""
                mem_str = mem_str.upper()
                if mem_str.endswith('G'):
                    return int(mem_str[:-1]) * 1024 * 1024 * 1024
                elif mem_str.endswith('M'):
                    return int(mem_str[:-1]) * 1024 * 1024
                return int(mem_str)

            max_32bit = parse_memory(HmsJython.MAX_MEMORY_32BIT)
            if parse_memory(max_memory) > max_32bit:
                logger.warning(f"32-bit Java detected, capping max_memory from {max_memory} to {HmsJython.MAX_MEMORY_32BIT}")
                max_memory = HmsJython.MAX_MEMORY_32BIT
            if parse_memory(initial_memory) > max_32bit:
                initial_memory = HmsJython.INITIAL_MEMORY_32BIT

        # Build classpath (semicolon-separated on Windows, colon on Unix)
        path_sep = ";" if os.name == 'nt' else ":"
        classpath = f"{hms_jar}{path_sep}{lib_dir / '*'}"

        # Native library paths for GDAL, HDF, etc.
        bin_path = hms_install_path / "bin"
        gdal_path = bin_path / "gdal"
        hdf_path = bin_path / "hdf"
        java_lib_path = f"{bin_path}{path_sep}{gdal_path}{path_sep}{hdf_path}"

        # Build Java command
        cmd = [
            str(java_exe),
            f"-Xms{initial_memory}",
            f"-Xmx{max_memory}",
            "-Dsun.java2d.d3d=false",
            "-DMapPanel.NoVolatileImage=true",
            "-Dpython.import.site=false",
            f"-Dpython.path={classpath}",
            f"-Djava.library.path={java_lib_path}",
        ]

        # Add any additional JVM options
        if additional_java_opts:
            cmd.extend(additional_java_opts)

        # Add classpath and main class
        cmd.extend([
            "-classpath", classpath,
            "hms.Hms",
            "-script", str(script_path.resolve())
        ])

        # Environment variables (required for GDAL/HDF native libraries)
        env = os.environ.copy()
        env["HMS_ROOT"] = str(hms_install_path) + os.sep

        # Update PATH for native libraries
        taudem_path = bin_path / "taudem"
        new_path = f"{bin_path}{path_sep}{gdal_path}{path_sep}{taudem_path}"
        env["PATH"] = f"{new_path}{path_sep}{env.get('PATH', '')}"

        # GDAL-specific environment
        env["PROJ_LIB"] = str(gdal_path / "projlib")
        env["GDAL_DATA"] = str(gdal_path / "gdal-data")

        logger.debug(f"Java command: {' '.join(cmd)}")

        # Execute
        try:
            result = subprocess.run(
                cmd,
                capture_output=True,
                timeout=timeout,
                text=True,
                env=env,
                cwd=str(script_path.parent)
            )

            stdout = result.stdout if result.stdout else ""
            stderr = result.stderr if result.stderr else ""

            # Check for success indicators
            # Note: HMS may return non-zero but still complete successfully
            # The JythonHms.Exit(0) call generates a SystemExit that shows in stderr
            computation_completed = "Computation completed" in stdout
            project_opened = "Project opened" in stdout
            has_error = "Error" in stderr and "SystemExit: 0" not in stderr

            success = (result.returncode == 0 or computation_completed) and not has_error

            return success, stdout, stderr

        except subprocess.TimeoutExpired:
            logger.error(f"HMS execution timed out after {timeout} seconds")
            return False, "", f"Timeout after {timeout} seconds"

        except Exception as e:
            logger.error(f"Error executing HMS: {e}")
            return False, "", str(e)

    # -------------------------------------------------------------------------
    # Public Execution Method
    # -------------------------------------------------------------------------

    @staticmethod
    @log_call
    def execute_script(
        script_content: str,
        hms_exe_path: Union[str, Path],
        working_dir: Optional[Union[str, Path]] = None,
        timeout: int = 3600,
        capture_output: bool = True,
        max_memory: str = None,
        initial_memory: str = None,
        additional_java_opts: Optional[List[str]] = None
    ) -> Tuple[bool, str, str]:
        """
        Execute a Jython script using HEC-HMS via direct Java invocation.

        Uses direct Java invocation for reliable execution across all HMS 4.x
        versions with configurable memory settings for large models.

        Args:
            script_content: Jython script content to execute
            hms_exe_path: Path to HEC-HMS executable or installation directory
            working_dir: Working directory for execution (temp dir if None)
            timeout: Maximum execution time in seconds (default: 3600 = 1 hour)
            capture_output: Whether to capture stdout/stderr (default: True)
            max_memory: Maximum JVM heap size (default: "4G")
                        Examples: "4G", "8G", "16G" for large models
            initial_memory: Initial JVM heap size (default: "128M")
            additional_java_opts: Extra JVM options
                        Examples: ["-XX:+UseG1GC", "-XX:MaxGCPauseMillis=200"]

        Returns:
            Tuple of (success: bool, stdout: str, stderr: str)

        Raises:
            RuntimeError: If HMS version < 4.0 (no scripting support)
            FileNotFoundError: If HMS installation not found

        Example:
            >>> # Default execution (4GB heap)
            >>> success, stdout, stderr = HmsJython.execute_script(
            ...     script, hms_exe_path, working_dir=project_path
            ... )

            >>> # Large model (16GB heap)
            >>> success, stdout, stderr = HmsJython.execute_script(
            ...     script, hms_exe_path, working_dir=project_path,
            ...     max_memory="16G"
            ... )

            >>> # Custom JVM options
            >>> success, stdout, stderr = HmsJython.execute_script(
            ...     script, hms_exe_path,
            ...     max_memory="8G",
            ...     additional_java_opts=["-XX:+UseG1GC"]
            ... )
        """
        # Set defaults
        if max_memory is None:
            max_memory = HmsJython.DEFAULT_MAX_MEMORY
        if initial_memory is None:
            initial_memory = HmsJython.DEFAULT_INITIAL_MEMORY

        hms_exe_path = Path(hms_exe_path)

        # Determine HMS installation directory
        if hms_exe_path.is_file():
            hms_install_path = hms_exe_path.parent
        elif hms_exe_path.is_dir():
            hms_install_path = hms_exe_path
        else:
            raise FileNotFoundError(f"HMS path not found: {hms_exe_path}")

        # Verify installation exists
        hms_jar = hms_install_path / "hms.jar"
        if not hms_jar.exists():
            raise FileNotFoundError(
                f"Invalid HMS installation - hms.jar not found at: {hms_install_path}"
            )

        # Get and check version
        version = HmsJython._get_hms_version(hms_install_path)
        HmsJython._check_version_supported(version)
        version_str = HmsJython._format_version(version)
        is_3x = HmsJython._is_hms_3x(version)

        # Adjust defaults for HMS 3.x (32-bit JVM)
        if is_3x:
            if max_memory == HmsJython.DEFAULT_MAX_MEMORY:
                max_memory = "512M"  # 32-bit default
            if initial_memory == HmsJython.DEFAULT_INITIAL_MEMORY:
                initial_memory = "32M"

        # Create script file
        if working_dir:
            working_dir = Path(working_dir)
            working_dir.mkdir(parents=True, exist_ok=True)
            script_path = working_dir / "hms_script.py"
            cleanup_script = False
        else:
            temp_dir = tempfile.mkdtemp(prefix="hms_commander_")
            script_path = Path(temp_dir) / "hms_script.py"
            cleanup_script = True

        # Write script to file
        script_path.write_text(script_content, encoding='utf-8')

        logger.info(f"Executing HMS {version_str} via direct Java invocation")
        logger.info(f"Script: {script_path}")
        logger.info(f"Memory: -Xms{initial_memory} -Xmx{max_memory}")

        try:
            # Route to appropriate execution method based on HMS version
            if is_3x:
                logger.info(f"Using HMS 3.x execution path (32-bit Java)")
                success, stdout, stderr = HmsJython._execute_via_java_3x(
                    script_path=script_path,
                    hms_install_path=hms_install_path,
                    timeout=timeout,
                    max_memory=max_memory,
                    initial_memory=initial_memory,
                )
            else:
                # HMS 4.x execution
                success, stdout, stderr = HmsJython._execute_via_java(
                    script_path=script_path,
                    hms_install_path=hms_install_path,
                    timeout=timeout,
                    max_memory=max_memory,
                    initial_memory=initial_memory,
                    additional_java_opts=additional_java_opts
                )

            if success:
                logger.info(f"HMS {version_str} script executed successfully")
            else:
                logger.error(f"HMS {version_str} script failed")
                if stderr:
                    logger.error(f"stderr: {stderr[:500]}")
                if stdout and "ERROR" in stdout:
                    logger.error(f"stdout errors: {stdout[:500]}")

            return success, stdout, stderr

        finally:
            # Cleanup temp files
            if cleanup_script and script_path.exists():
                try:
                    script_path.unlink()
                    script_path.parent.rmdir()
                except (OSError, PermissionError) as e:
                    logger.debug(f"Could not clean up temp script: {e}")

    # -------------------------------------------------------------------------
    # Script Generation Methods
    # -------------------------------------------------------------------------

    # Python 2 compatible script header (for HMS 3.x)
    SCRIPT_HEADER_PY2 = '''"""
HEC-HMS Jython Script (Python 2 compatible)
Generated by hms-commander
Generated at: {timestamp}
"""

from hms.model import JythonHms
import sys

'''

    @staticmethod
    @log_call
    def generate_compute_script(
        project_path: Union[str, Path],
        run_name: str,
        output_dss: Optional[Union[str, Path]] = None,
        save_project: bool = True,
        hms_object=None,
        python2_compatible: bool = False
    ) -> str:
        """
        Generate a Jython script to compute a simulation run.

        Args:
            project_path: Path to the HEC-HMS project folder
            run_name: Name of the simulation run to execute
            output_dss: Optional custom output DSS file path
            save_project: Whether to save project after computation
            hms_object: Optional HmsPrj instance
            python2_compatible: If True, generate Python 2 syntax for HMS 3.x

        Returns:
            Jython script content as string

        Example:
            >>> script = HmsJython.generate_compute_script(
            ...     "C:/Projects/MyProject",
            ...     "Run 1"
            ... )
            >>> # For HMS 3.x
            >>> script = HmsJython.generate_compute_script(
            ...     "C:/Projects/MyProject",
            ...     "Run 1",
            ...     python2_compatible=True
            ... )
        """
        project_path = Path(project_path)
        project_name = project_path.name

        if python2_compatible:
            return HmsJython._generate_compute_script_py2(
                project_path, project_name, run_name, save_project
            )

        # Python 3 compatible script (HMS 4.x)
        script = HmsJython.SCRIPT_HEADER.format(
            timestamp=datetime.now().isoformat()
        )

        # Open project
        script += f'''
# Open the HEC-HMS project
project_path = r"{project_path}"
project_name = "{project_name}"

try:
    JythonHms.OpenProject(project_name, project_path)
    print("Project opened successfully: " + project_name)
except Exception as e:
    print("Error opening project: " + str(e))
    JythonHms.Exit(1)
'''

        # Set output DSS if specified
        if output_dss:
            script += f'''
# Set custom output DSS file
try:
    # Note: This may require specific JythonHms methods depending on HMS version
    pass  # Placeholder for DSS output configuration
except Exception as e:
    print("Warning: Could not set output DSS: " + str(e))
'''

        # Compute the run
        script += f'''
# Compute the simulation run
run_name = "{run_name}"
try:
    JythonHms.Compute(run_name)
    print("Computation completed for: " + run_name)
except Exception as e:
    print("Error during computation: " + str(e))
    JythonHms.Exit(1)
'''

        # Save project if requested
        if save_project:
            script += '''
# Save the project (method name varies by HMS version)
try:
    # Try different method names for different HMS versions
    if hasattr(JythonHms, 'saveProject'):
        JythonHms.saveProject()
        print("Project saved successfully")
    elif hasattr(JythonHms, 'SaveProject'):
        JythonHms.SaveProject()
        print("Project saved successfully")
    else:
        print("Note: Project save method not available - results stored in DSS")
except Exception as e:
    print("Warning: Could not save project: " + str(e))
'''

        script += HmsJython.SCRIPT_FOOTER
        return script

    @staticmethod
    def _generate_compute_script_py2(
        project_path: Path,
        project_name: str,
        run_name: str,
        save_project: bool = True
    ) -> str:
        """
        Generate a Python 2 compatible Jython script for HMS 3.x.

        Uses Python 2 syntax:
        - print statement instead of print()
        - except Exception, e: instead of except Exception as e:
        """
        script = HmsJython.SCRIPT_HEADER_PY2.format(
            timestamp=datetime.now().isoformat()
        )

        # Python 2 compatible code
        script += f'''
# Open the HEC-HMS project
project_path = r"{project_path}"
project_name = "{project_name}"

try:
    JythonHms.OpenProject(project_name, project_path)
    print "Project opened successfully: " + project_name
except Exception, e:
    print "Error opening project: " + str(e)
    JythonHms.Exit(1)

# Compute the simulation run
run_name = "{run_name}"
try:
    JythonHms.Compute(run_name)
    print "Computation completed for: " + run_name
except Exception, e:
    print "Error during computation: " + str(e)
    JythonHms.Exit(1)
'''

        if save_project:
            script += '''
# Save the project
try:
    JythonHms.SaveAllProjectComponents()
    print "Project saved successfully"
except Exception, e:
    print "Warning: Could not save project: " + str(e)
'''

        script += '''
# Exit cleanly
JythonHms.Exit(0)
'''
        return script

    @staticmethod
    @log_call
    def generate_batch_compute_script(
        project_path: Union[str, Path],
        run_names: List[str],
        save_after_each: bool = False,
        hms_object=None
    ) -> str:
        """
        Generate a Jython script to compute multiple simulation runs.

        Args:
            project_path: Path to the HEC-HMS project folder
            run_names: List of simulation run names to execute
            save_after_each: Whether to save project after each run
            hms_object: Optional HmsPrj instance

        Returns:
            Jython script content as string
        """
        project_path = Path(project_path)
        project_name = project_path.name

        script = HmsJython.SCRIPT_HEADER.format(
            timestamp=datetime.now().isoformat()
        )

        # Open project
        script += f'''
# Open the HEC-HMS project
project_path = r"{project_path}"
project_name = "{project_name}"

try:
    JythonHms.OpenProject(project_name, project_path)
    print("Project opened successfully: " + project_name)
except Exception as e:
    print("Error opening project: " + str(e))
    JythonHms.Exit(1)

# List of runs to compute
run_names = {run_names}
results = {{}}

# Compute each run
for run_name in run_names:
    try:
        print("Computing: " + run_name)
        JythonHms.Compute(run_name)
        results[run_name] = "Success"
        print("Completed: " + run_name)
'''

        if save_after_each:
            script += '''
        JythonHms.SaveProject()
'''

        script += '''
    except Exception as e:
        results[run_name] = "Failed: " + str(e)
        print("Error computing " + run_name + ": " + str(e))

# Print summary
print("\\n=== Computation Summary ===")
for run_name, status in results.items():
    print(run_name + ": " + status)

# Save project at the end
try:
    JythonHms.SaveProject()
    print("Project saved successfully")
except Exception as e:
    print("Warning: Could not save project: " + str(e))
'''

        script += HmsJython.SCRIPT_FOOTER
        return script

    @staticmethod
    @log_call
    def generate_parameter_modification_script(
        project_path: Union[str, Path],
        basin_name: str,
        modifications: Dict[str, Dict[str, Any]],
        run_name: Optional[str] = None,
        hms_object=None
    ) -> str:
        """
        Generate a Jython script to modify basin parameters.

        Args:
            project_path: Path to the HEC-HMS project folder
            basin_name: Name of the basin model to modify
            modifications: Dictionary mapping element names to parameter changes
                Example: {"Subbasin-1": {"CurveNumber": 75, "InitialAbstraction": 0.5}}
            run_name: Optional run to execute after modifications
            hms_object: Optional HmsPrj instance

        Returns:
            Jython script content as string

        Example:
            >>> mods = {"Subbasin-1": {"CurveNumber": 80}}
            >>> script = HmsJython.generate_parameter_modification_script(
            ...     "C:/Projects/MyProject",
            ...     "Basin 1",
            ...     mods,
            ...     run_name="Run 1"
            ... )
        """
        project_path = Path(project_path)
        project_name = project_path.name

        script = HmsJython.SCRIPT_HEADER.format(
            timestamp=datetime.now().isoformat()
        )

        script += f'''
# Open the HEC-HMS project
project_path = r"{project_path}"
project_name = "{project_name}"

try:
    JythonHms.OpenProject(project_name, project_path)
    print("Project opened successfully: " + project_name)
except Exception as e:
    print("Error opening project: " + str(e))
    JythonHms.Exit(1)

# Open the basin model
basin_name = "{basin_name}"
try:
    basin = JythonHms.OpenBasinModel(basin_name)
    print("Basin model opened: " + basin_name)
except Exception as e:
    print("Error opening basin model: " + str(e))
    JythonHms.Exit(1)

# Modify parameters
modifications = {modifications}

for element_name, params in modifications.items():
    try:
        element = basin.getElement(element_name)
        if element is None:
            print("Warning: Element not found: " + element_name)
            continue

        for param_name, param_value in params.items():
            try:
                # Use reflection or direct setter based on parameter
                setter_name = "set" + param_name
                setter = getattr(element, setter_name, None)
                if setter:
                    setter(param_value)
                    print("Set " + element_name + "." + param_name + " = " + str(param_value))
                else:
                    print("Warning: No setter for " + param_name)
            except Exception as e:
                print("Error setting " + param_name + ": " + str(e))

    except Exception as e:
        print("Error modifying " + element_name + ": " + str(e))

# Save modifications
try:
    JythonHms.SaveProject()
    print("Project saved with modifications")
except Exception as e:
    print("Error saving project: " + str(e))
    JythonHms.Exit(1)
'''

        # Optionally run simulation
        if run_name:
            script += f'''
# Compute simulation with modified parameters
run_name = "{run_name}"
try:
    JythonHms.Compute(run_name)
    print("Computation completed: " + run_name)
except Exception as e:
    print("Error during computation: " + str(e))
    JythonHms.Exit(1)

# Save results
try:
    JythonHms.SaveProject()
    print("Results saved")
except Exception as e:
    print("Warning: Could not save results: " + str(e))
'''

        script += HmsJython.SCRIPT_FOOTER
        return script

    @staticmethod
    @log_call
    def generate_calibration_script(
        project_path: Union[str, Path],
        run_name: str,
        parameters: Dict[str, float],
        basin_name: Optional[str] = None,
        hms_object=None
    ) -> str:
        """
        Generate a Jython script for parameter calibration.

        This creates a script that modifies specified parameters and runs
        a simulation - useful for optimization/calibration workflows.

        Args:
            project_path: Path to the HEC-HMS project folder
            run_name: Name of the simulation run
            parameters: Dictionary of parameter names to values
            basin_name: Name of basin model (auto-detected if None)
            hms_object: Optional HmsPrj instance

        Returns:
            Jython script content as string
        """
        project_path = Path(project_path)
        project_name = project_path.name

        script = HmsJython.SCRIPT_HEADER.format(
            timestamp=datetime.now().isoformat()
        )

        script += f'''
# Calibration script for HEC-HMS
# Parameters: {parameters}

project_path = r"{project_path}"
project_name = "{project_name}"
run_name = "{run_name}"

# Open project
try:
    JythonHms.OpenProject(project_name, project_path)
except Exception as e:
    print("CALIBRATION_ERROR: " + str(e))
    JythonHms.Exit(1)

# Apply calibration parameters
calibration_params = {parameters}
print("CALIBRATION_PARAMS: " + str(calibration_params))

# TODO: Apply parameters to appropriate elements
# This requires knowledge of which elements each parameter applies to

# Run simulation
try:
    JythonHms.Compute(run_name)
    print("CALIBRATION_COMPLETE: " + run_name)
except Exception as e:
    print("CALIBRATION_ERROR: " + str(e))
    JythonHms.Exit(1)

# Save
JythonHms.SaveProject()
'''

        script += HmsJython.SCRIPT_FOOTER
        return script

    # -------------------------------------------------------------------------
    # Utility Methods
    # -------------------------------------------------------------------------

    @staticmethod
    @log_call
    def save_script(
        script_content: str,
        output_path: Union[str, Path]
    ) -> str:
        """
        Save a generated script to a file.

        Args:
            script_content: Jython script content
            output_path: Path to save the script

        Returns:
            Path to saved script file
        """
        output_path = Path(output_path)
        output_path.parent.mkdir(parents=True, exist_ok=True)

        with open(output_path, 'w', encoding='utf-8') as f:
            f.write(script_content)

        logger.info(f"Script saved to: {output_path}")
        return str(output_path)

    @staticmethod
    @log_call
    def find_hms_executable(
        version: Optional[str] = None,
        search_paths: Optional[List[Union[str, Path]]] = None
    ) -> Optional[Path]:
        """
        Find HEC-HMS installation on the system.

        Args:
            version: Specific version to find (e.g., "4.9", "4.10")
            search_paths: Additional paths to search

        Returns:
            Path to HEC-HMS installation directory, or None if not found
        """
        # Common installation paths
        default_paths = [
            Path("C:/Program Files/HEC/HEC-HMS"),
            Path("C:/Program Files (x86)/HEC/HEC-HMS"),
            Path(os.environ.get("HEC_HMS_HOME", "")) if os.environ.get("HEC_HMS_HOME") else None,
        ]

        # Filter None values
        default_paths = [p for p in default_paths if p]

        # Add custom search paths
        if search_paths:
            default_paths.extend([Path(p) for p in search_paths])

        # Search for HMS installation
        for base_path in default_paths:
            if not base_path.exists():
                continue

            # If version specified, look in version subdirectory
            if version:
                versioned_path = base_path / version
                if versioned_path.exists() and (versioned_path / "hms.jar").exists():
                    logger.info(f"Found HEC-HMS {version}: {versioned_path}")
                    return versioned_path

            # Otherwise search for any version
            for item in base_path.iterdir():
                if item.is_dir() and (item / "hms.jar").exists():
                    logger.info(f"Found HEC-HMS: {item}")
                    return item

        logger.warning("HEC-HMS installation not found")
        return None

    @staticmethod
    def get_jython_api_reference() -> str:
        """
        Get reference documentation for the JythonHms API.

        Returns:
            String with API reference information
        """
        return """
HEC-HMS JythonHms API Reference
================================

Basic Operations:
    JythonHms.OpenProject(name, path)     - Open a project
    JythonHms.OpenBasinModel(name)        - Open a basin model
    JythonHms.Compute(runName)            - Run a simulation
    JythonHms.SaveProject()               - Save the project
    JythonHms.Exit(code)                  - Exit (0=success, 1=error)

Basin Model Access:
    basin = JythonHms.OpenBasinModel(name)
    element = basin.getElement(elementName)

Common Element Methods:
    element.getArea()                     - Get subbasin area
    element.setArea(value)                - Set subbasin area
    element.getCurveNumber()              - Get SCS curve number
    element.setCurveNumber(value)         - Set SCS curve number
    element.getTimeOfConcentration()      - Get Tc
    element.setTimeOfConcentration(value) - Set Tc

Memory Configuration (via execute_script):
    max_memory="4G"     - Default, suitable for small/medium models
    max_memory="8G"     - Medium/large models
    max_memory="16G"    - Large models with many subbasins
    max_memory="32G"    - Very large models (requires 64-bit Java)

Custom JVM Options:
    additional_java_opts=["-XX:+UseG1GC"]           - Use G1 garbage collector
    additional_java_opts=["-XX:MaxGCPauseMillis=200"]  - Limit GC pauses
"""
