"""Output format handlers for segmentation predictions.

Supports saving predictions as:
- PT: PyTorch dict {video_name: tensor}
- JSON: Structured segments with timestamps
- ELAN: EAF XML file compatible with ELAN annotation software
- Spotted ELAN: EAF with per-dictionary gloss rank/score tiers

Adapted from segmenterv2/segmenter/inference/output.py.
"""

from __future__ import annotations

import json
from enum import Enum
from pathlib import Path
from typing import TYPE_CHECKING
from xml.etree.ElementTree import Element, ElementTree, SubElement, indent

import torch

if TYPE_CHECKING:
    from sltk.embedding.pipeline import SpottingResult


class OutputFormat(Enum):
    PT = "pt"
    JSON = "json"
    ELAN = "elan"


def _guess_mime(path: Path) -> str:
    """Return MIME type for common video formats."""
    ext = path.suffix.lower()
    return {
        ".mp4": "video/mp4",
        ".mov": "video/quicktime",
        ".avi": "video/x-msvideo",
        ".mkv": "video/x-matroska",
        ".webm": "video/webm",
        ".mpg": "video/mpeg",
        ".mpeg": "video/mpeg",
        ".wav": "audio/x-wav",
    }.get(ext, "video/mp4")


def _relative_media_path(media: Path, eaf: Path) -> str | None:
    """Try to compute a relative path from the EAF to the media file."""
    try:
        return str(media.resolve().relative_to(eaf.resolve().parent))
    except ValueError:
        return None


def save_predictions(
    predictions: dict[str, torch.Tensor],
    output_path: Path,
    fmt: OutputFormat,
    fps: float = 25.0,
    media_path: Path | str | None = None,
) -> Path:
    """Save predictions in the specified format.

    Args:
        predictions: Dict mapping video name to label tensor (0=OUT, 1=IN, 2=BEGIN).
        output_path: Path to save the output file.
        fmt: Output format (PT, JSON, or ELAN).
        fps: Frames per second for time conversion.
        media_path: Optional media file path (used for ELAN format).

    Returns:
        Path to the saved file.
    """
    output_path = Path(output_path)
    if fmt == OutputFormat.PT:
        return save_as_pt(predictions, output_path)
    elif fmt == OutputFormat.JSON:
        return save_as_json(predictions, output_path, fps)
    elif fmt == OutputFormat.ELAN:
        return save_as_elan(predictions, output_path, fps, media_path=media_path)
    else:
        raise ValueError(f"Unknown output format: {fmt}")


def save_as_pt(
    predictions: dict[str, torch.Tensor],
    output_path: Path,
) -> Path:
    """Save as PyTorch dict {video_name: tensor}."""
    output_path = Path(output_path)
    torch.save(predictions, output_path)
    return output_path


def _extract_segments_from_labels(labels: torch.Tensor) -> list[dict]:
    """Extract sign segments from a label tensor."""
    segments = []
    start = None

    for i, val in enumerate(labels.tolist()):
        if val in (1, 2):
            if start is None:
                start = i
        else:
            if start is not None:
                segments.append({"start_frame": start, "end_frame": i - 1})
                start = None

    if start is not None:
        segments.append({"start_frame": start, "end_frame": len(labels) - 1})

    return segments


def save_as_json(
    predictions: dict[str, torch.Tensor],
    output_path: Path,
    fps: float = 25.0,
) -> Path:
    """Save as structured JSON with segments and timestamps."""
    output_path = Path(output_path)
    result = {}

    for video_name, labels in predictions.items():
        segments = _extract_segments_from_labels(labels)

        for seg in segments:
            seg["start_sec"] = round(seg["start_frame"] / fps, 4)
            seg["end_sec"] = round(seg["end_frame"] / fps, 4)

        result[video_name] = {
            "fps": fps,
            "num_frames": len(labels),
            "segments": segments,
        }

    with open(output_path, "w") as f:
        json.dump(result, f, indent=2)

    return output_path


