import concurrent.futures
import re
import sys
from itertools import chain
from pathlib import Path
from typing import Optional

import click
import xarray as xr

from wrf_ensembly import experiment, external, nco, processors
from wrf_ensembly import statistics as stats
from wrf_ensembly import utils
from wrf_ensembly.click_utils import GroupWithStartEndPrint, pass_experiment_path
from wrf_ensembly.console import logger
from wrf_ensembly.processors import ProcessingContext, process_file_with_pipeline


@click.group(name="postprocess", cls=GroupWithStartEndPrint)
def postprocess_cli():
    pass


@postprocess_cli.command()
@pass_experiment_path
def print_variables_to_keep(experiment_path: Path):
    """
    Prints which variables will be kept in the wrfout files after postprocessing.
    This is useful to check if the variables you want to keep are actually kept or check
    how the regex filters are applied.

    The variables are defined in the `PostprocessConfig:variables_to_keep` variable of the
    configuration file.

    Cycle 0 output for member 0 must exist for this command to work.
    """

    logger.setup("postprocess-print_variables_to_keep", experiment_path)
    exp = experiment.Experiment(experiment_path)

    if not exp.cfg.postprocess.variables_to_keep:
        logger.warning("`variables_to_keep` is not set in the config, exiting")
        sys.exit(0)

    # Find a sample forecast file, cycle 0, member 0
    scratch_forecast_dir = exp.paths.scratch_forecasts_path(0)
    forecast_files = list(scratch_forecast_dir.rglob("member_00/wrfout*"))
    if len(forecast_files) == 0:
        logger.error("No forecast files found, exiting")
        sys.exit(1)

    # Use the first file
    forecast_file = forecast_files[0]
    logger.info(f"Using {forecast_file} as sample file")
    ds = xr.open_dataset(forecast_file, decode_times=False)

    # Add variables that would normally be generated by xwrf_post
    for new_var in [
        "air_potential_temperature",
        "air_pressure",
        "wind_east",
        "wind_north",
        "air_density",
    ]:
        ds[new_var] = xr.zeros_like(ds["THM"])
    ds = ds.rename({"XTIME": "t"})

    filters = [re.compile(x) for x in exp.cfg.postprocess.variables_to_keep]
    logger.info("Variables in output file:")
    for var in ds.data_vars:
        for f in filters:
            if f.match(str(var)):
                logger.info(f"{f} -> {var}: ({ds[var].dims}) ({ds[var].dtype})")
                break


@postprocess_cli.command()
@click.option(
    "--cycle",
    type=int,
    help="Cycle to compute statistics for. Will compute for current cycle if missing.",
)
@click.option(
    "--jobs",
    type=click.IntRange(min=0, max=None),
    help="How many files to process in parallel",
)
@pass_experiment_path
def process_pipeline(experiment_path: Path, cycle: Optional[int], jobs: Optional[int]):
    """
    Apply the configured data processor pipeline to output files.

    This command replaces apply-scripts with a more efficient plugin-based approach
    that processes data in memory using a configurable pipeline of DataProcessor instances.
    """

    logger.setup("postprocess-process-pipeline", experiment_path)
    exp = experiment.Experiment(experiment_path)

    if cycle is None:
        cycle = exp.current_cycle_i

    logger.info(f"Cycle: {exp.cycles[cycle]}")

    jobs = utils.determine_jobs(jobs)
    logger.info(f"Using {jobs} jobs")

    try:
        pipeline = processors.create_pipeline_from_config(
            exp.cfg.postprocess.processors
        )
        logger.info(
            f"Created pipeline with processors: {[p.name for p in pipeline.processors]}"
        )
    except Exception as e:
        logger.error(f"Failed to create processor pipeline: {e}")
        sys.exit(1)

    scratch_forecast_dir = exp.paths.scratch_forecasts_path(cycle)
    scratch_analysis_dir = exp.paths.scratch_analysis_path(cycle)

    # Clean any old _post files first
    for f in scratch_forecast_dir.rglob("wrfout*_post"):
        logger.debug(f"Removing old file {f}")
        f.unlink()
    for f in scratch_analysis_dir.rglob("wrfout*_post"):
        logger.debug(f"Removing old file {f}")
        f.unlink()

    files_to_process = list(scratch_analysis_dir.rglob("member_*/wrfout*")) + list(
        scratch_forecast_dir.rglob("member_*/wrfout*")
    )
    # Filter out any _post files that might still exist, just in case?
    files_to_process = [f for f in files_to_process if not f.name.endswith("_post")]

    logger.info(f"Found {len(files_to_process)} files to process")

    if len(files_to_process) == 0:
        logger.warning("No files found to process")
        return

    # Create processor arguments for each file
    process_args = []
    for file_path in files_to_process:
        # Extract metadata from file path
        member_match = re.search(r"member_(\d+)", str(file_path))
        member_id = int(member_match.group(1)) if member_match else 0

        # Create output file path with _post suffix
        output_file = file_path.parent / f"{file_path.name}_post"

        context = ProcessingContext(
            member=member_id,
            cycle=cycle,
            input_file=file_path,
            output_file=output_file,
            config=exp.cfg,
        )

        process_args.append(
            (file_path, output_file, exp.cfg.postprocess.processors, context)
        )

    logger.info(f"Processing {len(files_to_process)} files through pipeline")
    with concurrent.futures.ProcessPoolExecutor(max_workers=jobs) as executor:
        futures = [
            executor.submit(process_file_with_pipeline, args) for args in process_args
        ]

        # Wait for all to complete and handle any errors
        for future in concurrent.futures.as_completed(futures):
            try:
                future.result()
            except Exception as e:
                logger.error(f"Pipeline processing failed: {e}")
                sys.exit(1)

    logger.info("Pipeline processing completed successfully")


