"""
General utility functions which are not specific to one data type.
"""

from __future__ import annotations

import copy
import logging
import warnings
from collections.abc import Callable, Iterable
from contextlib import suppress
from functools import singledispatch
from pathlib import Path

import obspy
import obspy.core.event as ev
import pandas as pd
from obspy.core.event import Catalog, Event, ResourceIdentifier, WaveformStreamID

import obsplus
from obsplus.constants import (
    EVENT_NAME_STRUCTURE,
    EVENT_PATH_STRUCTURE,
    catalog_component,
    catalog_or_event,
    event_clientable_type,
)
from obsplus.exceptions import ValidationError
from obsplus.interfaces import EventClient
from obsplus.utils.bank import EVENT_EXT, _get_time_values
from obsplus.utils.misc import _get_path, yield_obj_parent_attr
from obsplus.utils.pd import get_seed_id_series
from obsplus.utils.time import _get_event_origin_time, to_utc


def duplicate_events(
    event: catalog_or_event, resource_generator: Callable | None = None
) -> catalog_or_event:
    """
    Duplicate an event.

    Creates a copy of the event, and assigns new resource IDs to all
    internal objects (ie picks, station_magnitudes, etc.). Will not
    assign new resource_id to those representing external resources
    (eg analyst, velocity models, etc.).

    Parameters
    ----------
    event
        The event to duplicate.
    resource_generator:
        A callable that can be used to generate resource IDs.
    """
    new = copy.deepcopy(event)
    gen_func = resource_generator or ResourceIdentifier
    id_map = {}  # map old id to new
    # first pass through bind resource id to parent if attr is resource_id
    for rid, parent, attr in yield_obj_parent_attr(new, ResourceIdentifier):
        if attr == "resource_id":
            try:
                prefix = rid.prefix
            except AttributeError:
                prefix = "smi:local"
            new_rid = gen_func(referred_object=parent, prefix=prefix)
            id_map[rid.id] = new_rid
            setattr(parent, attr, new_rid)
    # second pass, swap out resource ids with the same id as those in id_map
    for rid, parent, attr in yield_obj_parent_attr(new, ResourceIdentifier):
        if rid.id in id_map and attr != "resource_id":
            setattr(parent, attr, id_map[rid.id])
    return new


def prune_events(events: catalog_or_event) -> Catalog:
    """
    Remove all the unreferenced rejected objects from an event or catalog.

    This function first creates a copy of the event/catalog. Then it looks
    all objects with rejected evaluation status', which are not referred to
    by any un-rejected object, and removes them.

    Parameters
    ----------
    events
        The events to iterate and modify

    Returns
    -------
    A Catalog with the pruned events.

    """
    from obsplus.events.validate import validate_catalog

    # ensure we have something iterable
    events = [events] if isinstance(events, ev.Event) else events
    out = []

    def _get_edges_rids_opa(event):
        """
        Return a list of edges (resource_id_parent, resource_id_child),
        a set of rejected resource_ids and a dict of
        {resource_id: (obj, parent, attr)}.
        """
        edges = []  # list of tuples (parent, child)
        rejected_rid = set()  # {resource_id, }
        rejected_opa = {}  # {resource_id: [obj, parent, attr]}
        # first make a list objects with eval status as well as all other
        # objects they refer to, using resource_ids
        opa_iter = yield_obj_parent_attr(event, has_attr="evaluation_status")
        for obj, parent, attr in opa_iter:
            rid = str(obj.resource_id)
            is_rejected = obj.evaluation_status == "rejected"
            if is_rejected:
                rejected_rid.add(rid)
                rejected_opa[rid] = (obj, parent, attr)
            # now recurse object and find all objects this one refers to
            opa_iter_2 = yield_obj_parent_attr(obj, cls=ev.ResourceIdentifier)
            for obj2, _, _ in opa_iter_2:
                edges.append((rid, str(obj2)))
        return edges, rejected_rid, rejected_opa

    def _remove_object(obj, parent, attr):
        """
        Remove the object.
        """
        maybe_obj = getattr(parent, attr)
        # if maybe_obj is obj:
        #     breakpoint()
        #     setattr(parent, attr, None)
        # else:
        assert isinstance(maybe_obj, list)
        maybe_obj.remove(obj)

    # iterate events, find and destroy rejected orphans (that sounds bad...)
    for event in events:
        event = event.copy()
        validate_catalog(event)
        edges, rejected_rid, rejected_opa = _get_edges_rids_opa(event)
        df = pd.DataFrame(edges, columns=["parent", "child"])
        # filter out non rejected children
        df = df[df.child.isin(rejected_opa)]
        # iterate rejected children IDs and search for any non-rejected parent
        for rid, dff in df.groupby("child"):
            if dff.parent.isin(rejected_rid).all():
                _remove_object(*rejected_opa[rid])
        out.append(event)

    return Catalog(out)


