# piker: trading gear for hackers
# Copyright (C) Tyler Goodlet (in stewardship for pikers)

# This program is free software: you can redistribute it and/or
# modify it under the terms of the GNU Affero General Public License
# as published by the Free Software Foundation, either version 3
# of the License, or (at your option) any later version.

# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.

# You should have received a copy of the GNU Affero General Public
# License along with this program. If not, see
# <https://www.gnu.org/licenses/>.

'''
Chart-local gap-overlay detection, rendering and IPC messages.

'''
from __future__ import annotations
from typing import (
    Literal,
    TYPE_CHECKING,
)

import numpy as np

from piker.types import Struct
from ._annotate import GapAnnotations

if TYPE_CHECKING:
    from PyQt6.QtWidgets import QGraphicsScene
    from pyqtgraph import PlotItem

    from piker.ui.qt import QGraphicsItem
    from ._chart import ChartPlotWidget
    from ._dataviz import Viz
    from ._display import DisplayState


# `DisplayState` currently exposes exactly one realtime and one
# historical OHLC chart at these sampling periods.
RT_GAP_TIMEFRAME: float = 1.
HIST_GAP_TIMEFRAME: float = 60.


class GapSpec(Struct, frozen=True):
    '''
    Render-ready geometry for one OHLC time gap.

    '''
    start_pos: tuple[float, float]
    end_pos: tuple[float, float]
    arrow_x: float
    arrow_y: float
    pointing: Literal['up', 'down']
    start_time: float
    end_time: float
    time: float
    color: str = 'dad_blue'
    alpha: int = 169


class SetGapOverlay(Struct, frozen=True, tag=True):
    '''
    Request replacement of one named gap-overlay layer.

    `GapOverlayMngr.apply()` accepts this type locally and through
    the chart actor's typed IPC endpoint. The receiver assigns
    ownership so remote teardown cannot remove chart-local overlays.

    '''
    fqme: str
    timeframe: float
    specs: list[GapSpec]
    visible: bool = True
    request_id: str = ''


class GapOverlay(Struct, frozen=True, tag=True):
    '''
    Gap-overlay state returned after applying `SetGapOverlay`.

    '''
    fqme: str
    timeframe: float
    visible: bool
    gap_count: int
    aid: int|None = None
    error: str|None = None
    request_id: str = ''


GapOverlayPayload = (
    SetGapOverlay
    | GapOverlay
    | list[str]
    | None
)


def gap_specs_from_ohlcv(
    array: np.ndarray,
    period_s: float,

) -> list[GapSpec]:
    '''
    Build render specs for positive timestamp gaps in sorted OHLCV.

    Invalid and non-positive timestamps are excluded. Timestamp steps
    are detected in one vectorized pass; only actual gaps enter the
    small Python object-construction loop.

    '''
    if len(array) < 2:
        return []

    times: np.ndarray = array['time'].astype(float, copy=False)
    valid: np.ndarray = (
        np.isfinite(times)
        & (times > 0)
    )
    rows: np.ndarray = array[valid]
    if len(rows) < 2:
        return []

    times = rows['time'].astype(float, copy=False)
    gap_ends: np.ndarray = np.flatnonzero(
        np.diff(times) > period_s
    ) + 1
    if not len(gap_ends):
        return []

    gap_starts: np.ndarray = gap_ends - 1
    left: np.ndarray = rows[gap_starts]
    right: np.ndarray = rows[gap_ends]

    start_xs: np.ndarray = (
        left['index'].astype(float, copy=False)
        + 0.9
    )
    end_xs: np.ndarray = (
        right['index'].astype(float, copy=False)
        + 0.1
    )
    start_ys: np.ndarray = left['close'].astype(
        float,
        copy=False,
    )
    end_ys: np.ndarray = right['open'].astype(
        float,
        copy=False,
    )

    specs: list[GapSpec] = []
    start_x: np.floating
    end_x: np.floating
    start_y: np.floating
    end_y: np.floating
    start_time: np.floating
    end_time: np.floating
    for (
        start_x,
        end_x,
        start_y,
        end_y,
        start_time,
        end_time,
    ) in zip(
        start_xs,
        end_xs,
        start_ys,
        end_ys,
        left['time'],
        right['time'],
        strict=True,
    ):
        specs.append(GapSpec(
            start_pos=(float(start_x), float(start_y)),
            end_pos=(float(end_x), float(end_y)),
            arrow_x=float(end_x - 0.1),
            arrow_y=float(start_y),
            pointing=(
                'down'
                if start_y > end_y
                else 'up'
            ),
            start_time=float(start_time),
            end_time=float(end_time),
            time=float(end_time),
        ))

    return specs