@postprocess_cli.command()
@click.option(
    "--cycle",
    type=int,
    help="Cycle to compute statistics for. Will compute for all current cycle if missing.",
)
@click.option(
    "--jobs",
    type=click.IntRange(min=0, max=None),
    help="How many files to process in parallel",
)
@pass_experiment_path
def statistics(
    experiment_path: Path,
    cycle: Optional[int],
    jobs: int,
):
    """
    Calculates the ensemble mean and standard deviation from the forecast/analysis files of given cycle.
    This function reads the `*_post` files created by the `wrf_post` and, optionally, `apply_scripts` commands.
    """

    logger.setup("postprocess-statistics", experiment_path)
    exp = experiment.Experiment(experiment_path)

    if cycle is None:
        cycle = exp.current_cycle_i
    logger.info(f"Cycle: {exp.cycles[cycle]}")

    jobs = utils.determine_jobs(jobs)
    logger.info(f"Using {jobs} jobs")

    # Escape hatch when only one member is present
    if exp.cfg.assimilation.n_members == 1:
        logger.warning(
            "Only one member, copying file over unchanged, no standard deviation."
        )

        analysis_path = exp.paths.scratch_analysis_path(cycle) / "member_00"
        forecast_path = exp.paths.scratch_forecasts_path(cycle) / "member_00"
        print(forecast_path)
        for f in chain(
            analysis_path.rglob("wrfout*_post"),
            forecast_path.rglob("wrfout*_post"),
        ):
            print(f)
            target = f.parent.parent / f.name.replace("_post", "_mean")
            utils.copy(f, target)

        return

    # We have to use `stats.compute_ensemble_statistics` once for each output file in the cycle (both forecast and analysis)
    # First, find out how many output files we have by checking the first member
    args: list[tuple[list[Path], tuple[Path, Path]]] = []

    scratch_forecast_dir = exp.paths.scratch_forecasts_path(cycle)
    forecast_filenames = [
        x.name for x in scratch_forecast_dir.rglob("member_00/wrfout*_post")
    ]
    for name in forecast_filenames:
        output_mean_file = scratch_forecast_dir / name.replace("_post", "_mean")
        output_sd_file = scratch_forecast_dir / name.replace("_post", "_sd")

        input_files = [
            scratch_forecast_dir / f"member_{i:02d}/{name}"
            for i in range(exp.cfg.assimilation.n_members)
        ]

        args.append((input_files, (output_mean_file, output_sd_file)))

    # Now do the same for the analysis files
    scratch_analysis_dir = exp.paths.scratch_analysis_path(cycle)
    analysis_filenames = [
        x.name for x in scratch_analysis_dir.rglob("member_00/wrfout*_post")
    ]
    for name in analysis_filenames:
        output_mean_file = scratch_analysis_dir / name.replace("_post", "_mean")
        output_sd_file = scratch_analysis_dir / name.replace("_post", "_sd")

        input_files = [
            scratch_analysis_dir / f"member_{i:02d}/{name}"
            for i in range(exp.cfg.assimilation.n_members)
        ]

        args.append((input_files, (output_mean_file, output_sd_file)))

    logger.info(f"Found {len(args)} filesets to process")

    # Call `compute_ensemble_statistics` for each fileset, in parallel
    with concurrent.futures.ProcessPoolExecutor(max_workers=jobs) as executor:
        futures = []
        for input_files, output_files in args:
            futures.append(
                executor.submit(
                    stats.compute_ensemble_statistics,
                    input_files,
                    output_files[0],
                    output_files[1],
                )
            )

        for res in concurrent.futures.as_completed(futures):
            res.result()

    logger.info("All files processed successfully")


