"""
Precondition module — Data preprocessing utilities for seismograms.

Provides component merging, time alignment, and data trimming/splicing.
"""

import logging

import obspy
from obspy import UTCDateTime

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Merge components (> 3 traces)
# ---------------------------------------------------------------------------

def merge_components(
    input_path: str,
    output_path: str,
    method: int = 1,
    fill_value: str = "interpolate",
) -> "obspy.Stream":
    """Merge traces when a stream has more than 3 components.

    Parameters
    ----------
    input_path : str
        Path to input MiniSEED file.
    output_path : str
        Path to save the merged output.
    method : int
        ObsPy merge method (default 1).
    fill_value : str
        Fill value strategy (default ``'interpolate'``).

    Returns
    -------
    obspy.Stream
        The (possibly merged) stream.
    """
    st = obspy.read(input_path)
    n_traces = len(st)
    logger.info("Read %d trace(s) from %s", n_traces, input_path)

    if n_traces > 3:
        logger.info("Merging %d components …", n_traces)
        st.merge(method=method, interpolation_samples=-1, fill_value=fill_value)
        st.write(output_path, format="MSEED")
        logger.info("Merged to %d trace(s), saved to %s", len(st), output_path)
    else:
        logger.info("Stream already has ≤ 3 components; copying as-is.")
        st.write(output_path, format="MSEED")

    return st


# ---------------------------------------------------------------------------
# Align time windows
# ---------------------------------------------------------------------------

def align_times(
    input_path: str,
    output_path: str,
) -> "obspy.Stream":
    """Trim all traces to the common (overlapping) time window.

    Parameters
    ----------
    input_path : str
        Path to input MiniSEED file.
    output_path : str
        Path to save the aligned output.

    Returns
    -------
    obspy.Stream
        Trimmed stream with identical start/end times across traces.
    """
    st = obspy.read(input_path)
    logger.info("Read %d trace(s) from %s", len(st), input_path)

    starttime = max(tr.stats.starttime for tr in st)
    endtime = min(tr.stats.endtime for tr in st)

    if starttime >= endtime:
        raise ValueError(
            f"No overlapping time window: latest start {starttime} ≥ earliest end {endtime}"
        )

    logger.info("Common window: %s → %s", starttime, endtime)
    st2 = st.slice(starttime, endtime)
    st2.write(output_path, format="MSEED")
    logger.info("Aligned stream saved to %s", output_path)
    return st2


# ---------------------------------------------------------------------------
# Trim & splice (remove a time segment)
# ---------------------------------------------------------------------------

def trim_and_splice(
    input_path: str,
    output_path: str,
    cut_start: str,
    cut_end: str,
    merge_method: int = 1,
    fill_value: str = "interpolate",
) -> "obspy.Stream":
    """Remove a time segment from a stream and splice the remaining parts.

    Parameters
    ----------
    input_path : str
        Path to input MiniSEED file.
    output_path : str
        Path to save the trimmed & spliced output.
    cut_start : str
        ISO-8601 datetime string for start of the cut window.
    cut_end : str
        ISO-8601 datetime string for end of the cut window.

    Returns
    -------
    obspy.Stream
        The spliced stream.
    """
    st = obspy.read(input_path)
    logger.info("Read %d trace(s) from %s", len(st), input_path)

    # Auto-merge if > 3 components
    if len(st) > 3:
        logger.info("Auto-merging %d traces before trimming …", len(st))
        st.merge(fill_value=fill_value)

    d1 = UTCDateTime(cut_start)
    d2 = UTCDateTime(cut_end)
    di = st[0].stats.starttime
    df = st[0].stats.endtime

    logger.info("Cutting segment %s → %s", d1, d2)

    st1 = st.copy()
    st1.trim(di, d1)
    st2 = st.copy()
    st2.trim(d2, df)

    st_merged = st1 + st2
    st_merged.merge(method=merge_method, fill_value=fill_value)

    st_merged.write(output_path, format="MSEED")
    logger.info("Spliced stream saved to %s (%d traces)", output_path, len(st_merged))
    return st_merged
