"""
Unified annotation format I/O for SLTK.

Provides reading, writing, and validation for the SLTK unified
annotation format, enabling consistent data access across all
sign language datasets.
"""

from __future__ import annotations

import json
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

from sltk.data.core import Sample, Segment, SegmentList
from sltk.exceptions import ConfigurationError, SLTKFileNotFoundError

# Schema version
SCHEMA_VERSION = "1.0.0"

# Standard tier names
TIER_GLOSS = "gloss"
TIER_GLOSS_RH = "gloss-rh"
TIER_GLOSS_LH = "gloss-lh"
TIER_MOUTHING = "mouthing"
TIER_TRANSLATION = "translation"
TIER_CAPTION = "caption"
TIER_SUBTITLE = "subtitle"
TIER_SIGN_BOUNDARY = "sign-boundary"
TIER_DEFAULT = "default"

STANDARD_TIERS = {
    TIER_GLOSS,
    TIER_GLOSS_RH,
    TIER_GLOSS_LH,
    TIER_MOUTHING,
    TIER_TRANSLATION,
    TIER_CAPTION,
    TIER_SUBTITLE,
    TIER_SIGN_BOUNDARY,
    TIER_DEFAULT,
}


def read_unified_annotation(
    path: str | Path,
    load_video_path: bool = False,
    video_root: Path | None = None,
) -> Sample:
    """
    Read a unified annotation file and return a Sample.

    Args:
        path: Path to the unified annotation JSON file
        load_video_path: Whether to resolve video_relative_path to absolute path
        video_root: Root directory for video files (required if load_video_path=True)

    Returns:
        Sample instance with annotation data

    Example:
        >>> sample = read_unified_annotation("annotations/bslcp/BF01n.json")
        >>> print(sample.id)
        'BF01n'
        >>> print(len(sample.segments))
        142
    """
    path = Path(path)

    if not path.exists():
        raise SLTKFileNotFoundError(str(path))

    with open(path, encoding="utf-8") as f:
        data = json.load(f)

    # Parse segments
    segments = None
    if "segments" in data and data["segments"]:
        segment_list = []
        for seg_data in data["segments"]:
            segment = Segment(
                start=seg_data["start"],
                end=seg_data["end"],
                label=seg_data.get("label"),
                confidence=seg_data.get("confidence"),
                tier=seg_data.get("tier", TIER_DEFAULT),
                metadata=seg_data.get("metadata", {}),
            )
            segment_list.append(segment)
        segments = SegmentList(segment_list)

    # Resolve video path if requested
    video_path = None
    if load_video_path and data.get("video_relative_path"):
        if video_root is None:
            raise ConfigurationError("video_root required when load_video_path=True")
        video_path = video_root / data["video_relative_path"]

    # Build metadata
    metadata = data.get("metadata", {})
    if "source" in data:
        metadata["source"] = data["source"]
    if "duration_sec" in data:
        metadata["duration_sec"] = data["duration_sec"]
    if "fps" in data:
        metadata["fps"] = data["fps"]

    return Sample(
        id=data["sample_id"],
        video_path=video_path,
        poses=None,
        segments=segments,
        glosses=data.get("glosses"),
        text=data.get("text"),
        signer_id=data.get("signer_id"),
        dataset=data.get("dataset"),
        split=data.get("split"),
        metadata=metadata,
    )