@postprocess_cli.command()
@click.option(
    "--cycle",
    type=int,
    help="Cycle to compute statistics for. Will compute for all current cycle if missing.",
)
@click.option(
    "--jobs",
    type=click.IntRange(min=0, max=None),
    default=4,
    help="How many NCO commands to execute in parallel",
)
@pass_experiment_path
def concatenate(
    experiment_path: Path,
    cycle: Optional[int],
    jobs: int,
):
    """
    Concatenates all output files (mean and standard deviation) into two files, one for analysis
    and one for forecast. It uses the `_mean` and `_sd` files created by the `statistics` command.
    """

    logger.setup("postprocess-statistics", experiment_path)
    exp = experiment.Experiment(experiment_path)

    if cycle is None:
        cycle = exp.current_cycle_i

    commands = []

    # Prepare compression related arguments
    cmp_args = []
    if len(exp.cfg.postprocess.ppc_filter) > 0:
        cmp_args = ["--ppc", exp.cfg.postprocess.ppc_filter]
    if len(exp.cfg.postprocess.compression_filters) > 0:
        cmp_args.append(f"--cmp={exp.cfg.postprocess.compression_filters}")

    if len(cmp_args) == 0:
        logger.warning("No compression filters set, output files will be uncompressed")
    else:
        logger.info(f"Using compression filters: {' '.join(cmp_args)}")

    # Find all forecast files
    forecast_dir = exp.paths.forecast_path(cycle)
    scratch_forecast_dir = exp.paths.scratch_forecasts_path(cycle)
    forecast_files = sorted(scratch_forecast_dir.rglob("*_mean"))
    if len(forecast_files) > 0:
        commands.append(
            nco.concatenate(
                exp.cfg.postprocess.ncrcat_cmd,
                forecast_files,
                forecast_dir / f"forecast_mean_cycle_{cycle:03d}.nc",
                cmp_args,
            )
        )
    forecast_files = sorted(scratch_forecast_dir.rglob("*_sd"))
    if len(forecast_files) > 0:
        commands.append(
            nco.concatenate(
                exp.cfg.postprocess.ncrcat_cmd,
                forecast_files,
                forecast_dir / f"forecast_sd_cycle_{cycle:03d}.nc",
                cmp_args,
            )
        )

    # Find all analysis files
    analysis_dir = exp.paths.analysis_path(cycle)
    scratch_analysis_dir = exp.paths.scratch_analysis_path(cycle)
    analysis_files = sorted(scratch_analysis_dir.rglob("*_mean"))
    if len(analysis_files) > 0:
        commands.append(
            nco.concatenate(
                exp.cfg.postprocess.ncrcat_cmd,
                analysis_files,
                analysis_dir / f"analysis_mean_cycle_{cycle:03d}.nc",
                cmp_args,
            )
        )
    analysis_files = sorted(scratch_analysis_dir.rglob("*_sd"))
    if len(analysis_files) > 0:
        commands.append(
            nco.concatenate(
                exp.cfg.postprocess.ncrcat_cmd,
                analysis_files,
                analysis_dir / f"analysis_sd_cycle_{cycle:03d}.nc",
                cmp_args,
            )
        )

    # Concatenate per-member if enabled
    if exp.cfg.postprocess.keep_per_member:
        for i in range(exp.cfg.assimilation.n_members):
            forecast_files = list(
                scratch_forecast_dir.glob(f"member_{i:02d}/wrfout*_post")
            )
            commands.append(
                nco.concatenate(
                    exp.cfg.postprocess.ncrcat_cmd,
                    sorted(forecast_files),
                    forecast_dir / f"forecast_member_{i:02d}_cycle_{cycle:03d}.nc",
                    cmp_args,
                )
            )

    failure = False
    logger.info(
        f"Executing {len(commands)} nco commands in parallel, using {jobs} jobs"
    )
    for res in external.run_in_parallel(commands, jobs):
        if res.returncode != 0:
            logger.error(f"nco command failed with exit code {res.returncode}")
            logger.error(res.output)
            failure = True

    if failure:
        logger.error("One or more nco commands failed, exiting")
        sys.exit(1)


@postprocess_cli.command()
@click.option(
    "--cycle",
    type=int,
    help="Cycle to clean up. Will clean for current cycle if missing.",
)
@click.option(
    "--remove-wrfout",
    default=True,
    is_flag=True,
    help="Remove the raw wrfout files",
)
@pass_experiment_path
def clean(experiment_path: Path, cycle: Optional[int], remove_wrfout: bool):
    """
    Clean up the scratch directory for the given cycle. Use after running the other
    postprocessing commands to save disk space.
    """

    logger.setup("postprocess-statistics", experiment_path)
    exp = experiment.Experiment(experiment_path)

    if cycle is None:
        cycle = exp.current_cycle_i

    logger.info(f"Cleaning scratch for cycle {cycle}")

    scratch_dirs = [
        exp.paths.scratch_forecasts_path(cycle),
        exp.paths.scratch_analysis_path(cycle),
    ]
    for dir in scratch_dirs:
        if remove_wrfout:
            for f in dir.rglob("wrfout*"):
                logger.info(f"Removing {f}")
                f.unlink()
