#!/usr/bin/env python3
#
# Copyright (c) 2026 Stormux
#
# This library is free software; you can redistribute it and/or
# modify it under the terms of the GNU Lesser General Public
# License as published by the Free Software Foundation; either
# version 2.1 of the License, or (at your option) any later version.
#
# Forked from Orca screen reader.
# Cthulhu project: https://git.stormux.org/storm/cthulhu

"""Helpers for document-wide AT-SPI text selections."""

from __future__ import annotations

import functools
from dataclasses import dataclass, field
from typing import Any, Optional, TYPE_CHECKING

import gi

gi.require_version("Atspi", "2.0")
from gi.repository import Atspi
from gi.repository import GLib

from . import debug
from .ax_object import AXObject
from .ax_text import AXText

if TYPE_CHECKING:
    from .script_utilities import Utilities


@dataclass
class AXSelectionRange:
    start_object: Atspi.Accessible
    start_offset: int
    end_object: Atspi.Accessible
    end_offset: int
    start_is_active: bool = False

    def anchor_endpoint(self) -> tuple[Any, int]:
        if self.start_is_active:
            return self.end_object, self.end_offset
        return self.start_object, self.start_offset

    def focus_endpoint(self) -> tuple[Any, int]:
        if self.start_is_active:
            return self.start_object, self.start_offset
        return self.end_object, self.end_offset

    def active_endpoint(self) -> tuple[Any, int]:
        if self.start_is_active:
            return self.start_object, self.start_offset
        return self.end_object, self.end_offset


@dataclass
class AXSelectionSnapshot:
    document: Optional[Atspi.Accessible]
    source_object: Optional[Atspi.Accessible]
    ranges: list[AXSelectionRange] = field(default_factory=list)
    anchor_focus_pairs: list[tuple[Any, int, Any, int]] = field(default_factory=list)
    entire_document_selected: bool = False
    is_authoritative: bool = False


@dataclass
class AXProjectedSelection:
    obj: Optional[Atspi.Accessible]
    local_ranges: list[tuple[int, int]] = field(default_factory=list)
    text: str = ""
    start_offset: int = 0
    end_offset: int = 0