def write_unified_annotation(
    sample: Sample,
    path: str | Path,
    source_format: str | None = None,
    source_path: str | None = None,
    video_relative_path: str | None = None,
    duration_sec: float | None = None,
    fps: float | None = None,
) -> Path:
    """
    Write a Sample to a unified annotation file.

    Args:
        sample: Sample to write
        path: Output path for JSON file
        source_format: Original annotation format (for provenance)
        source_path: Original annotation file path (for provenance)
        video_relative_path: Relative path to video file
        duration_sec: Duration in seconds (overrides sample.duration)
        fps: Frames per second

    Returns:
        Path to the written file

    Example:
        >>> sample = Sample(id="BF01n", dataset="bslcp", segments=my_segments)
        >>> write_unified_annotation(sample, "annotations/bslcp/BF01n.json")
    """
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)

    # Determine duration
    if duration_sec is None:
        duration_sec = sample.duration
    if duration_sec is None and sample.metadata:
        duration_sec = sample.metadata.get("duration_sec")

    # Determine fps
    if fps is None and sample.metadata:
        fps = sample.metadata.get("fps")

    # Build segments list
    segments_data: list[dict[str, Any]] = []
    if sample.segments:
        for seg in sample.segments:
            seg_data: dict[str, Any] = {
                "start": seg.start,
                "end": seg.end,
            }
            if seg.label is not None:
                seg_data["label"] = seg.label
            if seg.tier != TIER_DEFAULT:
                seg_data["tier"] = seg.tier
            if seg.confidence is not None:
                seg_data["confidence"] = seg.confidence
            if seg.metadata:
                seg_data["metadata"] = seg.metadata
            segments_data.append(seg_data)

    # Build the annotation document
    data: dict[str, Any] = {
        "schema_version": SCHEMA_VERSION,
        "sample_id": sample.id,
        "dataset": sample.dataset or "unknown",
    }

    # Optional fields (only include if set)
    if video_relative_path:
        data["video_relative_path"] = video_relative_path
    if duration_sec is not None:
        data["duration_sec"] = duration_sec
    if fps is not None:
        data["fps"] = fps
    if sample.signer_id:
        data["signer_id"] = sample.signer_id
    if sample.split:
        data["split"] = sample.split
    if segments_data:
        data["segments"] = segments_data
    if sample.glosses:
        data["glosses"] = sample.glosses
    if sample.text:
        data["text"] = sample.text

    # Add metadata (excluding internal fields)
    if sample.metadata:
        metadata = {
            k: v for k, v in sample.metadata.items() if k not in {"source", "duration_sec", "fps"}
        }
        if metadata:
            data["metadata"] = metadata

    # Add source information
    if source_format or source_path:
        data["source"] = {
            "format": source_format,
            "path": source_path,
            "converted_at": datetime.now(timezone.utc).isoformat() + "Z",
        }

    with open(path, "w", encoding="utf-8") as f:
        json.dump(data, f, indent=2, ensure_ascii=False)

    return path


def validate_annotation(
    data: dict[str, Any] | str | Path,
    strict: bool = False,
) -> tuple[bool, list[str]]:
    """
    Validate an annotation document against the schema.

    Args:
        data: Annotation data as dict, JSON string, or file path
        strict: If True, use jsonschema for full validation (requires jsonschema)

    Returns:
        Tuple of (is_valid, list_of_errors)

    Example:
        >>> is_valid, errors = validate_annotation("annotations/bslcp/BF01n.json")
        >>> if not is_valid:
        ...     print("Errors:", errors)
    """
    errors = []

    # Load data if needed
    if isinstance(data, (str, Path)):
        path = Path(data)
        if path.exists():
            try:
                with open(path, encoding="utf-8") as f:
                    data = json.load(f)
            except json.JSONDecodeError as e:
                return False, [f"Invalid JSON in file {path}: {e}"]
        else:
            return False, [f"File not found: {path}"]

    if isinstance(data, str):
        try:
            data = json.loads(data)
        except json.JSONDecodeError as e:
            return False, [f"Invalid JSON: {e}"]

    if not isinstance(data, dict):
        return False, [f"Expected dict, got {type(data).__name__}"]

    # Basic validation (always performed)
    required_fields = {"schema_version", "sample_id", "dataset"}
    missing = required_fields - set(data.keys())
    if missing:
        errors.append(f"Missing required fields: {missing}")

    # Schema version check
    if data.get("schema_version") != SCHEMA_VERSION:
        errors.append(
            f"Unsupported schema version: {data.get('schema_version')} "
            f"(expected {SCHEMA_VERSION})"
        )

    # Validate segments
    if "segments" in data:
        if not isinstance(data["segments"], list):
            errors.append("segments must be a list")
        else:
            for i, seg in enumerate(data["segments"]):
                seg_errors = _validate_segment(seg, i)
                errors.extend(seg_errors)

    # Validate split value
    valid_splits = {"train", "val", "test", "dev", None}
    if "split" in data and data["split"] not in valid_splits:
        errors.append(f"Invalid split value: {data['split']}")

    # Strict validation with jsonschema
    if strict:
        try:
            import jsonschema

            schema_path = (
                Path(__file__).parent.parent
                / "data"
                / "annotations"
                / "schema"
                / "sltk_annotation_v1.0.schema.json"
            )

            if schema_path.exists():
                with open(schema_path) as f:
                    schema = json.load(f)

                try:
                    jsonschema.validate(data, schema)
                except jsonschema.ValidationError as e:
                    errors.append(f"Schema validation error: {e.message}")
            else:
                errors.append(f"Schema file not found: {schema_path}")

        except ImportError:
            errors.append("jsonschema not installed for strict validation")

    return len(errors) == 0, errors