@singledispatch
def strip_events(
    events: catalog_or_event, reject_evaluation_status: Iterable = "rejected"
) -> catalog_or_event:
    """
    Removes all derivative data and rejected objects from an event or catalog

    This is a nuclear option for when processing goes horribly wrong. It will
    only keep picks and amplitudes that are not rejected in addition to the
    first event description for the event.

    Parameters
    ----------
    events
        The events to strip
    reject_evaluation_status
        Reject picks and amplitudes that have this as an evaluation status
        (accepts either a single value or a list)

    Returns
    -------
    The stripped events
    """
    # Shouldn't get here
    raise TypeError("Must pass an obspy Catalog or obspy Event")


@strip_events.register(Catalog)
def _strip_catalog(cat, reject_evaluation_status="rejected") -> Catalog:
    """Act on a catalog of events"""
    # Make sure this returns a new catalog
    out = Catalog()
    for eve in cat:
        out.append(strip_events(eve, reject_evaluation_status=reject_evaluation_status))
    return out


@strip_events.register(Event)
def _strip_event(eve, reject_evaluation_status="rejected") -> Event:
    """Strip down a single event"""
    # Make sure this returns a copy of the events
    eve = eve.copy()
    # Remove derivative data
    for att in ["origins", "magnitudes", "station_magnitudes", "focal_mechanisms"]:
        setattr(eve, att, [])
    # Unset preferred anything
    for att in ["origin", "magnitude", "focal_mechanism"]:
        setattr(eve, f"preferred_{att}_id", None)
    # Filter the picks
    if isinstance(reject_evaluation_status, str):
        reject_evaluation_status = [reject_evaluation_status]
    eve.picks = [
        p for p in eve.picks if p.evaluation_status not in reject_evaluation_status
    ]
    # Filter the amplitudes
    amps = []
    for amp in eve.amplitudes:
        # Reject if the evaluation status is in the reject list
        if amp.evaluation_status in reject_evaluation_status:
            continue
        if amp.pick_id:
            pick = amp.pick_id.get_referred_object()
            # Reject if the evaluation status of the pick tied to the amplitude is in
            # the reject list
            if pick and pick.evaluation_status in reject_evaluation_status:
                continue
        amps.append(amp)
    eve.amplitudes = amps
    # Filter the event descriptions
    if len(eve.event_descriptions):
        eve.event_descriptions = [eve.event_descriptions[0]]
    return eve


def bump_creation_version(obj):
    """
    Bump the version in an object's CreationInfo and add creation time.

    Take an obspy object and try to bump creation info. This is done
    by first creating info if obj.creation_info is None, then addition
    creation_time and bumping the version by one
    """
    if not hasattr(obj, "creation_info"):  # nothing to do if no CI
        msg = f"{obj} has no creation_info attribute"
        logging.warning(msg)
        return
    # get creation info
    now = obspy.UTCDateTime.now()
    ci = obj.creation_info
    if ci is None:
        ci = obspy.core.event.CreationInfo()
    ci.creation_time = now
    ci.version = _bump_version(ci.version)
    obj.creation_info = ci


def _bump_version(version):
    """Bump the version in the creation info"""
    if isinstance(version, str):
        split = [int(x) for x in version.split(".")]
        split[-1] += 1
        version = ".".join([str(x) for x in split])
    elif version is None:
        version = "0.0.0"
    return version


