"""
Simulation Main File Writer for IWFM models.

This module provides the main writer for IWFM simulation main input files,
orchestrating the writing of the main control file that references all
component input files.
"""

from __future__ import annotations

import logging
from dataclasses import dataclass
from datetime import datetime
from pathlib import Path
from typing import TYPE_CHECKING, Any

from pyiwfm.io.writer_base import TemplateWriter
from pyiwfm.templates.engine import TemplateEngine

if TYPE_CHECKING:
    from pyiwfm.core.model import IWFMModel

logger = logging.getLogger(__name__)


@dataclass
class SimulationMainConfig:
    """
    Configuration for simulation main file writing.

    Attributes
    ----------
    output_dir : Path
        Base output directory for simulation files
    """

    output_dir: Path

    # File names
    main_file: str = "Simulation_MAIN.IN"
    preprocessor_bin: str = "PreProcessor.bin"

    # Component file paths (relative to output_dir)
    gw_main: str = "GW\\GW_MAIN.dat"
    stream_main: str = "Stream\\Stream_MAIN.dat"
    lake_main: str = "Lake\\Lake_MAIN.dat"
    rootzone_main: str = "RootZone\\RootZone_MAIN.dat"
    swshed_main: str = ""  # SmallWatershed
    unsatzone_main: str = ""  # Unsaturated zone
    irig_frac: str = "IrigFrac.dat"
    supply_adjust: str = "SupplyAdjust.dat"
    precip_file: str = "Precip.dat"
    et_file: str = "ET.dat"
    kc_file: str = ""  # Crop coefficient

    # Simulation period (default 10-year simulation)
    begin_date: str = "09/30/1990_24:00"
    end_date: str = "09/30/2000_24:00"
    time_step: str = "1DAY"
    restart: int = 0

    # Output options
    istrt: int = 0  # Restart file generation
    kdeb: int = 0  # Debug output
    cache_size: int = 500000

    # Solution scheme
    matrix_solver: int = 2  # 1=SOR, 2=Conjugate gradient
    relaxation: float = 1.0
    max_iterations: int = 1500
    max_supply_iter: int = 50
    convergence_head: float = 0.0001
    convergence_volume: float = 0.001
    convergence_supply: float = 0.001

    # Supply adjustment (11 = adjust both GW and diversions)
    supply_adjust_option: int = 11

    # Title lines
    title1: str = "IWFM"
    title2: str = "Generated by pyiwfm"
    title3: str = ""

    @property
    def main_path(self) -> Path:
        """Get the main file path."""
        return self.output_dir / self.main_file


class SimulationMainWriter(TemplateWriter):
    """
    Writer for IWFM Simulation Main Input File.

    Writes the main simulation control file that references all component files.

    Example
    -------
    >>> from pyiwfm.io.simulation_writer import SimulationMainWriter, SimulationMainConfig
    >>> config = SimulationMainConfig(output_dir=Path("model/Simulation"))
    >>> writer = SimulationMainWriter(model, config)
    >>> path = writer.write_main()
    """

    def __init__(
        self,
        model: IWFMModel,
        config: SimulationMainConfig,
        template_engine: TemplateEngine | None = None,
    ) -> None:
        """
        Initialize the simulation main writer.

        Parameters
        ----------
        model : IWFMModel
            Model to write
        config : SimulationMainConfig
            Output file configuration
        template_engine : TemplateEngine, optional
            Custom template engine
        """
        super().__init__(config.output_dir, template_engine)
        self.model = model
        self.config = config

    @property
    def format(self) -> str:
        return "iwfm_simulation"

    def write(self, data: Any = None) -> None:
        """Write all simulation files."""
        self.write_main()

    def write_main(self) -> Path:
        """
        Write the simulation main input file.

        Returns
        -------
        Path
            Path to written file
        """
        output_path = self.config.main_path
        self._ensure_dir(output_path)

        content = self._render_simulation_main()

        output_path.write_text(content)
        logger.info(f"Wrote simulation main file: {output_path}")
        return output_path

    def _render_simulation_main(self) -> str:
        """Render the simulation main file using Jinja2 template."""
        generation_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")

        context = {
            "generation_time": generation_time,
            "title1": self.config.title1,
            "title2": self.config.title2,
            "title3": self.config.title3,
            "preprocessor_bin": self.config.preprocessor_bin,
            "gw_main": self.config.gw_main,
            "stream_main": self.config.stream_main,
            "lake_main": self.config.lake_main,
            "rootzone_main": self.config.rootzone_main,
            "swshed_main": self.config.swshed_main,
            "unsatzone_main": self.config.unsatzone_main,
            "irig_frac": self.config.irig_frac,
            "supply_adjust": self.config.supply_adjust,
            "precip_file": self.config.precip_file,
            "et_file": self.config.et_file,
            "kc_file": self.config.kc_file,
            "begin_date": self.config.begin_date,
            "restart": self.config.restart,
            "time_step": self.config.time_step,
            "end_date": self.config.end_date,
            "istrt": self.config.istrt,
            "kdeb": self.config.kdeb,
            "cache_size": self.config.cache_size,
            "matrix_solver": self.config.matrix_solver,
            "relaxation": self.config.relaxation,
            "max_iterations": self.config.max_iterations,
            "max_supply_iter": self.config.max_supply_iter,
            "convergence_head": self.config.convergence_head,
            "convergence_volume": self.config.convergence_volume,
            "convergence_supply": self.config.convergence_supply,
            "supply_adjust_option": self.config.supply_adjust_option,
        }

        return self._engine.render_template("simulation/main.j2", **context)


def write_simulation_main(
    model: IWFMModel,
    output_dir: Path | str,
    config: SimulationMainConfig | None = None,
) -> Path:
    """
    Write simulation main file for a model.

    Parameters
    ----------
    model : IWFMModel
        Model to write
    output_dir : Path or str
        Output directory
    config : SimulationMainConfig, optional
        File configuration

    Returns
    -------
    Path
        Path to written file
    """
    output_dir = Path(output_dir)

    if config is None:
        config = SimulationMainConfig(output_dir=output_dir)
    else:
        config.output_dir = output_dir

    writer = SimulationMainWriter(model, config)
    return writer.write_main()