def save_as_elan(
    predictions: dict[str, torch.Tensor],
    output_path: Path,
    fps: float = 25.0,
    media_path: Path | str | None = None,
) -> Path:
    """Save as ELAN EAF XML file.

    Args:
        predictions: Dict mapping video name to label tensor.
        output_path: Path to save the EAF file.
        fps: Frames per second for time conversion.
        media_path: Optional path to the media file. If provided, a
            MEDIA_DESCRIPTOR is added so ELAN can open the video.
    """
    output_path = Path(output_path)

    root = Element("ANNOTATION_DOCUMENT")
    root.set("xmlns:xsi", "http://www.w3.org/2001/XMLSchema-instance")
    root.set("xsi:noNamespaceSchemaLocation", "http://www.mpi.nl/tools/elan/EAFv3.0.xsd")
    root.set("AUTHOR", "segmenter_v2")
    root.set("FORMAT", "3.0")

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

    if media_path is not None:
        media_path = Path(media_path).resolve()
        md = SubElement(header, "MEDIA_DESCRIPTOR")
        md.set("MEDIA_URL", media_path.as_uri())
        md.set("MIME_TYPE", _guess_mime(media_path))
        relative = _relative_media_path(media_path, output_path)
        if relative:
            md.set("RELATIVE_MEDIA_URL", f"./{relative}")

    time_order = SubElement(root, "TIME_ORDER")
    ts_id = 1

    for video_name, labels in predictions.items():
        segments = _extract_segments_from_labels(labels)

        tier = SubElement(root, "TIER")
        tier.set("LINGUISTIC_TYPE_REF", "default-lt")
        tier.set("TIER_ID", f"{video_name}_segmentation")

        for seg_idx, seg in enumerate(segments):
            start_ms = int(seg["start_frame"] / fps * 1000)
            end_ms = int((seg["end_frame"] + 1) / fps * 1000)

            ts_start = SubElement(time_order, "TIME_SLOT")
            ts_start.set("TIME_SLOT_ID", f"ts{ts_id}")
            ts_start.set("TIME_VALUE", str(start_ms))
            ts_start_id = f"ts{ts_id}"
            ts_id += 1

            ts_end = SubElement(time_order, "TIME_SLOT")
            ts_end.set("TIME_SLOT_ID", f"ts{ts_id}")
            ts_end.set("TIME_VALUE", str(end_ms))
            ts_end_id = f"ts{ts_id}"
            ts_id += 1

            annotation = SubElement(tier, "ANNOTATION")
            alignable = SubElement(annotation, "ALIGNABLE_ANNOTATION")
            alignable.set("ANNOTATION_ID", f"a{video_name}_{seg_idx}")
            alignable.set("TIME_SLOT_REF1", ts_start_id)
            alignable.set("TIME_SLOT_REF2", ts_end_id)
            value = SubElement(alignable, "ANNOTATION_VALUE")
            value.text = "SIGN"

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

    indent(root)
    tree = ElementTree(root)
    tree.write(output_path, encoding="unicode", xml_declaration=True)

    return output_path