def make_origins(
    events: catalog_or_event,
    inventory: obspy.Inventory,
    depth: float = 1.0,
    phase_hints: Iterable | None = ("P", "p"),
) -> catalog_or_event:
    """
    Iterate a catalog or single events and ensure each has an origin.

    If no origins are found for an event, create one with the time set to
    the earliest pick and the location set to the location of the first hit
    station. Events are modified in place.

    This may be useful for location codes that need a starting location.

    Parameters
    ----------
    events
        The events to scan and add origins were necessary.
    inventory
        An inventory object which contains all the stations referenced in
        quakeml elements of events.
    depth
        The default depth for created origins. Should be in meters. See the
        obspy docs for Origin or the quakeml standard for more details.
    phase_hints
        List of acceptable phase hints to use for identifying the earliest
        pick. By default will only search for "P" or "p" phase hints.

    Returns
    -------
    Either a Catalog or Event object (same as input).
    """
    # ensure input is an iterable of events
    cat = [events] if isinstance(events, Event) else events
    # load inv dataframe and make sure it has a seed_id column
    df = obsplus.stations_to_df(inventory)
    nslc_series = get_seed_id_series(df)
    for event in cat:
        if not event.origins:  # make new origin
            picks = event.picks_to_df()
            picks = picks.loc[
                (~(picks["evaluation_status"] == "rejected"))
                & (picks["phase_hint"].isin(phase_hints))
            ]
            if not len(picks):
                msg = f"{event} has no acceptable picks to create origin"
                raise ValidationError(msg)
            # get first pick, determine time/station used
            first_pick = picks.loc[picks["time"].idxmin()]
            seed_id = first_pick["seed_id"]
            # find channel corresponding to pick
            df_chan = df[nslc_series == seed_id]
            if not len(df_chan):
                raise ValidationError(f"{seed_id} not found in inventory")
            ser = df_chan.iloc[0]
            # create origin
            ori = _create_first_pick_origin(first_pick, ser, depth=depth)
            event.origins.append(ori)
    return events


def _create_first_pick_origin(first_pick, channel_ser, depth):
    """Create an origin based on first pick and a channel series."""
    msg = (
        "origin fixed to location and time of earliest pick by "
        f"obsplus version {obsplus.__last_version__}"
    )
    comment = ev.Comment(text=msg)
    odict = dict(
        time=to_utc(first_pick["time"]),
        latitude=channel_ser["latitude"],
        longitude=channel_ser["longitude"],
        depth=depth,
        time_fixed=True,
        comments=[comment],
    )
    return ev.Origin(**odict)


def get_seed_id(obj: catalog_component) -> str:
    """
    Get the NSLC associated with a station-specific object.

    Parameters
    ----------
    obj
        The object for which to retrieve the seed id. Can be anything that
        has a waveform_id attribute or refers to an object with a
        waveform_id attribute.

    Returns
    -------
    str :
        The seed_id in the form of "network.station.location.channel"
    """
    if isinstance(obj, WaveformStreamID):
        # Get the seed id
        return obj.get_seed_string()
    if isinstance(obj, ResourceIdentifier):
        # Get the next nested object
        return get_seed_id(obj.get_referred_object())
    # The order of this list matters! It should first try waveform_id and then
    # go from the shallowest nested attribute to the deepest
    attrs = ["waveform_id", "station_magnitude_id", "amplitude_id", "pick_id"]
    # Make sure the object has at least one of the required attributes
    if not len(set(attrs).intersection(set(vars(obj)))):
        raise TypeError(f"cannot retrieve seed id for objects of type {type(obj)}")
    # Loop over each of the attributes, if it exists and is not None,
    # go down a level until it finds a seed id
    for att in attrs:
        val = getattr(obj, att, None)
        if val is not None:
            with suppress((TypeError, AttributeError)):
                return get_seed_id(val)
    # If it makes it this far, it could not find a non-None attribute
    # raise assertion error so this still works in validators
    assert 0, f"Unable to fetch a seed id for {obj.resource_id}"


