"""
JVM configuration for jgo.

Handles JVM arguments like heap size, GC options, and system properties.
"""

from __future__ import annotations

import psutil

from ._gc import get_default_gc_options


def _flatten_properties(d: dict, prefix: str = "") -> dict:
    """
    Flatten a nested dict into dotted string keys.

    TOML's dotted key notation (e.g., ``app.name = "foo"`` under a table)
    produces nested dicts.  This converts them back to dotted keys so they
    can be passed as ``-Dapp.name=foo`` JVM arguments.
    """
    result = {}
    for key, value in d.items():
        full_key = f"{prefix}.{key}" if prefix else key
        if isinstance(value, dict):
            result.update(_flatten_properties(value, full_key))
        else:
            result[full_key] = value
    return result


class JVMConfig:
    """
    Configuration for JVM execution.

    Handles memory settings, GC options, system properties, and other JVM arguments.
    """

    def __init__(
        self,
        max_heap: str | None = None,
        min_heap: str | None = None,
        gc_options: list[str] | None = None,
        system_properties: dict[str, str] | None = None,
        extra_args: list[str] | None = None,
        auto_heap: bool = True,
        default_gc: str | None = None,
    ):
        """
        Initialize JVM configuration.

        Args:
            max_heap: Maximum heap size (e.g., "2G", "512M"). If None and auto_heap=True,
                     auto-detects based on system memory.
            min_heap: Minimum heap size (e.g., "512M")
            gc_options: Garbage collection options (e.g., ["-XX:+UseG1GC"]).
                       If None, smart defaults based on Java version will be used.
            system_properties: System properties as dict (e.g., {"foo": "bar"} -> -Dfoo=bar)
            extra_args: Additional JVM arguments
            auto_heap: If True and max_heap is None, auto-detect max heap size
            default_gc: DEPRECATED. Legacy default GC option. Use gc_options instead.
        """
        self.max_heap = max_heap
        self.min_heap = min_heap
        self.gc_options = gc_options  # Keep as None to distinguish from empty list
        self.system_properties = system_properties or {}
        self.extra_args = extra_args or []
        self.auto_heap = auto_heap
        self.default_gc = default_gc  # Kept for backward compatibility

    def to_jvm_args(self, java_version: int | None = None) -> list[str]:
        """
        Convert configuration to JVM arguments list.

        Args:
            java_version: Java major version for smart defaults (e.g., 8, 11, 17, 21).
                         Required if gc_options is None and you want version-aware defaults.

        Returns:
            List of JVM arguments suitable for java command
        """
        args = []

        # Add GC options with precedence:
        # 1. Explicit gc_options (could be empty list to disable)
        # 2. Legacy default_gc (backward compatibility)
        # 3. Smart defaults based on Java version
        # 4. Fallback to G1GC (safe for Java 8+)
        if self.gc_options is not None:
            # Explicit gc_options set (could be empty list to disable)
            args.extend(self.gc_options)
        elif self.default_gc:
            # Legacy default_gc parameter (backward compatibility)
            args.append(self.default_gc)
        elif java_version is not None:
            # Smart default based on Java version
            args.extend(get_default_gc_options(java_version))
        else:
            # Fallback: use G1GC (safe for Java 8+, most common case)
            args.append("-XX:+UseG1GC")

        # Add heap settings
        if self.min_heap:
            args.append(f"-Xms{self.min_heap}")

        max_heap = self.max_heap
        if max_heap is None and self.auto_heap:
            max_heap = self._auto_detect_max_heap()

        if max_heap:
            args.append(f"-Xmx{max_heap}")

        # Add system properties (flatten nested dicts from TOML dotted keys)
        for key, value in _flatten_properties(self.system_properties).items():
            args.append(f"-D{key}={value}")

        # Add extra args
        args.extend(self.extra_args)

        return args

    def _auto_detect_max_heap(self) -> str:
        """
        Auto-detect maximum heap size based on system memory.

        Sets max heap to half of system memory, rounded to a reasonable value.
        Uses GB for systems with >2GB RAM, MB otherwise.

        Returns:
            Heap size string (e.g., "2G", "512M")
        """
        total_memory = psutil.virtual_memory().total

        # Use GB for systems with >2GB, otherwise MB
        if total_memory > 2 * 1024**3:
            # GB mode
            max_heap_gb = max(total_memory // (1024**3) // 2, 1)
            return f"{max_heap_gb}G"
        else:
            # MB mode
            max_heap_mb = max(total_memory // (1024**2) // 2, 1)
            return f"{max_heap_mb}M"

    def with_system_property(self, key: str, value: str) -> JVMConfig:
        """
        Return a new JVMConfig with an additional system property.

        Args:
            key: Property key
            value: Property value

        Returns:
            New JVMConfig instance
        """
        new_props = self.system_properties.copy()
        new_props[key] = value

        return JVMConfig(
            max_heap=self.max_heap,
            min_heap=self.min_heap,
            gc_options=self.gc_options.copy() if self.gc_options is not None else None,
            system_properties=new_props,
            extra_args=self.extra_args.copy(),
            auto_heap=self.auto_heap,
            default_gc=self.default_gc,
        )

    def with_extra_arg(self, arg: str) -> JVMConfig:
        """
        Return a new JVMConfig with an additional JVM argument.

        Args:
            arg: JVM argument to add

        Returns:
            New JVMConfig instance
        """
        new_args = self.extra_args.copy()
        new_args.append(arg)

        return JVMConfig(
            max_heap=self.max_heap,
            min_heap=self.min_heap,
            gc_options=self.gc_options.copy() if self.gc_options is not None else None,
            system_properties=self.system_properties.copy(),
            extra_args=new_args,
            auto_heap=self.auto_heap,
            default_gc=self.default_gc,
        )