def _validate_segment(seg: dict[str, Any], index: int) -> list[str]:
    """Validate a single segment."""
    errors = []
    prefix = f"Segment {index}"

    if not isinstance(seg, dict):
        return [f"{prefix}: must be a dictionary"]

    # Required fields
    if "start" not in seg:
        errors.append(f"{prefix}: missing 'start'")
    elif not isinstance(seg["start"], (int, float)):
        errors.append(f"{prefix}: 'start' must be a number")
    elif seg["start"] < 0:
        errors.append(f"{prefix}: 'start' must be >= 0")

    if "end" not in seg:
        errors.append(f"{prefix}: missing 'end'")
    elif not isinstance(seg["end"], (int, float)):
        errors.append(f"{prefix}: 'end' must be a number")

    # Start/end consistency
    if "start" in seg and "end" in seg:
        if isinstance(seg["start"], (int, float)) and isinstance(seg["end"], (int, float)):
            if seg["end"] < seg["start"]:
                errors.append(f"{prefix}: 'end' must be >= 'start'")

    # Optional tier validation
    if "tier" in seg and seg["tier"] not in STANDARD_TIERS:
        # Warn but don't error - allow custom tiers
        pass

    # Confidence range
    if "confidence" in seg and seg["confidence"] is not None:
        if not isinstance(seg["confidence"], (int, float)):
            errors.append(f"{prefix}: 'confidence' must be a number")
        elif not 0 <= seg["confidence"] <= 1:
            errors.append(f"{prefix}: 'confidence' must be between 0 and 1")

    return errors


def read_manifest(path: str | Path) -> dict[str, Any]:
    """
    Read a dataset manifest file.

    Args:
        path: Path to manifest.json

    Returns:
        Manifest data as dictionary
    """
    path = Path(path)

    if not path.exists():
        raise SLTKFileNotFoundError(str(path))

    with open(path, encoding="utf-8") as f:
        return json.load(f)


def write_manifest(
    dataset: str,
    samples: list[dict[str, Any]],
    output_path: str | Path,
    source_format: str | None = None,
    tier_mapping: dict[str, str] | None = None,
    statistics: dict[str, Any] | None = None,
    notes: str | None = None,
) -> Path:
    """
    Write a dataset manifest file.

    Args:
        dataset: Dataset name
        samples: List of sample summaries (sample_id, split, num_segments, etc.)
        output_path: Output path for manifest.json
        source_format: Original annotation format
        tier_mapping: Mapping from original tier names to standard tiers
        statistics: Aggregate statistics for the dataset
        notes: Optional notes about the conversion

    Returns:
        Path to written manifest file
    """
    output_path = Path(output_path)
    output_path.parent.mkdir(parents=True, exist_ok=True)

    manifest: dict[str, Any] = {
        "dataset": dataset,
        "version": SCHEMA_VERSION,
        "created_at": datetime.now(timezone.utc).isoformat() + "Z",
        "samples": samples,
    }

    if source_format:
        manifest["source_format"] = source_format
    if tier_mapping:
        manifest["tier_mapping"] = tier_mapping
    if statistics:
        manifest["statistics"] = statistics
    if notes:
        manifest["notes"] = notes

    with open(output_path, "w", encoding="utf-8") as f:
        json.dump(manifest, f, indent=2, ensure_ascii=False)

    return output_path


