#!/usr/bin/env python
"""
CMORisation script for variable {{ variable }}
Generated automatically by ACCESS-MOPPy batch processing
"""

import os
import glob
import logging
import sys
from pathlib import Path
import dask.distributed as dd

from access_moppy import ACCESS_ESM_CMORiser
from access_moppy.executors.dask_config import recommend_dask_config
from access_moppy.tracking import TaskTracker


def main():
    import dask
    dask.config.set({
        'distributed.worker.memory.target': 0.6,
        'distributed.worker.memory.spill': 0.7,
        'distributed.worker.memory.pause': 0.8,
    })

    # The Dask client is created *after* input files are discovered, so it can
    # be sized to this variable's actual read footprint (recommend_dask_config
    # reads PBS_NCPUS / PBS_MEM_GB, so it adapts to whatever the job requested).
    client = None
    tracker = None

    try:
        # Get values from environment
        variable = os.environ['VARIABLE']
        db_path = os.environ['CMOR_TRACKER_DB']
        experiment_id = os.environ['EXPERIMENT_ID']
        source_id = os.environ['SOURCE_ID']
        variant_label = os.environ['VARIANT_LABEL']
        grid_label = os.environ['GRID_LABEL']
        activity_id = os.environ['ACTIVITY_ID']
        input_folder = os.environ['INPUT_FOLDER']
        output_folder = os.environ['OUTPUT_FOLDER']
        drs_root = os.environ.get('DRS_ROOT') or None
        cmip_version = os.environ.get('CMIP_VERSION', 'CMIP6')

        # File patterns — explicit overrides from the batch config take precedence;
        # if none is supplied the mapping-driven auto-discovery is used instead.
        file_patterns = {{ config.get('file_patterns', {}) | tojson }}
        model_id = os.environ.get('MODEL_ID', 'ACCESS-ESM1-6')

        # Optional year-range filtering for auto-discovery (pi-control / spinup runs).
        # Accepts integers or zero-padded strings (e.g. "0001").
        start_year = {{ config['start_year'] | tojson if 'start_year' in config else 'None' }}
        end_year   = {{ config['end_year']   | tojson if 'end_year'   in config else 'None' }}

        # Optional parent-experiment metadata. Keys left out fall back to the
        # package defaults, which describe a CMIP6 ACCESS-ESM1-5 piControl parent
        # and are rejected by the CMIP7 controlled vocabularies.
        parent_info = {{ config.get('parent_info', {}) | tojson }} or None

        # The tracker is opened before file discovery so that a discovery failure is
        # still recorded against this task instead of being lost.
        tracker = TaskTracker(Path(db_path))

        if tracker.is_done(variable, experiment_id):
            print(f'Skipped: {variable} (already done)')
            return

        tracker.mark_running(variable, experiment_id)

        # File discovery is driven by the CMIP6-style 'table.variable' name, so a
        # CMIP7 branded name must be translated first. The CMORiser still receives
        # the original name, which it needs for the CMIP7 DRS and branding metadata.
        discovery_name = variable
        if cmip_version == 'CMIP7':
            from access_moppy.utilities import _get_cmip7_to_cmip6_mapping
            discovery_name = _get_cmip7_to_cmip6_mapping(variable) or variable

        # Variables whose mapping is self-contained (an internal calculation, or a
        # formula that pulls every field from a bundled resource) have no input to
        # discover. Running discovery for them can only fail — on a missing fx
        # pattern, or an unsubstitutable {model_var} — for files nothing would read.
        from access_moppy.utilities import is_self_contained_variable
        self_contained = is_self_contained_variable(variable, model_id)

        input_files = None
        if self_contained:
            print(f'{variable} needs no input files (self-contained mapping); skipping file discovery')
        else:
            pattern = file_patterns.get(variable)
            if pattern:
                patterns = pattern if isinstance(pattern, list) else [pattern]
                input_files = []
                for p in patterns:
                    input_files.extend(glob.glob(os.path.join(input_folder, p.lstrip("/")), recursive=True))
            else:
                # Auto-discover using the model mapping file_discovery config
                from access_moppy.file_discovery import (
                    FileDiscoveryError,
                    discover_files,
                )
                try:
                    input_files = [str(f) for f in discover_files(
                        input_folder, discovery_name, model_id=model_id,
                        start_year=start_year, end_year=end_year,
                    )]
                except FileDiscoveryError as exc:
                    raise ValueError(
                        f"Could not auto-discover files for '{variable}': {exc}\n"
                        "Add a 'file_patterns' entry in the batch config or a "
                        "'file_pattern' field to the variable's mapping entry."
                    ) from exc

            if not input_files:
                from access_moppy.file_discovery import _diagnose_no_files
                diag = _diagnose_no_files(
                    Path(input_folder), discovery_name, model_id,
                    start_year=start_year, end_year=end_year,
                )
                raise ValueError(f"No files found for variable '{variable}'. {diag}")

        print(f'Processing {variable} with {len(input_files or [])} files')

        # Size the Dask client to this variable's footprint, then start it.
        # processes=True: each worker is a separate process with its own HDF5
        # lock, so netCDF4 reads run in parallel across workers (a threaded
        # client serialises them on the global HDF5 lock and collapses to ~1
        # core). Worker count / memory adapt to PBS_NCPUS / PBS_MEM_GB.
        dask_config = recommend_dask_config(discovery_name, input_files or [], model_id)
        client = dd.Client(
            processes=True,
            silence_logs=logging.WARNING,
            **dask_config,
        )
        print(f'Dask dashboard: {client.dashboard_link}')

        # Run CMORisation
        cmoriser = ACCESS_ESM_CMORiser(
            input_data=input_files,
            compound_name=variable,
            experiment_id=experiment_id,
            source_id=source_id,
            variant_label=variant_label,
            grid_label=grid_label,
            activity_id=activity_id,
            output_path=output_folder,
            drs_root=drs_root,
            cmip_version=cmip_version,
            parent_info=parent_info,
            enable_resampling={{ 'True' if (config['enable_resampling'] | default(true)) else 'False' }},
            # Stream the write in time-chunks instead of materialising the whole
            # variable in one .compute(). Without this the write is a single
            # monolithic gather ('finalize') over the entire time series, which
            # exceeds a worker's memory on long (multi-century) daily runs and
            # kills the job. Chunked writing bounds per-worker memory regardless
            # of dataset length.
            enable_chunking={{ 'True' if (config['enable_chunking'] | default(true)) else 'False' }},
        )

        cmoriser.run()
        cmoriser.write()

        tracker.mark_done(variable, experiment_id)
        print(f'Completed: {variable}')

    except Exception as e:
        import traceback
        traceback.print_exc(file=sys.stderr)
        print(f'Error processing {variable}: {e}', file=sys.stderr)
        try:
            tracker.mark_failed(variable, experiment_id, str(e))
        except:
            pass
        sys.exit(1)

    finally:
        # Release the sqlite handle before the process exits so the journal
        # file (DELETE mode) is unlinked promptly. On Lustre, lingering
        # handles can leave .db-journal behind and confuse downstream tools.
        if tracker is not None:
            try:
                tracker.close()
            except Exception:
                pass
        if client is not None and client.status != 'closed':
            client.close()

if __name__ == '__main__':
    main()
