"""
XML-preserving ELAN .eaf round-trip editing.

Unlike ``write_eaf`` which generates fresh XML (dropping controlled
vocabularies, linguistic types, referencing annotations, etc.), this
module wraps the parsed ``ElementTree`` and mutates it in-place so
that every element the original file contained is preserved through
read-modify-write cycles.

Usage::

    doc = ElanDocument.open("annotations.eaf")
    doc.add_segment("Gloss", 1.5, 3.0, "HELLO")
    doc.save()
"""

from __future__ import annotations

import logging
import os
import tempfile
from dataclasses import dataclass
from datetime import datetime
from pathlib import Path
from xml.etree import ElementTree as ET

from sltk.exceptions import SLTKFileNotFoundError

_logger = logging.getLogger(__name__)

# Safe XML parsing — same pattern as elan.py
try:
    from defusedxml import ElementTree as SafeET
except ImportError:
    _logger.warning(
        "defusedxml is not installed. XML parsing will use stdlib "
        "xml.etree.ElementTree, which is vulnerable to XML bomb attacks. "
        "Install defusedxml: pip install defusedxml"
    )
    SafeET = ET  # type: ignore[assignment,misc]


# ---------------------------------------------------------------------------
# Return-type dataclasses
# ---------------------------------------------------------------------------


@dataclass(frozen=True)
class TierInfo:
    """Metadata about a single ELAN tier."""

    name: str
    linguistic_type: str
    participant: str
    annotator: str
    segment_count: int


@dataclass(frozen=True)
class SegmentInfo:
    """A single annotation with its tier context."""

    annotation_id: str
    tier: str
    start: float  # seconds
    end: float  # seconds
    label: str


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _ms_to_sec(ms: int) -> float:
    return ms / 1000.0


def _sec_to_ms(sec: float) -> int:
    return round(sec * 1000)


def _get_mime_type(path: Path) -> str:
    suffix = path.suffix.lower()
    mime_types = {
        ".mp4": "video/mp4",
        ".avi": "video/x-msvideo",
        ".mov": "video/quicktime",
        ".webm": "video/webm",
        ".mkv": "video/x-matroska",
        ".wav": "audio/x-wav",
        ".mp3": "audio/mpeg",
    }
    return mime_types.get(suffix, "video/mp4")


# ---------------------------------------------------------------------------
# Main class
# ---------------------------------------------------------------------------