def save_spotted_elan(
    spot_result: SpottingResult,
    output_path: Path,
    fps: float = 25.0,
    media_path: Path | str | None = None,
) -> Path:
    """Save spotting results as ELAN EAF with per-rank gloss and score tiers.

    Creates tiers like "Rank-1", "Rank-2", ... with gloss labels, and
    "Score-1", "Score-2", ... with similarity scores for each segment.

    Args:
        spot_result: SpottingResult from pipeline.spot().
        output_path: Path to save the EAF file.
        fps: Frames per second for time conversion.
        media_path: Optional path to the media file.
    """
    output_path = Path(output_path)
    output_path.parent.mkdir(parents=True, exist_ok=True)

    root = Element("ANNOTATION_DOCUMENT")
    root.set("xmlns:xsi", "http://www.w3.org/2001/XMLSchema-instance")
    root.set(
        "xsi:noNamespaceSchemaLocation",
        "http://www.mpi.nl/tools/elan/EAFv3.0.xsd",
    )
    root.set("AUTHOR", "signrep_spotter")
    root.set("FORMAT", "3.0")

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

    if media_path is not None:
        media_path = Path(media_path).resolve()
        md = SubElement(header, "MEDIA_DESCRIPTOR")
        md.set("MEDIA_URL", media_path.as_uri())
        md.set("MIME_TYPE", _guess_mime(media_path))
        relative = _relative_media_path(media_path, output_path)
        if relative:
            md.set("RELATIVE_MEDIA_URL", f"./{relative}")

    time_order = SubElement(root, "TIME_ORDER")
    ts_id = 1

    if not spot_result.segments:
        # Empty result
        lt = SubElement(root, "LINGUISTIC_TYPE")
        lt.set("LINGUISTIC_TYPE_ID", "default-lt")
        lt.set("TIME_ALIGNABLE", "true")
        lt.set("GRAPHIC_REFERENCES", "false")
        indent(root)
        tree = ElementTree(root)
        tree.write(output_path, encoding="unicode", xml_declaration=True)
        return output_path

    # Determine max rank from first segment
    max_rank = len(spot_result.segments[0].top_glosses) if spot_result.segments else 5

    # Create tier elements: Rank-N for glosses, Score-N for similarity
    gloss_tiers: dict[int, Element] = {}
    score_tiers: dict[int, Element] = {}
    for rank in range(1, max_rank + 1):
        gt = SubElement(root, "TIER")
        gt.set("LINGUISTIC_TYPE_REF", "default-lt")
        gt.set("TIER_ID", f"Rank-{rank}")
        gloss_tiers[rank] = gt

        st = SubElement(root, "TIER")
        st.set("LINGUISTIC_TYPE_REF", "default-lt")
        st.set("TIER_ID", f"Score-{rank}")
        score_tiers[rank] = st

    # Add annotations per segment
    for seg in spot_result.segments:
        start_ms = seg.start_ms
        end_ms = seg.end_ms

        # Create shared time slots for this segment
        ts_start_id = f"ts{ts_id}"
        ts_start = SubElement(time_order, "TIME_SLOT")
        ts_start.set("TIME_SLOT_ID", ts_start_id)
        ts_start.set("TIME_VALUE", str(start_ms))
        ts_id += 1

        ts_end_id = f"ts{ts_id}"
        ts_end = SubElement(time_order, "TIME_SLOT")
        ts_end.set("TIME_SLOT_ID", ts_end_id)
        ts_end.set("TIME_VALUE", str(end_ms))
        ts_id += 1

        for gl in seg.top_glosses:
            rank = gl.get("rank", 1)
            if rank > max_rank:
                continue
            gloss = gl.get("gloss", "")
            similarity = gl.get("similarity", 0.0)

            # Gloss annotation
            ann = SubElement(gloss_tiers[rank], "ANNOTATION")
            alignable = SubElement(ann, "ALIGNABLE_ANNOTATION")
            alignable.set("ANNOTATION_ID", f"g{seg.segment_id}_r{rank}")
            alignable.set("TIME_SLOT_REF1", ts_start_id)
            alignable.set("TIME_SLOT_REF2", ts_end_id)
            val = SubElement(alignable, "ANNOTATION_VALUE")
            val.text = str(gloss)

            # Score annotation
            ann_s = SubElement(score_tiers[rank], "ANNOTATION")
            alignable_s = SubElement(ann_s, "ALIGNABLE_ANNOTATION")
            alignable_s.set("ANNOTATION_ID", f"s{seg.segment_id}_r{rank}")
            alignable_s.set("TIME_SLOT_REF1", ts_start_id)
            alignable_s.set("TIME_SLOT_REF2", ts_end_id)
            val_s = SubElement(alignable_s, "ANNOTATION_VALUE")
            val_s.text = f"{similarity:.4f}"

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

    indent(root)
    tree = ElementTree(root)
    tree.write(output_path, encoding="unicode", xml_declaration=True)

    return output_path
