#!/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.base import CMORiser
from access_moppy.executors.dask_config import (
    peak_worker_memory_mb,
    recommend_dask_config,
    record_measured_peak,
)
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
    dask_config = None
    variable = None
    experiment_id = None
    discovery_name = None
    model_id = None
    input_files = 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.get('GRID_LABEL') or None
        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
        staging_path = os.environ.get('CMOR_STAGING_DIR') 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' }}
        split_years = {{ config.get('split_years', 'auto') | tojson }}
        source_partition_years = {{ config.get('source_partition_years') | tojson if config.get('source_partition_years') is not none else 'None' }}
        resume = {{ 'True' if config.get('resume', False) else 'False' }}

        # 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
        resume_checkpoint = 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:
                    # glob.glob's result order is filesystem-dependent (readdir
                    # order), not alphabetical/chronological, so it must be
                    # sorted explicitly to match the auto-discovery path below
                    # (discover_files always returns a sorted list). Without
                    # this, input_files[0] -- what the Dask worker-memory probe
                    # reads to size this job -- could be an arbitrary file from
                    # run to run, making the sizing tier non-deterministic for
                    # otherwise-identical jobs.
                    input_files.extend(sorted(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}")

            if resume:
                from access_moppy.batch_cmoriser import find_resume_checkpoint
                resume_checkpoint = find_resume_checkpoint(
                    drs_root or output_folder,
                    variable,
                    experiment_id=experiment_id,
                    source_id=source_id,
                    variant_label=variant_label,
                    input_files=input_files,
                    end_year=end_year,
                    split_years=split_years,
                )
                if resume_checkpoint and resume_checkpoint.complete:
                    tracker.mark_done(variable, experiment_id)
                    print(f'Completed: {variable} (all output splits already exist)')
                    return
                if resume_checkpoint and resume_checkpoint.next_year is not None:
                    from access_moppy.file_discovery import _extract_year_from_path
                    input_files = [
                        path for path in input_files
                        if (_extract_year_from_path(Path(path)) or 0) >= resume_checkpoint.next_year
                    ]
                    start_year = resume_checkpoint.next_year
                    print(
                        f'Resuming {variable} at year {start_year} '
                        f'in version v{resume_checkpoint.version_date}'
                    )

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

        input_partitions = [(None, input_files)]
        if source_partition_years is not None:
            if isinstance(source_partition_years, bool) or not isinstance(source_partition_years, int) or source_partition_years <= 0:
                raise ValueError('source_partition_years must be a positive integer')

            from access_moppy.defaults import DEFAULT_CHUNK_YEARS
            from access_moppy.base import _canonical_frequency
            from access_moppy.file_discovery import partition_files_by_year
            from access_moppy.utilities import load_model_mappings

            canonical_frequency = _canonical_frequency(discovery_name)
            if self_contained:
                print(
                    f'Source partitioning skipped for {discovery_name}: self-contained '
                    'mappings need no input files'
                )
            elif canonical_frequency not in ('mon', 'day'):
                print(
                    f'Source partitioning skipped for {discovery_name}: only monthly and '
                    f'daily variables support source partitioning (frequency is {canonical_frequency!r})'
                )
            else:
                _, mapped_name = discovery_name.split('.', 1)
                mapping_entry = load_model_mappings(discovery_name, model_id).get(mapped_name, {})
                calculation_type = (mapping_entry.get('calculation') or {}).get('type')
                if calculation_type != 'direct':
                    print(
                        f'Source partitioning skipped for {discovery_name}: calculation type '
                        f'{calculation_type!r} is not a direct mapping'
                    )
                else:
                    resolved_split_years = (
                        DEFAULT_CHUNK_YEARS.get(canonical_frequency)
                        if split_years == 'auto'
                        else split_years
                    )
                    if not isinstance(resolved_split_years, int) or isinstance(resolved_split_years, bool):
                        raise ValueError(
                            'source_partition_years requires split_years to resolve to an integer'
                        )
                    if source_partition_years % resolved_split_years != 0:
                        raise ValueError(
                            'source_partition_years must be a multiple of the resolved split_years '
                            f'({resolved_split_years})'
                        )

                    input_partitions = partition_files_by_year(
                        input_files, source_partition_years
                    )
                    print(
                        f'Source partitioning enabled: {len(input_partitions)} partitions '
                        f'of up to {source_partition_years} years'
                    )

        # 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,
            enable_chunking={{ 'True' if (config['enable_chunking'] | default(true)) else 'False' }},
            max_chunk_size_mb={{ config.get('max_chunk_size_mb', 128.0) }},
            write_prefetch={{ config.get('write_prefetch', 4) }},
        )
        client = dd.Client(
            processes=True,
            silence_logs=logging.WARNING,
            **dask_config,
        )
        print(f'Dask dashboard: {client.dashboard_link}')

        # Run each source partition through an independent CMORiser so its
        # Dask graph contains only the files needed for that period.
        written_files = []
        output_summaries = []
        shared_version_date = (
            resume_checkpoint.version_date
            if resume and resume_checkpoint and resume_checkpoint.version_date
            else None
        )
        for partition_index, (partition_start, partition_files) in enumerate(input_partitions):
            if partition_start is not None:
                print(
                    f'Processing source partition starting {partition_start}: '
                    f'{len(partition_files)} files'
                )
            cmoriser = ACCESS_ESM_CMORiser(
                input_data=partition_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,
                staging_path=staging_path,
                cmip_version=cmip_version,
                parent_info=parent_info,
                enable_resampling={{ 'True' if (config['enable_resampling'] | default(true)) else 'False' }},
                enable_chunking={{ 'True' if (config['enable_chunking'] | default(true)) else 'False' }},
                chunk_size_mb={{ config.get('chunk_size_mb', 4.0) }},
                max_chunk_size_mb={{ config.get('max_chunk_size_mb', 128.0) }},
                write_prefetch={{ config.get('write_prefetch', 4) }},
                split_years=split_years,
            )
            if shared_version_date is None:
                shared_version_date = cmoriser.cmoriser.version_date
            else:
                cmoriser.cmoriser.version_date = shared_version_date
            {% if config.get('qc_plots', False) %}
            # Preserve the existing behavior of generating one initial snapshot.
            cmoriser.cmoriser.enable_qc_plots = partition_index == 0
            cmoriser.cmoriser.qc_comparison_store = {{ config.get('cmip6_comparison_store') | tojson if config.get('cmip6_comparison_store') is not none else 'None' }}
            cmoriser.cmoriser.qc_preferred_member = {{ config.get('preferred_cmip6_member') | tojson if config.get('preferred_cmip6_member') is not none else 'None' }}
            {% endif %}
            {% if config.get('compliance_check', False) %}

            # Validate the first file this variable publishes, as soon as it is
            # published: write() calls this hook before writing any further
            # split, so a non-compliant file aborts the variable there (renamed
            # aside, report kept next to this script) rather than after the
            # whole time series has been written.
            if partition_index == 0:
                from access_moppy.qc.compliance import enforce_compliance

                def _record_compliance(record, _variable=variable, _experiment_id=experiment_id):
                    # Keep the verdict on the task row so it reaches the batch
                    # report. Without this the result exists only as a JSON
                    # report beside this script, which nothing downstream reads.
                    if tracker is not None:
                        tracker.set_compliance(_variable, _experiment_id, record)

                cmoriser.cmoriser.first_write_hook = lambda first_file: enforce_compliance(
                    first_file,
                    report_dir=Path(__file__).resolve().parent,
                    cmip_version=cmip_version,
                    min_weight={{ config.get('compliance_check_min_weight', 3) | tojson }},
                    suites={{ config.get('compliance_check_suite') | tojson if config.get('compliance_check_suite') is not none else 'None' }},
                    on_record=_record_compliance,
                )
            {% endif %}

            try:
                cmoriser.run()
                cmoriser.write()
                if isinstance(cmoriser.written_files, list):
                    written_files.extend(cmoriser.written_files)
                if isinstance(getattr(cmoriser, 'output_summary', None), dict):
                    output_summaries.append(cmoriser.output_summary)
            finally:
                try:
                    if getattr(cmoriser.cmoriser, 'ds', None) is not None:
                        cmoriser.cmoriser.ds.close()
                finally:
                    cmoriser.close()
        {% if config.get('qc_plots', False) %}

        # The first snapshot is generated during write(), before its NetCDF
        # split leaves JobFS. Generate one full-period timeseries after all
        # published splits are available.
        _qc_output_dir = Path(output_folder) / 'qc_plots'
        _comparison_store = {{ config.get('cmip6_comparison_store') | tojson if config.get('cmip6_comparison_store') is not none else 'None' }}
        _preferred_member = {{ config.get('preferred_cmip6_member') | tojson if config.get('preferred_cmip6_member') is not none else 'None' }}
        if len(written_files) > 1:
            from access_moppy.qc.plots import generate_qc_plots_for_split_files as _gen_combined_plots
            try:
                _combined_dir = _gen_combined_plots(
                    sorted(written_files),
                    qc_dir=_qc_output_dir,
                    comparison_store=_comparison_store,
                    preferred_member=_preferred_member,
                )
                if _combined_dir:
                    print(f'QC combined timeseries written -> {_combined_dir}')
            except Exception as _qc_exc:
                print(f'QC combined timeseries warning: {_qc_exc}', file=sys.stderr)
        {% endif %}

        if output_summaries:
            if len(output_summaries) == 1:
                output_summary = dict(output_summaries[0])
            else:
                total_bytes = sum(int(s.get('total_bytes', 0)) for s in output_summaries)
                files = []
                for summary in output_summaries:
                    files.extend(summary.get('files', []))
                output_summary = {
                    'file_count': len(written_files),
                    'total_bytes': total_bytes,
                    'total_size': CMORiser._format_bytes(total_bytes),
                    'files': files,
                }
                # One variable can span several CMORisers; keep the worst
                # outcome each gate produced across all of them.
                gates = CMORiser.merge_gate_results(
                    summary.get('gates', {}) for summary in output_summaries
                )
                if gates:
                    output_summary['gates'] = gates
            tracker.set_output_summary(variable, experiment_id, output_summary)
            print(
                f"Produced {variable}: {output_summary['file_count']} files, "
                f"{output_summary['total_size']}"
            )
        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:
        # Record each worker's peak RSS against the sizing this job picked, so
        # a batch's worker_memory_json rows can be compared against the
        # recommend_dask_config floors to see whether workers were over- or
        # under-provisioned. Best-effort: a dead/unreachable client just
        # yields an empty result, never blocks teardown.
        if client is not None and client.status != 'closed':
            try:
                worker_peaks = peak_worker_memory_mb(client)
                if worker_peaks and tracker is not None and variable and experiment_id:
                    tracker.set_worker_memory(variable, experiment_id, {
                        'n_workers': (dask_config or {}).get('n_workers'),
                        'memory_limit_per_worker': (dask_config or {}).get('memory_limit'),
                        'peak_rss_mb': worker_peaks,
                    })
                # Cross-experiment calibration (opt-in via
                # MOPPY_WORKER_MEMORY_HISTORY): feeds this run's measured peak
                # back so future runs of the same variable can size workers
                # from reality instead of the static file-probe heuristic.
                # n_input_files records the scale this measurement covers, so
                # a short run can never be used to under-size a longer one.
                if worker_peaks and discovery_name and model_id:
                    record_measured_peak(
                        discovery_name, model_id, worker_peaks,
                        len(input_files or []),
                    )
            except Exception:
                pass
            client.close()
        # 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 __name__ == '__main__':
    main()