@singledispatch
def get_event_client(events: event_clientable_type) -> EventClient:
    """
    Extract an event client from various inputs.

    Parameters
    ----------
    events
        Any of the following:
            * A path to an obspy-readable event file
            * A path to a directory of obspy-readable event files
            * An `obspy.Catalog` instance
            * An instance of :class:`~obsplus.EventBank`
            * Any other object that has a `get_events` method

    Raises
    ------
    TypeError
        If an event client cannot be determined from the input.
    """
    if not isinstance(events, EventClient):
        msg = f"an event client could not be extracted from {events}"
        raise TypeError(msg)
    return events


@get_event_client.register(str)
@get_event_client.register(Path)
def _catalog_to_client(path):
    """
    Turn a str or directory into a client.

    If a single file, try to read as events, if directory init event bank.
    """
    path = Path(path)  # ensure we are working with a path
    assert path.exists()
    if path.is_dir():
        return get_event_client(obsplus.EventBank(path))
    else:
        return get_event_client(obspy.read_events(str(path)))


@get_event_client.register(ev.Event)
def _event_to_catalog(event):
    return get_event_client(obspy.Catalog(events=[event]))


def get_preferred(event: Event, what: str, init_empty=False):
    """
    get the preferred object (eg origin, magnitude) from the event.

    If not defined use the last in the list. If list is empty init empty
    object.

    Parameters
    ----------
    event: obspy.core.event.Event
        The instance for which the preferred should be sought.
    what: the preferred item to get
        Can either be "magnitude", "origin", or "focal_mechanism".
    init_empty
        If True, rather than return None when no preferred object is found
        create an empty object of the appropriate class and return it.
    """

    def _none_or_empty():
        """Return None or an empty object of correct type."""
        if init_empty:
            return getattr(obspy.core.event, what.capitalize())()
        else:
            return None

    pref_type = {
        "magnitude": ev.Magnitude,
        "origin": ev.Origin,
        "focal_mechanism": ev.FocalMechanism,
    }

    prefname = "preferred_" + what
    whats = what + "s"
    obj = getattr(event, prefname)()
    if obj is None:  # get end of list
        pid = getattr(event, prefname + "_id")
        if pid is None:  # no preferred id set, return last in list
            try:  # if there is None return an empty one
                obj = getattr(event, whats)[-1]
            except IndexError:  # object has no whats (eg magnitude)
                # TODO why not return None here?
                return _none_or_empty()
        else:  # there is an id, it has just come detached, try to find it
            potentials = {x.resource_id.id: x for x in getattr(event, whats)}
            obj = potentials.get(pid.id, None)
            if obj is None:  # it wasn't found in the potentials.
                var = (pid.id, whats, str(event))
                warnings.warn("cannot find {} in {} for event {}".format(*var))
                try:
                    obj = getattr(event, whats)[-1]
                except IndexError:  # there are no "whats", return None
                    return _none_or_empty()
    # make sure the correct output type is returned
    cls = pref_type[what]
    assert isinstance(obj, cls), f"{type(obj)} is not a {cls}, wrong type returned"
    return obj


def _summarize_event(
    event: ev.Event,
    path: str | None = None,
    name: str | None = None,
    path_struct: str | None = None,
    name_struct: str | None = None,
) -> dict:
    """
    Function to extract info from events for indexing.

    Parameters
    ----------
    event
        The event object
    path
        Other Parameters to the file
    name
        Name of the file
    path_struct
        directory structure to create
    name_struct
    """
    res_id = str(event.resource_id)
    out = {
        "ext": EVENT_EXT,
        "event_id": res_id,
        "event_id_short": res_id[-5:],
        "event_id_end": res_id.split("/")[-1],
    }
    t1 = _get_event_origin_time(event)
    out.update(_get_time_values(t1))
    path_struct = path_struct if path_struct is not None else EVENT_PATH_STRUCTURE
    name_struct = name_struct or EVENT_NAME_STRUCTURE
    out.update(_get_path(out, path, name, path_struct, name_struct))
    return out