def generate_manifest(
    annotation_dir: str | Path,
    dataset: str,
    source_format: str | None = None,
    tier_mapping: dict[str, str] | None = None,
) -> dict[str, Any]:
    """
    Generate a manifest from a directory of annotation files.

    Args:
        annotation_dir: Directory containing unified annotation JSON files
        dataset: Dataset name
        source_format: Original annotation format
        tier_mapping: Mapping from original tier names to standard tiers

    Returns:
        Generated manifest dictionary (also writes manifest.json to annotation_dir)
    """
    annotation_dir = Path(annotation_dir)

    if not annotation_dir.exists():
        raise SLTKFileNotFoundError(str(annotation_dir))

    samples = []
    all_glosses = set()
    all_signers = set()
    all_tiers = set()
    total_segments = 0
    total_duration = 0.0
    splits_count: dict[str, int] = {}

    # Process all annotation files
    for json_file in sorted(annotation_dir.glob("*.json")):
        if json_file.name == "manifest.json":
            continue

        with open(json_file, encoding="utf-8") as f:
            data = json.load(f)

        sample_info = {
            "sample_id": data["sample_id"],
        }

        if "split" in data:
            sample_info["split"] = data["split"]
            splits_count[data["split"]] = splits_count.get(data["split"], 0) + 1

        if "segments" in data:
            num_segs = len(data["segments"])
            sample_info["num_segments"] = num_segs
            total_segments += num_segs

            for seg in data["segments"]:
                if seg.get("label"):
                    all_glosses.add(seg["label"])
                if seg.get("tier"):
                    all_tiers.add(seg["tier"])

        if "duration_sec" in data and data["duration_sec"]:
            sample_info["duration_sec"] = data["duration_sec"]
            total_duration += data["duration_sec"]

        if "signer_id" in data and data["signer_id"]:
            sample_info["signer_id"] = data["signer_id"]
            all_signers.add(data["signer_id"])

        samples.append(sample_info)

    # Build statistics
    statistics = {
        "total_samples": len(samples),
        "total_segments": total_segments,
        "unique_glosses": len(all_glosses),
        "unique_signers": len(all_signers),
        "total_duration_sec": round(total_duration, 2),
        "tiers": sorted(list(all_tiers)),
    }

    if splits_count:
        statistics["samples_by_split"] = splits_count

    # Write manifest
    manifest_path = annotation_dir / "manifest.json"
    write_manifest(
        dataset=dataset,
        samples=samples,
        output_path=manifest_path,
        source_format=source_format,
        tier_mapping=tier_mapping,
        statistics=statistics,
    )

    # Return the full manifest
    return {
        "dataset": dataset,
        "version": SCHEMA_VERSION,
        "source_format": source_format,
        "tier_mapping": tier_mapping,
        "samples": samples,
        "statistics": statistics,
    }


def list_samples(
    annotation_dir: str | Path,
    split: str | None = None,
) -> list[str]:
    """
    List sample IDs from a dataset's unified annotations.

    Args:
        annotation_dir: Directory containing unified annotation JSON files
        split: Optional split filter ('train', 'val', 'test')

    Returns:
        List of sample IDs
    """
    annotation_dir = Path(annotation_dir)
    manifest_path = annotation_dir / "manifest.json"

    if manifest_path.exists():
        # Use manifest for efficiency
        manifest = read_manifest(manifest_path)
        samples = manifest.get("samples", [])

        if split:
            return [s["sample_id"] for s in samples if s.get("split") == split]
        return [s["sample_id"] for s in samples]

    # Fall back to scanning files
    sample_ids = []
    for json_file in sorted(annotation_dir.glob("*.json")):
        if json_file.name == "manifest.json":
            continue

        if split:
            with open(json_file, encoding="utf-8") as f:
                data = json.load(f)
            if data.get("split") == split:
                sample_ids.append(data["sample_id"])
        else:
            # Extract sample_id from filename
            sample_ids.append(json_file.stem)

    return sample_ids
