"""Build operational monthly series from observed history and LT=0..6 forecasts."""
from __future__ import annotations

import argparse
import logging
import os
from datetime import datetime, timezone
from pathlib import Path

import numpy as np
import xarray as xr

from .download_ecmwf_baseline import parse_init

LOGGER = logging.getLogger(__name__)


def combine_history_forecast(observed: xr.DataArray, forecast: xr.DataArray,
                             init: str, history_months: int = 5) -> xr.DataArray:
    """Prepend observed months to LT=0..6, broadcasting history over members."""
    parse_init(init)
    if history_months < 1:
        raise ValueError("history_months must be positive")
    if "time" not in observed.dims or not np.issubdtype(observed.time.dtype, np.datetime64):
        raise ValueError("Observed data require a datetime time dimension")
    if "lead" not in forecast.dims or "valid_time" not in forecast.coords:
        raise ValueError("Forecast requires lead dimension and valid_time coordinate")
    if forecast.valid_time.dims != ("lead",):
        raise ValueError("Forecast valid_time must be one-dimensional along lead")
    expected_lead = np.arange(7)
    if not np.array_equal(forecast.lead.values, expected_lead):
        raise ValueError("Forecast lead must be exactly 0..6")
    init_month = np.datetime64(f"{init}-01", "M")
    expected_forecast = (init_month + np.arange(7).astype("timedelta64[M]"))
    actual_forecast = np.asarray(forecast.valid_time.values).astype("datetime64[M]")
    if not np.array_equal(actual_forecast, expected_forecast):
        raise ValueError("Forecast valid_time does not match init plus LT=0..6")
    expected_history = (init_month - np.timedelta64(history_months, "M")
                        + np.arange(history_months).astype("timedelta64[M]"))
    observed_month = np.asarray(observed.time.values).astype("datetime64[M]")
    indices = []
    for month in expected_history:
        found = np.flatnonzero(observed_month == month)
        if found.size != 1:
            raise ValueError(f"Observed history requires exactly one value for {month}")
        indices.append(int(found[0]))
    history = observed.isel(time=indices).assign_coords(
        time=expected_history.astype("datetime64[ns]"))
    forecast_time = forecast.assign_coords(
        time=("lead", expected_forecast.astype("datetime64[ns]")))
    forecast_time = forecast_time.swap_dims({"lead": "time"})
    for coordinate in ("lead", "valid_time"):
        if coordinate in forecast_time.coords:
            forecast_time = forecast_time.drop_vars(coordinate)
    extra_dims = [name for name in forecast_time.dims if name not in history.dims and name != "time"]
    for name in extra_dims:
        history = history.expand_dims({name: forecast_time[name]})
    order = list(forecast_time.dims)
    history = history.transpose(*order)
    history, forecast_time = xr.align(history, forecast_time, join="exact", exclude={"time"})
    result = xr.concat([history, forecast_time], dim="time").rename(
        forecast.name or observed.name or "operational_series")
    expected_all = np.concatenate([expected_history, expected_forecast]).astype("datetime64[ns]")
    if not np.array_equal(result.time.values.astype("datetime64[ns]"), expected_all):
        raise ValueError("Combined operational series is not monthly-contiguous")
    result.attrs = dict(forecast.attrs)
    result.attrs.update({
        "initial_time": str(np.datetime64(f"{init}-01", "ns")),
        "history_months": history_months,
        "history_source": "observations through month before initial time",
        "forecast_source": "LT=0..6",
        "ensemble_processing": "observed history broadcast to forecast dimensions; members retained",
        "created_utc": datetime.now(timezone.utc).isoformat(),
    })
    return result


def validate_output(path: Path, variable: str, init: str, history_months: int) -> dict:
    with xr.open_dataset(path) as dataset:
        if variable not in dataset:
            raise ValueError(f"Operational output is missing {variable!r}")
        data = dataset[variable]
        if data.sizes.get("time") != history_months + 7:
            raise ValueError("Operational output has wrong time length")
        months = data.time.values.astype("datetime64[M]")
        if np.any(np.diff(months).astype(int) != 1):
            raise ValueError("Operational output contains a monthly gap")
        if months[history_months] != np.datetime64(f"{init}-01", "M"):
            raise ValueError("Operational output initial month is misplaced")
        return {"variable": variable, "sizes": dict(data.sizes), "bytes": path.stat().st_size}


def run(observed_path: Path, forecast_path: Path, output: Path, init: str,
        observed_variable: str, forecast_variable: str, output_variable: str,
        history_months: int = 5, overwrite: bool = False,
        retry_partial: bool = False) -> tuple[Path, dict]:
    if output.exists() and not overwrite:
        return output, validate_output(output, output_variable, init, history_months)
    partial = output.with_suffix(output.suffix + ".part")
    if partial.exists() and not (overwrite or retry_partial):
        raise FileExistsError(f"Partial output retained: {partial}")
    if overwrite:
        output.unlink(missing_ok=True)
    if overwrite or retry_partial:
        partial.unlink(missing_ok=True)
    with xr.open_dataset(observed_path) as observed_ds, xr.open_dataset(forecast_path) as forecast_ds:
        if observed_variable not in observed_ds or forecast_variable not in forecast_ds:
            raise KeyError("Configured observation or forecast variable is absent")
        combined = combine_history_forecast(
            observed_ds[observed_variable], forecast_ds[forecast_variable], init, history_months
        ).rename(output_variable).load()
    combined.attrs.update({"observed_file": str(observed_path), "forecast_file": str(forecast_path)})
    output.parent.mkdir(parents=True, exist_ok=True)
    combined.to_dataset().to_netcdf(
        partial, encoding={output_variable: {"zlib": True, "complevel": 1, "shuffle": True}}
    )
    report = validate_output(partial, output_variable, init, history_months)
    os.replace(partial, output)
    return output, report


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--observed", required=True, type=Path)
    parser.add_argument("--forecast", required=True, type=Path)
    parser.add_argument("--output", required=True, type=Path)
    parser.add_argument("--init", required=True)
    parser.add_argument("--observed-variable", required=True)
    parser.add_argument("--forecast-variable", required=True)
    parser.add_argument("--output-variable", required=True)
    parser.add_argument("--history-months", type=int, default=5)
    parser.add_argument("--overwrite", action="store_true")
    parser.add_argument("--retry-partial", action="store_true")
    args = parser.parse_args()
    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
    target, report = run(
        args.observed, args.forecast, args.output, args.init,
        args.observed_variable, args.forecast_variable, args.output_variable,
        args.history_months, args.overwrite, args.retry_partial,
    )
    LOGGER.info("Validated %s: %s", target, report)


if __name__ == "__main__":
    main()