class GapOverlayMngr:
    '''
    Own named gap-overlay layers inside one chart actor.

    Local UI calls and remote annotation IPC both dispatch typed
    `SetGapOverlay` messages through `GapOverlayMngr.apply()`.

    '''
    def __init__(
        self,
        dss: dict[str, DisplayState],
        annots: dict[int, QGraphicsItem],

    ) -> None:
        self.dss: dict[str, DisplayState] = dss
        self.annots: dict[int, QGraphicsItem] = annots
        self._layers: dict[
            tuple[  # owner-scoped layer key
                str,  # owner (`chart-local` or `Context.cid`)
                int,  # `id(ChartPlotWidget)`
                str,  # fqme
                float,  # timeframe
            ],
            int,  # `id(GapAnnotations)`
        ] = {}
        self._visible: dict[
            tuple[  # actor-local `SetGapOverlay` key
                int,  # `id(ChartPlotWidget)`
                str,  # fqme
                float,  # timeframe
            ],
            bool,  # `SetGapOverlay.visible`
        ] = {}
        self._plots: dict[
            int,  # `id(GapAnnotations)`
            PlotItem,  # owning `Viz.plot`
        ] = {}

    def _remove_annot(
        self,
        aid: int,

    ) -> None:
        '''
        Remove one gap item from all manager and plot registries.

        '''
        annot: QGraphicsItem|None = self.annots.pop(
            aid,
            None,
        )
        plot: PlotItem|None = self._plots.pop(aid, None)
        if annot is None:
            return

        if plot is not None:
            plot.removeItem(annot)
            return

        # Defensive fallback for manager state created before the
        # owning `Viz.plot` registry was introduced.
        scene: QGraphicsScene|None = annot.scene()
        if scene:
            scene.removeItem(annot)

    def apply(
        self,
        req: SetGapOverlay,
        owner: str = 'chart-local',
        aids: set[int]|None = None,
        ds: DisplayState|None = None,

    ) -> GapOverlay:
        '''
        Replace one owner-scoped batch layer from a typed request.

        '''
        if ds is None:
            ds = self.dss.get(req.fqme)
            if ds is None:
                return GapOverlay(
                    fqme=req.fqme,
                    timeframe=req.timeframe,
                    visible=False,
                    gap_count=0,
                    request_id=req.request_id,
                    error=(
                        f'No display state for fqme={req.fqme}'
                    ),
                )
        try:
            chart: ChartPlotWidget = {
                HIST_GAP_TIMEFRAME: ds.hist_chart,
                RT_GAP_TIMEFRAME: ds.chart,
            }[req.timeframe]
        except KeyError:
            return GapOverlay(
                fqme=req.fqme,
                timeframe=req.timeframe,
                visible=False,
                gap_count=0,
                request_id=req.request_id,
                error=(
                    f'No chart for timeframe={req.timeframe}s'
                ),
            )

        layer_key: tuple[str, int, str, float] = (
            owner,
            id(chart),
            req.fqme,
            req.timeframe,
        )
        aid: int|None = self._layers.pop(layer_key, None)
        if aid is not None:
            if aids is not None:
                aids.discard(aid)
            self._remove_annot(aid)

        if owner == 'chart-local':
            self._visible[
                (id(chart), req.fqme, req.timeframe)
            ] = req.visible
        if (
            not req.visible
            or not req.specs
        ):
            return GapOverlay(
                fqme=req.fqme,
                timeframe=req.timeframe,
                visible=req.visible,
                gap_count=len(req.specs),
                request_id=req.request_id,
            )

        viz: Viz = chart.get_viz(req.fqme)
        gap_specs: list[dict] = []
        spec: GapSpec
        for spec in req.specs:
            gap_specs.append(spec.to_dict())
        gaps_item: GapAnnotations = GapAnnotations(
            gap_specs=gap_specs,
            array=viz.shm.array,
            color=req.specs[0].color,
            alpha=req.specs[0].alpha,
            arrow_size=10,
            fqme=req.fqme,
            timeframe=req.timeframe,
        )
        viz.plot.addItem(
            gaps_item,
            ignoreBounds=True,
        )
        gaps_item._chart = chart

        new_aid: int = id(gaps_item)
        self.annots[new_aid] = gaps_item
        self._plots[new_aid] = viz.plot
        self._layers[layer_key] = new_aid
        if aids is not None:
            aids.add(new_aid)
        return GapOverlay(
            fqme=req.fqme,
            timeframe=req.timeframe,
            visible=True,
            gap_count=len(req.specs),
            aid=new_aid,
            request_id=req.request_id,
        )

    def refresh(
        self,
        ds: DisplayState,
        timeframe: float,
        visible: bool = True,

    ) -> GapOverlay:
        '''
        Detect gaps from chart-local SHM and replace the local layer.

        '''
        viz: Viz = {
            HIST_GAP_TIMEFRAME: ds.hist_viz,
            RT_GAP_TIMEFRAME: ds.viz,
        }[timeframe]
        specs: list[GapSpec] = gap_specs_from_ohlcv(
            viz.shm.array,
            period_s=timeframe,
        )
        return self.apply(SetGapOverlay(
            fqme=ds.fqme,
            timeframe=timeframe,
            specs=specs,
            visible=visible,
        ), ds=ds)

    def toggle(
        self,
        ds: DisplayState,
        timeframe: float,

    ) -> GapOverlay:
        '''
        Toggle and refresh one chart-local gap-overlay layer.

        '''
        chart: ChartPlotWidget = {
            HIST_GAP_TIMEFRAME: ds.hist_chart,
            RT_GAP_TIMEFRAME: ds.chart,
        }[timeframe]
        visible: bool = not self._visible.get(
            (id(chart), ds.fqme, timeframe),
            False,
        )
        return self.refresh(
            ds=ds,
            timeframe=timeframe,
            visible=visible,
        )

    def remove_owner(
        self,
        owner: str,

    ) -> None:
        '''
        Remove all graphics and bookkeeping for one IPC owner.

        '''
        owner_keys: list[tuple[str, int, str, float]] = []
        key: tuple[str, int, str, float]
        for key in self._layers:
            if key[0] == owner:
                owner_keys.append(key)

        layer_key: tuple[str, int, str, float]
        for layer_key in owner_keys:
            aid: int = self._layers.pop(layer_key)
            self._remove_annot(aid)