class AXDocumentSelection:
    """Normalizes and applies document-wide text selections."""

    def __init__(self, utilities: "Utilities") -> None:
        self._utilities = utilities

    def empty_snapshot(
        self,
        source_object: Optional[Atspi.Accessible] = None,
        document: Optional[Atspi.Accessible] = None,
    ) -> AXSelectionSnapshot:
        return AXSelectionSnapshot(document=document, source_object=source_object)

    def get_cached_snapshot(self, obj: Optional[Atspi.Accessible]) -> AXSelectionSnapshot:
        document = self._utilities.getDocumentForObject(obj)
        state = self._utilities._script.pointOfReference.get("selectionState", {})
        snapshot = state.get("snapshot")
        if snapshot is None:
            return self.empty_snapshot(obj, document)
        if state.get("document") != document:
            return self.empty_snapshot(obj, document)
        return snapshot

    def refresh_snapshot(self, obj: Optional[Atspi.Accessible]) -> AXSelectionSnapshot:
        document = self._utilities.getDocumentForObject(obj)
        snapshot = self._fetch_snapshot(document, obj)
        state = self._utilities._script.pointOfReference.get("selectionState", {})
        if state.get("document") != document:
            state = {}
        state["document"] = document
        state["snapshot"] = snapshot
        state.setdefault("textSelections", {})
        self._utilities._script.pointOfReference["selectionState"] = state
        if snapshot.is_authoritative:
            self._utilities._script.pointOfReference["entireDocumentSelected"] = snapshot.entire_document_selected
        return snapshot

    def get_selected_objects(self, snapshot: AXSelectionSnapshot) -> list[Atspi.Accessible]:
        seen = set()
        objects = []
        for selectionRange in snapshot.ranges:
            for obj in self._objects_for_range(selectionRange):
                if hash(obj) in seen:
                    continue
                seen.add(hash(obj))
                objects.append(obj)
        return objects

    def get_projected_selection(
        self,
        obj: Optional[Atspi.Accessible],
        snapshot: Optional[AXSelectionSnapshot] = None,
    ) -> AXProjectedSelection:
        if not obj:
            return AXProjectedSelection(obj)

        snapshot = snapshot or self.get_cached_snapshot(obj)
        if not snapshot.is_authoritative:
            return AXProjectedSelection(obj)

        localRanges = []
        for selectionRange in snapshot.ranges:
            localRanges.extend(self._project_range_onto_object(selectionRange, obj))
        localRanges = self._merge_ranges(localRanges)

        fragments = []
        for startOffset, endOffset in localRanges:
            text = self._utilities.expandEOCs(obj, startOffset, endOffset)
            if text:
                fragments.append(text)

        if localRanges:
            startOffset = localRanges[0][0]
            endOffset = localRanges[-1][1]
        else:
            startOffset = endOffset = 0

        return AXProjectedSelection(
            obj=obj,
            local_ranges=localRanges,
            text=" ".join(fragments),
            start_offset=startOffset,
            end_offset=endOffset,
        )

    def get_all_selected_text(
        self,
        obj: Optional[Atspi.Accessible],
        snapshot: Optional[AXSelectionSnapshot] = None,
    ) -> tuple[str, int, int]:
        if not obj:
            return "", 0, 0

        snapshot = snapshot or self.get_cached_snapshot(obj)
        projected = self.get_projected_selection(obj, snapshot)
        if not snapshot.is_authoritative or not snapshot.ranges:
            return projected.text, projected.start_offset, projected.end_offset

        fragments = []
        for selectionRange in snapshot.ranges:
            text = self._range_text(selectionRange)
            if text:
                fragments.append(text)

        return " ".join(fragments), projected.start_offset, projected.end_offset

    def clear_selection(self, obj: Optional[Atspi.Accessible]) -> bool:
        if not obj:
            return False

        snapshot = self.refresh_snapshot(obj)
        if not snapshot.is_authoritative:
            return False

        projected = self.get_projected_selection(obj, snapshot)
        if not projected.local_ranges and not snapshot.entire_document_selected:
            return False

        return self._set_document_ranges(snapshot.document, [])

    def adjust_selection(self, obj: Optional[Atspi.Accessible], offset: int) -> bool:
        if not obj:
            return False

        snapshot = self.refresh_snapshot(obj)
        if not snapshot.is_authoritative:
            return False

        ranges = list(snapshot.ranges)
        anchorOffset = AXText.get_caret_offset(obj)
        rangeIndex = self._resolve_primary_range_index(ranges, obj, anchorOffset)
        if rangeIndex is None:
            rangeIndex = len(ranges)
            ranges.append(
                AXSelectionRange(
                    start_object=obj,
                    start_offset=anchorOffset,
                    end_object=obj,
                    end_offset=anchorOffset,
                    start_is_active=False,
                )
            )

        selectionRange = ranges[rangeIndex]
        anchorObject, anchorOffset = selectionRange.anchor_endpoint()
        updatedRange = self._build_range(anchorObject, anchorOffset, obj, offset)
        if updatedRange is None:
            return False

        ranges[rangeIndex] = updatedRange
        return self._set_document_ranges(snapshot.document, ranges)

    def _fetch_snapshot(
        self,
        document: Optional[Atspi.Accessible],
        source_object: Optional[Atspi.Accessible],
    ) -> AXSelectionSnapshot:
        if not document:
            return self.empty_snapshot(source_object, document)

        try:
            rawSelections = Atspi.Document.get_text_selections(document)
        except (AttributeError, GLib.GError, TypeError) as error:
            msg = f"AXDocumentSelection: Exception in get_text_selections: {error}"
            debug.printMessage(debug.LEVEL_INFO, msg, True)
            return self.empty_snapshot(source_object, document)

        ranges = []
        for rawSelection in rawSelections or []:
            selectionRange = self._normalize_range(document, rawSelection)
            if selectionRange is None:
                return self.empty_snapshot(source_object, document)
            ranges.append(selectionRange)

        ranges.sort(key=functools.cmp_to_key(self._compare_ranges))
        return AXSelectionSnapshot(
            document=document,
            source_object=source_object,
            ranges=ranges,
            anchor_focus_pairs=[
                (
                    selectionRange.anchor_endpoint()[0],
                    selectionRange.anchor_endpoint()[1],
                    selectionRange.focus_endpoint()[0],
                    selectionRange.focus_endpoint()[1],
                )
                for selectionRange in ranges
            ],
            entire_document_selected=self._is_entire_document_selected(ranges),
            is_authoritative=bool(ranges),
        )

    def _normalize_range(
        self,
        document: Atspi.Accessible,
        rawSelection: Any,
    ) -> Optional[AXSelectionRange]:
        startObject = getattr(rawSelection, "start_object", None)
        endObject = getattr(rawSelection, "end_object", None)
        startOffset = getattr(rawSelection, "start_offset", 0)
        endOffset = getattr(rawSelection, "end_offset", 0)
        startIsActive = bool(getattr(rawSelection, "start_is_active", False))

        if not startObject or not endObject:
            return None
        if not AXObject.supports_text(startObject) or not AXObject.supports_text(endObject):
            return None
        if self._utilities.getDocumentForObject(startObject) != document:
            return None
        if self._utilities.getDocumentForObject(endObject) != document:
            return None

        comparison = self._compare_positions(startObject, startOffset, endObject, endOffset)
        if comparison is None:
            return None
        if comparison > 0:
            startObject, endObject = endObject, startObject
            startOffset, endOffset = endOffset, startOffset
            startIsActive = not startIsActive

        if startObject == endObject and startOffset > endOffset:
            startOffset, endOffset = endOffset, startOffset
            startIsActive = not startIsActive

        if startObject == endObject and startOffset == endOffset:
            return None

        return AXSelectionRange(
            start_object=startObject,
            start_offset=startOffset,
            end_object=endObject,
            end_offset=endOffset,
            start_is_active=startIsActive,
        )

    def _is_entire_document_selected(self, ranges: list[AXSelectionRange]) -> bool:
        if len(ranges) != 1:
            return False

        selectionRange = ranges[0]
        if selectionRange.start_offset != 0:
            return False
        if selectionRange.end_offset != AXText.get_character_count(selectionRange.end_object):
            return False

        if self._find_previous_text_object(selectionRange.start_object) is not None:
            return False
        if self._find_next_text_object(selectionRange.end_object) is not None:
            return False

        return True

    def _find_previous_text_object(self, obj: Atspi.Accessible) -> Optional[Atspi.Accessible]:
        current = obj
        while current:
            current = self._utilities.findPreviousObject(current)
            if current is None:
                return None
            text = self._utilities.queryNonEmptyText(current)
            if text:
                return text
        return None

    def _find_next_text_object(self, obj: Atspi.Accessible) -> Optional[Atspi.Accessible]:
        current = obj
        while current:
            current = self._utilities.findNextObject(current)
            if current is None:
                return None
            text = self._utilities.queryNonEmptyText(current)
            if text:
                return text
        return None

    def _compare_ranges(self, range1: AXSelectionRange, range2: AXSelectionRange) -> int:
        comparison = self._compare_positions(
            range1.start_object,
            range1.start_offset,
            range2.start_object,
            range2.start_offset,
        )
        if comparison is None:
            return 0
        return comparison

    def _compare_positions(
        self,
        obj1: Atspi.Accessible,
        offset1: int,
        obj2: Atspi.Accessible,
        offset2: int,
    ) -> Optional[int]:
        if obj1 == obj2:
            if offset1 < offset2:
                return -1
            if offset1 > offset2:
                return 1
            return 0

        seen = {hash(obj1)}
        current = obj1
        while current:
            current = self._utilities.findNextObject(current)
            if current is None:
                break
            if hash(current) in seen:
                break
            if current == obj2:
                return -1
            seen.add(hash(current))

        seen = {hash(obj2)}
        current = obj2
        while current:
            current = self._utilities.findNextObject(current)
            if current is None:
                break
            if hash(current) in seen:
                break
            if current == obj1:
                return 1
            seen.add(hash(current))

        return None

    def _objects_for_range(self, selectionRange: AXSelectionRange) -> list[Atspi.Accessible]:
        objects = [selectionRange.start_object]
        if selectionRange.start_object == selectionRange.end_object:
            return objects

        seen = {hash(selectionRange.start_object)}
        current = selectionRange.start_object
        while current and current != selectionRange.end_object:
            current = self._utilities.findNextObject(current)
            if current is None or hash(current) in seen:
                return []
            seen.add(hash(current))
            if self._utilities.queryNonEmptyText(current):
                objects.append(current)

        if not objects or objects[-1] != selectionRange.end_object:
            return []

        return objects

    def _project_range_onto_object(
        self,
        selectionRange: AXSelectionRange,
        obj: Atspi.Accessible,
    ) -> list[tuple[int, int]]:
        rangeObjects = self._objects_for_range(selectionRange)
        if obj not in rangeObjects:
            return []

        if selectionRange.start_object == selectionRange.end_object == obj:
            startOffset = min(selectionRange.start_offset, selectionRange.end_offset)
            endOffset = max(selectionRange.start_offset, selectionRange.end_offset)
            if startOffset < endOffset:
                return [(startOffset, endOffset)]
            return []

        if obj == selectionRange.start_object:
            startOffset = selectionRange.start_offset
            endOffset = AXText.get_character_count(obj)
            return [(startOffset, endOffset)] if startOffset < endOffset else []

        if obj == selectionRange.end_object:
            startOffset = 0
            endOffset = selectionRange.end_offset
            return [(startOffset, endOffset)] if startOffset < endOffset else []

        endOffset = AXText.get_character_count(obj)
        return [(0, endOffset)] if endOffset else []

    @staticmethod
    def _merge_ranges(ranges: list[tuple[int, int]]) -> list[tuple[int, int]]:
        if not ranges:
            return []

        merged = []
        for startOffset, endOffset in sorted(ranges):
            if not merged or startOffset > merged[-1][1]:
                merged.append((startOffset, endOffset))
                continue
            merged[-1] = (merged[-1][0], max(merged[-1][1], endOffset))

        return merged

    def _range_text(self, selectionRange: AXSelectionRange) -> str:
        fragments = []
        for obj in self._objects_for_range(selectionRange):
            if obj == selectionRange.start_object == selectionRange.end_object:
                startOffset = selectionRange.start_offset
                endOffset = selectionRange.end_offset
            elif obj == selectionRange.start_object:
                startOffset = selectionRange.start_offset
                endOffset = AXText.get_character_count(obj)
            elif obj == selectionRange.end_object:
                startOffset = 0
                endOffset = selectionRange.end_offset
            else:
                startOffset = 0
                endOffset = AXText.get_character_count(obj)

            if startOffset >= endOffset:
                continue

            text = self._utilities.expandEOCs(obj, startOffset, endOffset)
            if text:
                fragments.append(text)

        return " ".join(fragments)

    def _resolve_primary_range_index(
        self,
        ranges: list[AXSelectionRange],
        obj: Atspi.Accessible,
        caretOffset: int,
    ) -> Optional[int]:
        containing = []
        projected = []
        active = []
        for index, selectionRange in enumerate(ranges):
            localRanges = self._project_range_onto_object(selectionRange, obj)
            if localRanges:
                projected.append(index)
                for startOffset, endOffset in localRanges:
                    if startOffset <= caretOffset <= endOffset:
                        containing.append(index)
                        break
            if selectionRange.active_endpoint()[0] == obj:
                active.append(index)

        if len(containing) == 1:
            return containing[0]
        if len(projected) == 1:
            return projected[0]
        if len(active) == 1:
            return active[0]
        return None

    def _build_range(
        self,
        anchorObject: Atspi.Accessible,
        anchorOffset: int,
        focusObject: Atspi.Accessible,
        focusOffset: int,
    ) -> Optional[AXSelectionRange]:
        comparison = self._compare_positions(anchorObject, anchorOffset, focusObject, focusOffset)
        if comparison is None:
            return None
        if comparison <= 0:
            return AXSelectionRange(
                start_object=anchorObject,
                start_offset=anchorOffset,
                end_object=focusObject,
                end_offset=focusOffset,
                start_is_active=False,
            )

        return AXSelectionRange(
            start_object=focusObject,
            start_offset=focusOffset,
            end_object=anchorObject,
            end_offset=anchorOffset,
            start_is_active=True,
        )

    def _set_document_ranges(
        self,
        document: Optional[Atspi.Accessible],
        ranges: list[AXSelectionRange],
    ) -> bool:
        if not document:
            return False

        selections = []
        for selectionRange in ranges:
            selection = Atspi.TextSelection()
            selection.start_object = selectionRange.start_object
            selection.start_offset = selectionRange.start_offset
            selection.end_object = selectionRange.end_object
            selection.end_offset = selectionRange.end_offset
            selection.start_is_active = selectionRange.start_is_active
            selections.append(selection)

        try:
            result = Atspi.Document.set_text_selections(document, selections)
        except (AttributeError, GLib.GError, TypeError) as error:
            msg = f"AXDocumentSelection: Exception in set_text_selections: {error}"
            debug.printMessage(debug.LEVEL_INFO, msg, True)
            return False

        return bool(result)