class ElanDocument:
    """XML-preserving ELAN .eaf document wrapper.

    All mutations happen on the in-memory ``ElementTree``.  Call
    :meth:`save` to write the result atomically back to disk.
    """

    def __init__(
        self,
        tree: ET.ElementTree,
        path: Path | None,
        file_mtime: float | None,
    ) -> None:
        self._tree = tree
        root = tree.getroot()
        assert root is not None, "ElementTree has no root element"
        self._root = root
        self._path = path
        self._file_mtime = file_mtime

        # Build indices
        self._time_slots: dict[str, int] = {}  # slot_id -> ms
        self._next_ts_id: int = 1
        self._next_ann_id: int = 1
        self._index_time_slots()
        self._index_annotation_ids()

    # ------------------------------------------------------------------
    # Construction
    # ------------------------------------------------------------------

    @classmethod
    def open(cls, path: str | Path) -> ElanDocument:
        """Parse an existing ``.eaf`` file.

        Stores the file's mtime so that :meth:`check_disk_conflict` can
        detect external modifications before saving.
        """
        path = Path(path)
        if not path.exists():
            raise SLTKFileNotFoundError(str(path))

        tree = SafeET.parse(str(path))
        file_mtime = os.path.getmtime(path)
        return cls(tree, path, file_mtime)

    @classmethod
    def new(cls, video_path: str | Path | None = None) -> ElanDocument:
        """Create a minimal valid ``.eaf`` document from scratch."""
        root = ET.Element("ANNOTATION_DOCUMENT")
        root.set("AUTHOR", "SLTK")
        root.set(
            "DATE",
            datetime.now().strftime("%Y-%m-%dT%H:%M:%S+00:00"),
        )
        root.set("FORMAT", "3.0")
        root.set("VERSION", "3.0")
        root.set("xmlns:xsi", "http://www.w3.org/2001/XMLSchema-instance")
        root.set(
            "xsi:noNamespaceSchemaLocation",
            "http://www.mpi.nl/tools/elan/EAFv3.0.xsd",
        )

        header = ET.SubElement(root, "HEADER")
        header.set("MEDIA_FILE", "")
        header.set("TIME_UNITS", "milliseconds")

        if video_path is not None:
            video_path = Path(video_path)
            media_desc = ET.SubElement(header, "MEDIA_DESCRIPTOR")
            media_desc.set("MEDIA_URL", f"file:///{video_path.absolute()}")
            media_desc.set("MIME_TYPE", _get_mime_type(video_path))
            media_desc.set("RELATIVE_MEDIA_URL", f"./{video_path.name}")

        ET.SubElement(root, "TIME_ORDER")

        ling = ET.SubElement(root, "LINGUISTIC_TYPE")
        ling.set("LINGUISTIC_TYPE_ID", "default-lt")
        ling.set("TIME_ALIGNABLE", "true")
        ling.set("GRAPHIC_REFERENCES", "false")

        tree = ET.ElementTree(root)
        return cls(tree, None, None)

    # ------------------------------------------------------------------
    # Internal indexing
    # ------------------------------------------------------------------

    def _time_order(self) -> ET.Element:
        el = self._root.find("TIME_ORDER")
        if el is None:
            raise ValueError("Invalid ELAN file: missing TIME_ORDER")
        return el

    def _index_time_slots(self) -> None:
        self._time_slots.clear()
        max_num = 0
        for ts in self._time_order():
            ts_id = ts.get("TIME_SLOT_ID", "")
            ts_val = ts.get("TIME_VALUE")
            if ts_id and ts_val is not None:
                self._time_slots[ts_id] = int(ts_val)
                if ts_id.startswith("ts"):
                    try:
                        max_num = max(max_num, int(ts_id[2:]))
                    except ValueError:
                        pass
        self._next_ts_id = max_num + 1

    def _index_annotation_ids(self) -> None:
        max_num = 0
        for ann in self._root.iter("ALIGNABLE_ANNOTATION"):
            ann_id = ann.get("ANNOTATION_ID", "")
            if ann_id.startswith("a"):
                try:
                    max_num = max(max_num, int(ann_id[1:]))
                except ValueError:
                    pass
        self._next_ann_id = max_num + 1

    # ------------------------------------------------------------------
    # Time-slot helpers
    # ------------------------------------------------------------------

    def _reverse_ts_map(self) -> dict[int, str]:
        """ms -> first slot_id with that value."""
        rev: dict[int, str] = {}
        for sid, ms in self._time_slots.items():
            if ms not in rev:
                rev[ms] = sid
        return rev

    def _get_or_create_ts(self, ms: int) -> str:
        """Return existing slot_id for *ms*, or create a new one."""
        rev = self._reverse_ts_map()
        if ms in rev:
            return rev[ms]

        slot_id = f"ts{self._next_ts_id}"
        self._next_ts_id += 1

        ts_elem = ET.SubElement(self._time_order(), "TIME_SLOT")
        ts_elem.set("TIME_SLOT_ID", slot_id)
        ts_elem.set("TIME_VALUE", str(ms))
        self._time_slots[slot_id] = ms
        return slot_id

    def _remove_ts_if_orphaned(self, slot_id: str) -> None:
        """Remove a time slot element if no annotation references it."""
        for ann in self._root.iter("ALIGNABLE_ANNOTATION"):
            if ann.get("TIME_SLOT_REF1") == slot_id or ann.get("TIME_SLOT_REF2") == slot_id:
                return  # still referenced
        # Remove from XML
        time_order = self._time_order()
        for ts in list(time_order):
            if ts.get("TIME_SLOT_ID") == slot_id:
                time_order.remove(ts)
                break
        self._time_slots.pop(slot_id, None)

    # ------------------------------------------------------------------
    # Annotation-ID helper
    # ------------------------------------------------------------------

    def _new_ann_id(self) -> str:
        ann_id = f"a{self._next_ann_id}"
        self._next_ann_id += 1
        return ann_id

    # ------------------------------------------------------------------
    # Tier helpers
    # ------------------------------------------------------------------

    def _find_tier(self, name: str) -> ET.Element:
        for tier in self._root.findall("TIER"):
            if tier.get("TIER_ID") == name:
                return tier
        raise KeyError(f"Tier not found: {name!r}")

    def _tier_insert_index(self) -> int:
        """Position just after the last TIER (or after TIME_ORDER)."""
        tiers = self._root.findall("TIER")
        children = list(self._root)
        if tiers:
            return children.index(tiers[-1]) + 1
        time_order = self._root.find("TIME_ORDER")
        if time_order is not None:
            return children.index(time_order) + 1
        return len(children)

    # ------------------------------------------------------------------
    # Public API — reading
    # ------------------------------------------------------------------

    def get_tiers(self) -> list[TierInfo]:
        """Return metadata for every tier in document order."""
        result: list[TierInfo] = []
        for tier in self._root.findall("TIER"):
            tid = tier.get("TIER_ID", "")
            result.append(
                TierInfo(
                    name=tid,
                    linguistic_type=tier.get("LINGUISTIC_TYPE_REF", ""),
                    participant=tier.get("PARTICIPANT", ""),
                    annotator=tier.get("ANNOTATOR", ""),
                    segment_count=len(tier.findall(".//ALIGNABLE_ANNOTATION")),
                )
            )
        return result

    def get_segments(
        self,
        tiers: list[str] | None = None,
    ) -> list[SegmentInfo]:
        """Read annotations, optionally filtered by tier names."""
        result: list[SegmentInfo] = []
        for tier in self._root.findall("TIER"):
            tid = tier.get("TIER_ID", "")
            if tiers is not None and tid not in tiers:
                continue
            for ann in tier.findall(".//ALIGNABLE_ANNOTATION"):
                ref1 = ann.get("TIME_SLOT_REF1", "")
                ref2 = ann.get("TIME_SLOT_REF2", "")
                if ref1 not in self._time_slots or ref2 not in self._time_slots:
                    continue
                val_elem = ann.find("ANNOTATION_VALUE")
                label = ""
                if val_elem is not None and val_elem.text:
                    label = val_elem.text.strip()
                result.append(
                    SegmentInfo(
                        annotation_id=ann.get("ANNOTATION_ID", ""),
                        tier=tid,
                        start=_ms_to_sec(self._time_slots[ref1]),
                        end=_ms_to_sec(self._time_slots[ref2]),
                        label=label,
                    )
                )
        return result

    # ------------------------------------------------------------------
    # Public API — mutation
    # ------------------------------------------------------------------

    def update_segment(
        self,
        ann_id: str,
        *,
        start: float | None = None,
        end: float | None = None,
        label: str | None = None,
    ) -> None:
        """Mutate an existing annotation in-place."""
        ann = self._find_annotation(ann_id)
        if start is not None:
            old_ref = ann.get("TIME_SLOT_REF1", "")
            new_slot = self._get_or_create_ts(_sec_to_ms(start))
            ann.set("TIME_SLOT_REF1", new_slot)
            if old_ref != new_slot:
                self._remove_ts_if_orphaned(old_ref)
        if end is not None:
            old_ref = ann.get("TIME_SLOT_REF2", "")
            new_slot = self._get_or_create_ts(_sec_to_ms(end))
            ann.set("TIME_SLOT_REF2", new_slot)
            if old_ref != new_slot:
                self._remove_ts_if_orphaned(old_ref)
        if label is not None:
            val_elem = ann.find("ANNOTATION_VALUE")
            if val_elem is None:
                val_elem = ET.SubElement(ann, "ANNOTATION_VALUE")
            val_elem.text = label

    def add_segment(
        self,
        tier: str,
        start: float,
        end: float,
        label: str,
    ) -> str:
        """Insert a new annotation into *tier*. Returns the annotation ID."""
        tier_elem = self._find_tier(tier)
        start_slot = self._get_or_create_ts(_sec_to_ms(start))
        end_slot = self._get_or_create_ts(_sec_to_ms(end))
        ann_id = self._new_ann_id()

        ann_wrapper = ET.SubElement(tier_elem, "ANNOTATION")
        align_ann = ET.SubElement(ann_wrapper, "ALIGNABLE_ANNOTATION")
        align_ann.set("ANNOTATION_ID", ann_id)
        align_ann.set("TIME_SLOT_REF1", start_slot)
        align_ann.set("TIME_SLOT_REF2", end_slot)
        val_elem = ET.SubElement(align_ann, "ANNOTATION_VALUE")
        val_elem.text = label

        return ann_id

    def delete_segment(self, ann_id: str) -> None:
        """Remove an annotation and clean up orphaned time slots."""
        ann = self._find_annotation(ann_id)
        ref1 = ann.get("TIME_SLOT_REF1", "")
        ref2 = ann.get("TIME_SLOT_REF2", "")

        # The ALIGNABLE_ANNOTATION sits inside an ANNOTATION wrapper
        # which sits inside a TIER.
        ann_wrapper = self._find_annotation_wrapper(ann_id)
        tier_elem = self._find_annotation_tier(ann_id)
        tier_elem.remove(ann_wrapper)

        self._remove_ts_if_orphaned(ref1)
        self._remove_ts_if_orphaned(ref2)

    def add_tier(
        self,
        name: str,
        linguistic_type: str = "default-lt",
    ) -> None:
        """Add a new empty tier."""
        # Check for duplicates
        for tier in self._root.findall("TIER"):
            if tier.get("TIER_ID") == name:
                raise ValueError(f"Tier already exists: {name!r}")

        tier_elem = ET.Element("TIER")
        tier_elem.set("LINGUISTIC_TYPE_REF", linguistic_type)
        tier_elem.set("TIER_ID", name)
        self._root.insert(self._tier_insert_index(), tier_elem)

        # Ensure the linguistic type element exists
        self._ensure_linguistic_type(linguistic_type)

    def rename_tier(self, old_name: str, new_name: str) -> None:
        """Rename a tier (updates TIER_ID attribute)."""
        tier = self._find_tier(old_name)
        tier.set("TIER_ID", new_name)

    def delete_tier(self, name: str) -> None:
        """Delete a tier, all its annotations, and orphaned time slots."""
        tier = self._find_tier(name)

        # Collect time slot refs before removal
        refs: list[str] = []
        for ann in tier.findall(".//ALIGNABLE_ANNOTATION"):
            r1 = ann.get("TIME_SLOT_REF1", "")
            r2 = ann.get("TIME_SLOT_REF2", "")
            if r1:
                refs.append(r1)
            if r2:
                refs.append(r2)

        self._root.remove(tier)

        # Cleanup orphaned time slots
        for ref in refs:
            self._remove_ts_if_orphaned(ref)

    def reorder_tiers(self, order: list[str]) -> None:
        """Reorder tiers in the XML to match *order*.

        Tier names not in *order* are appended at the end in their
        original order.
        """
        existing: dict[str, ET.Element] = {
            tid: t for t in self._root.findall("TIER") if (tid := t.get("TIER_ID")) is not None
        }
        # Remove all tiers from root
        for t in list(self._root.findall("TIER")):
            self._root.remove(t)

        insert_idx = self._tier_insert_index()

        # Re-insert in requested order
        ordered_names: list[str] = []
        for name in order:
            if name in existing:
                ordered_names.append(name)
        # Append any not mentioned
        for name in existing:
            if name not in ordered_names:
                ordered_names.append(name)

        for name in ordered_names:
            self._root.insert(insert_idx, existing[name])
            insert_idx += 1

    # ------------------------------------------------------------------
    # Conflict detection
    # ------------------------------------------------------------------

    def check_disk_conflict(self) -> bool:
        """Return ``True`` if the file on disk has been modified since open.

        Always returns ``False`` for documents created with :meth:`new`.
        """
        if self._path is None or self._file_mtime is None:
            return False
        if not self._path.exists():
            return True
        return os.path.getmtime(self._path) != self._file_mtime

    # ------------------------------------------------------------------
    # Save
    # ------------------------------------------------------------------

    def save(
        self,
        path: str | Path | None = None,
        auto_version: bool = False,
    ) -> Path:
        """Write the XML tree atomically to *path*.

        If *path* is ``None``, overwrites the original file.  When
        *auto_version* is ``True``, appends ``_v2``, ``_v3``, etc.
        instead of overwriting.

        The write is atomic: a temp file is created in the same
        directory, then ``os.replace()`` renames it into place.
        """
        if path is not None:
            dest = Path(path)
        elif self._path is not None:
            dest = self._path
        else:
            raise ValueError(
                "No path specified and document was created with new(). " "Pass a path to save()."
            )

        if auto_version:
            dest = self._versioned_path(dest)

        dest.parent.mkdir(parents=True, exist_ok=True)

        # Use XML declaration
        ET.indent(self._tree, space="    ")

        fd, tmp = tempfile.mkstemp(dir=str(dest.parent), suffix=".eaf.tmp")
        try:
            os.close(fd)
            self._tree.write(
                tmp,
                encoding="unicode",
                xml_declaration=True,
            )
            os.replace(tmp, dest)
        except BaseException:
            # Clean up temp file on any error
            try:
                os.unlink(tmp)
            except OSError:
                pass
            raise

        # Update stored mtime/path after successful save
        self._path = dest
        self._file_mtime = os.path.getmtime(dest)
        return dest

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    def _find_annotation(self, ann_id: str) -> ET.Element:
        for ann in self._root.iter("ALIGNABLE_ANNOTATION"):
            if ann.get("ANNOTATION_ID") == ann_id:
                return ann
        raise KeyError(f"Annotation not found: {ann_id!r}")

    def _find_annotation_wrapper(self, ann_id: str) -> ET.Element:
        """Find the ``<ANNOTATION>`` wrapper element for *ann_id*."""
        for tier in self._root.findall("TIER"):
            for wrapper in tier.findall("ANNOTATION"):
                align = wrapper.find("ALIGNABLE_ANNOTATION")
                if align is not None and align.get("ANNOTATION_ID") == ann_id:
                    return wrapper
        raise KeyError(f"Annotation wrapper not found: {ann_id!r}")

    def _find_annotation_tier(self, ann_id: str) -> ET.Element:
        """Find the ``<TIER>`` containing annotation *ann_id*."""
        for tier in self._root.findall("TIER"):
            for ann in tier.findall(".//ALIGNABLE_ANNOTATION"):
                if ann.get("ANNOTATION_ID") == ann_id:
                    return tier
        raise KeyError(f"No tier contains annotation: {ann_id!r}")

    def _ensure_linguistic_type(self, lt_id: str) -> None:
        """Add a LINGUISTIC_TYPE element if *lt_id* doesn't exist."""
        for lt in self._root.findall("LINGUISTIC_TYPE"):
            if lt.get("LINGUISTIC_TYPE_ID") == lt_id:
                return
        lt_elem = ET.Element("LINGUISTIC_TYPE")
        lt_elem.set("LINGUISTIC_TYPE_ID", lt_id)
        lt_elem.set("TIME_ALIGNABLE", "true")
        lt_elem.set("GRAPHIC_REFERENCES", "false")
        self._root.append(lt_elem)

    @staticmethod
    def _versioned_path(path: Path) -> Path:
        """Return ``path_v2.eaf``, ``path_v3.eaf``, etc."""
        stem = path.stem
        suffix = path.suffix
        parent = path.parent
        version = 2
        while True:
            candidate = parent / f"{stem}_v{version}{suffix}"
            if not candidate.exists():
                return candidate
            version += 1
            if version > 10000:
                raise RuntimeError(f"Too many versions for {path}")
