"""ON AIR artwork: real terminal images over an always-present ANSI mosaic."""

from __future__ import annotations

import base64
import contextlib
import itertools
import os
import random
import sys
from io import BytesIO
from typing import TYPE_CHECKING, cast

from rich.console import Group, RenderableType
from rich.control import Control
from rich.segment import ControlType, Segment
from rich.style import Style
from rich.text import Text
from textual.dom import NoScreen
from textual.geometry import Region
from textual.strip import Strip
from textual.widget import Widget

from .. import artwork
from .common import art_lines

_NULL_STYLE = Style()
_NULL_CONTROL = ((ControlType.CURSOR_FORWARD, 0),)

if TYPE_CHECKING:
    from ..app import BeefyPlayerApp


_IMAGE_IDS = itertools.count(random.randrange(1, 2**31 - 1))


def _tgp(data: str) -> str:
    return f"\x1b_G{data}\x1b\\"


def _iterm2_native() -> bool:
    """Are we talking to iTerm2 itself (rather than a compatible terminal)?

    WezTerm understands the same OSC 1337 payload but is fussier about
    parameters it does not know, so the iTerm2-only ones stay gated.
    """
    env = os.environ
    return env.get("LC_TERMINAL", "").lower() == "iterm2" or env.get(
        "TERM_PROGRAM", ""
    ).lower() in {"iterm.app", "iterm2"}


def _art_key(app) -> str:
    """Identity of what the mosaic is drawing (station + current track)."""
    station = app.current
    if station is None:
        return ""
    song = station.now_playing
    artist = song.artist if song is not None else ""
    title = song.title if song is not None else ""
    return f"{station.id}|{artist}|{title}"


def _terminal_wrap(sequence: str) -> str:
    """Wrap Kitty sequences for tmux DCS passthrough."""
    term = os.environ.get("TERM", "")
    if os.environ.get("TMUX") or term.startswith(("screen", "tmux")):
        return (
            "\x1bPtmux;\x1b"
            + sequence.replace("\x1b", "\x1b\x1b")
            + "\x1b\\"
        )
    return sequence


class KittyImage(Widget):
    """Direct Kitty graphics placement (no Unicode placeholders or probe)."""

    DEFAULT_CSS = """
    KittyImage {
        width: 100%;
        height: 100%;
    }
    """

    can_focus = False

    def __init__(self, image) -> None:
        super().__init__()
        self._image = image
        self._image_id: int | None = None
        self._delete_ids: list[int] = []
        self._placed_key: tuple[int, int, int] | None = None

    @property
    def image(self):
        return self._image

    @image.setter
    def image(self, value) -> None:
        if self._image_id is not None and value is not self._image:
            self._delete_ids.append(self._image_id)
            self._image_id = None
        self._image = value
        self._placed_key = None
        self.refresh(layout=True)

    def _sequences(self, width: int, height: int) -> list[str]:
        key = (id(self._image), width, height)
        if self._image_id is not None and self._placed_key == key:
            return []

        buffer = BytesIO()
        self._image.save(buffer, format="PNG")
        payload = base64.b64encode(buffer.getvalue()).decode("ascii")
        chunks = [payload[i : i + 4096] for i in range(0, len(payload), 4096)]

        commands: list[str] = []
        for old_id in self._delete_ids:
            commands.append(_terminal_wrap(_tgp(f"a=d,i={old_id},q=2")))
        self._delete_ids.clear()

        image_id = next(_IMAGE_IDS)
        for index, chunk in enumerate(chunks):
            more = 1 if index < len(chunks) - 1 else 0
            commands.append(
                _terminal_wrap(
                    _tgp(f"a=t,i={image_id},f=100,m={more},q=2;{chunk}")
                )
            )
        commands.append(
            _terminal_wrap(
                _tgp(f"a=p,i={image_id},c={width},r={height},C=1,q=2")
            )
        )
        self._image_id = image_id
        self._placed_key = key
        return commands

    def cleanup(self) -> None:
        ids = list(self._delete_ids)
        if self._image_id is not None:
            ids.append(self._image_id)
        self._image_id = None
        self._delete_ids.clear()
        self._placed_key = None
        stdout = sys.__stdout__
        if stdout is None:
            return
        for image_id in ids:
            with contextlib.suppress(Exception):
                stdout.write(
                    _terminal_wrap(_tgp(f"a=d,i={image_id},q=2"))
                )
        with contextlib.suppress(Exception):
            stdout.flush()

    def render_lines(self, crop: Region) -> list[Strip]:
        empty = [Strip([]) for _ in range(max(0, crop.height))]
        if self._image is None:
            return empty
        try:
            if not self.screen.is_active:
                return empty
        except NoScreen:
            return empty

        width = self.content_size.width
        height = self.content_size.height
        if width <= 0 or height <= 0:
            return empty

        commands = self._sequences(width, height)
        if not commands:
            return empty

        try:
            visible = self.screen.find_widget(self).visible_region
        except NoScreen:
            return empty

        segments: list[Segment] = [
            Segment(
                Control.move_to(visible.x, visible.y).segment.text,
                style=_NULL_STYLE,
            )
        ]
        for command in commands:
            segments.append(
                Segment(command, style=_NULL_STYLE, control=_NULL_CONTROL)
            )
        segments.append(
            Segment(
                Control.move_to(visible.right, visible.bottom).segment.text,
                style=_NULL_STYLE,
            )
        )
        last = Strip(segments, cell_length=max(1, crop.width))
        return [Strip([]) for _ in range(max(0, crop.height - 1))] + [last]


class ITerm2Image(Widget):
    """Inline image via OSC 1337, injected through Textual's render_lines().

    WezTerm and iTerm2 both implement this natively.  ``doNotMoveCursor=1``
    (iTerm2 only) keeps the terminal cursor sane; the widget then restores it
    to the bottom-right of the region.
    """

    DEFAULT_CSS = """
    ITerm2Image {
        width: 100%;
        height: 100%;
    }
    """

    can_focus = False

    def __init__(self, image) -> None:
        super().__init__()
        self._image = image
        self._cache_key: tuple[int, int, int] | None = None
        self._cache_data = ""

    @property
    def image(self):
        return self._image

    @image.setter
    def image(self, value) -> None:
        self._image = value
        self._cache_key = None
        self._cache_data = ""
        self.refresh(layout=True)

    def _sequence(self, width: int, height: int) -> str:
        key = (id(self._image), width, height)
        if key == self._cache_key:
            return self._cache_data
        buffer = BytesIO()
        self._image.save(buffer, format="PNG")
        payload = base64.b64encode(buffer.getvalue()).decode("ascii")
        # ``doNotMoveCursor`` keeps iTerm2 from scrolling/jumping the cursor
        # out from under Textual's compositor; we place the cursor ourselves.
        keep_cursor = "doNotMoveCursor=1;" if _iterm2_native() else ""
        data = (
            "\x1b]1337;File=inline=1;"
            f"{keep_cursor}"
            f"width={width};height={height};"
            f"preserveAspectRatio=0:{payload}\x07"
        )
        self._cache_key = key
        self._cache_data = data
        return data

    def render_lines(self, crop: Region) -> list[Strip]:
        empty = [Strip([]) for _ in range(max(0, crop.height))]
        if self._image is None:
            return empty
        try:
            if not self.screen.is_active:
                return empty
        except NoScreen:
            return empty

        width = self.content_size.width
        height = self.content_size.height
        if width <= 0 or height <= 0:
            return empty

        try:
            visible = self.screen.find_widget(self).visible_region
        except NoScreen:
            return empty

        segments = [
            Segment(
                Control.move_to(visible.x, visible.y).segment.text,
                style=_NULL_STYLE,
            ),
            Segment(
                self._sequence(width, height),
                style=_NULL_STYLE,
                control=_NULL_CONTROL,
            ),
            Segment(
                Control.move_to(visible.right, visible.bottom).segment.text,
                style=_NULL_STYLE,
            ),
        ]
        last = Strip(segments, cell_length=max(1, crop.width))
        return [Strip([]) for _ in range(max(0, crop.height - 1))] + [last]


class ArtPane(Widget):
    """Mosaic background plus an optional protocol-backed image overlay.

    The mosaic is always rendered by this widget.  If the terminal image
    protocol silently fails (or an image is transparent), the artwork area is
    still not blank.
    """

    can_focus = False

    def __init__(self) -> None:
        super().__init__()
        self._image_widget: Widget | None = None
        self._last_data: bytes | None = None
        self._last_size: tuple[int, int] = (0, 0)
        self._mosaic_key = ""

    # ------------------------------------------------------------------ #
    # mosaic background
    # ------------------------------------------------------------------ #

    def render(self) -> RenderableType:
        app = cast("BeefyPlayerApp", self.app)
        station = app.current
        if station is None:
            return Text("")
        width = max(1, self.size.width)
        height = max(1, self.size.height)
        return Group(*art_lines(_art_key(app), width, height))

    # ------------------------------------------------------------------ #
    # image overlay
    # ------------------------------------------------------------------ #

    def on_mount(self) -> None:
        self.sync_image()

    def on_resize(self, event) -> None:
        self.refresh()
        if self.display:
            self.sync_image()

    def on_unmount(self) -> None:
        self._drop_image_widget()

    def sync_image(self, *, force: bool = False) -> None:
        # Keep the mosaic behind the image in sync with the current track.
        self.refresh()

        app = cast("BeefyPlayerApp", self.app)
        station = app.current
        data = app.artwork_bytes
        width = max(1, self.size.width)
        height = max(1, self.size.height)

        # Repainting the mosaic overwrites the cells an inline (iTerm2) image
        # occupies, so re-emit the image whenever what is drawn underneath
        # changes -- even if the artwork bytes happen to be the same.
        mosaic_key = _art_key(app)
        if mosaic_key != self._mosaic_key:
            self._mosaic_key = mosaic_key
            force = True

        if station is not None and data and artwork.supported():
            if (
                not force
                and self._image_widget is not None
                and self._last_data is data
                and self._last_size == (width, height)
            ):
                return
            fitted = artwork.fit_image(data, width, height)
            if fitted is not None:
                self._show_image(fitted, data, width, height)
                return
        self._show_fallback()

    def _show_image(self, image, data: bytes, width: int, height: int) -> None:
        if self._image_widget is None:
            if artwork.PROTOCOL == "tgp":
                widget_class = KittyImage
            elif artwork.PROTOCOL == "iterm2":
                widget_class = ITerm2Image
            else:
                self._show_fallback()
                return
            try:
                widget = widget_class(image)
            except Exception:
                self._show_fallback()
                return
            widget.add_class("art-image")
            self._image_widget = widget
            self.mount(widget)
        else:
            self._image_widget.image = image

        self._last_data = data
        self._last_size = (width, height)

    def _show_fallback(self) -> None:
        self._drop_image_widget()

    def _drop_image_widget(self) -> None:
        widget = self._image_widget
        self._image_widget = None
        self._last_data = None
        self._last_size = (0, 0)
        if widget is None:
            return
        with contextlib.suppress(Exception):
            widget.cleanup()
        with contextlib.suppress(Exception):
            widget.image = None
        widget.remove()