"""I/O and reporting helpers for OmniSTEval.

Contains shared functions for loading hypothesis/reference data, dumping
instances and scores, and formatting reports. This centralises duplicated
logic used by the CLI and resegmentation pipeline.
"""
from typing import List, Tuple, Optional, Dict
import json
import os
import logging

from .data import (
    Instance,
    load_reference,
    load_hypothesis_jsonl,
    load_hypothesis_text,
    load_text_segmentation,
    get_segmentation_order,
    load_hypothesis_simulstream,
)
from . import __version__

logger = logging.getLogger(__name__)

def load_shortform_instances(
    hypothesis_file: str,
    ref_sentences_file: Optional[str],
    emission_cu_key: str,
    emission_ca_key: str,
    char_level: bool = False,
) -> List[Instance]:
    """Load SimulEval-style shortform JSONL + references and return Instances.

    Returns tuple instances.
    """
    with open(hypothesis_file, "r", encoding="utf-8") as f:
        hyp_lines = [json.loads(line.strip()) for line in f if line.strip()]
    assert hyp_lines, f"Hypothesis file is empty: {hypothesis_file}"

    all_have_reference = hyp_lines and all("reference" in h and len(h["reference"]) > 0 for h in hyp_lines)

    if ref_sentences_file is not None:
        if all_have_reference:
            logger.warning(
                f"Hypothesis lines already contain 'reference' fields, but {ref_sentences_file} is also provided. "
                "The reference sentences from the file will be used instead of those in the hypothesis lines."
            )
        with open(ref_sentences_file, "r", encoding="utf-8") as f:
            ref_sentences = [line.strip() for line in f]
        assert len(ref_sentences) == len(hyp_lines), (
            f"Number of references ({len(ref_sentences)}) does not match number of hypothesis lines ({len(hyp_lines)})."
        )
    elif all_have_reference:
        ref_sentences = [h["reference"].strip() for h in hyp_lines] 
    else:
        raise ValueError("No reference sentences provided. Each hypothesis line must contain a 'reference' field, or a separate reference sentences file must be provided.")

    instances: List[Instance] = []
    for i, (h, ref) in enumerate(zip(hyp_lines, ref_sentences)):
        prediction = h.get("prediction", "")
        cu_values = h.get(emission_cu_key, None)
        ca_values = h.get(emission_ca_key, None)
        source_length = h.get("source_length", None)

        info: dict = {
            "index": i,
            "prediction": prediction,
            "reference": ref,
        }
        if source_length is not None:
            info["source_length"] = source_length
        if cu_values is not None:
            info["emission_cu"] = cu_values
        if ca_values is not None:
            info["emission_ca"] = ca_values

        instances.append(Instance.from_dict(info, latency_unit=("char" if char_level else "word")))

    return instances


def load_resegmentation_inputs(
    speech_segmentation: Optional[str],
    text_segmentation: Optional[str],
    ref_sentences_file: str,
    hypothesis_file: str,
    hypothesis_format: str = "jsonl",
    char_level: bool = False,
    offset_delays: bool = False,
    fix_emission_ca_flag: bool = False,
    emission_cu_key: str = "delays",
    emission_ca_key: str = "elapsed",
    simulstream_config_file: Optional[str] = None,
) -> Tuple[List[List], List[List], list, List[str]]:
    """Load inputs required for resegmentation.

    Returns (ref_words, hyp_words, segmentation, ref_sentences).
    """
    if speech_segmentation is not None:
        ref_words, segmentation, ref_sentences = load_reference(
            speech_segmentation, ref_sentences_file, char_level, offset_delays
        )
        segmentation_order = get_segmentation_order(segmentation)

        if hypothesis_format == "jsonl":
            hyp_words = load_hypothesis_jsonl(
                hypothesis_file,
                char_level,
                segmentation_order,
                fix_emission_ca_flag,
                emission_cu_key,
                emission_ca_key,
            )
        elif hypothesis_format == "text":
            hyp_words = load_hypothesis_text(hypothesis_file, char_level, len(segmentation_order))
        elif hypothesis_format == "simulstream":
            if simulstream_config_file is None:
                raise ValueError("simulstream_config_file must be provided when hypothesis_format is 'simulstream'.")
            hyp_words = load_hypothesis_simulstream(
                hypothesis_file=hypothesis_file,
                char_level=char_level,
                segmentation_order=segmentation_order,
                eval_config_file=simulstream_config_file,
            )
        else:
            raise ValueError(f"Unknown hypothesis format: {hypothesis_format}")
    elif text_segmentation is not None:
        # text segmentation
        ref_words, segmentation, ref_sentences, num_documents = load_text_segmentation(
            text_segmentation, ref_sentences_file, char_level
        )
        hyp_words = load_hypothesis_text(hypothesis_file, char_level, num_documents)
    else:
        raise ValueError("Either speech_segmentation or text_segmentation must be provided.")

    return ref_words, hyp_words, segmentation, ref_sentences


def dump_instances_jsonl(instances_dict_list: List[dict], output_folder: str) -> None:
    os.makedirs(output_folder, exist_ok=True)
    with open(os.path.join(output_folder, "instances.resegmented.jsonl"), "w", encoding="utf-8") as f:
        for d in instances_dict_list:
            f.write(json.dumps(d, ensure_ascii=False) + "\n")


def _metric_display_name(key: str, is_longform: bool) -> str:
    """Return a human-friendly display name for a metric key.

    Mirrors the display logic used by the CLI report.
    """
    metric_display = {
        "bleu":              "BLEU",
        "chrf":              "chrF",
        "comet":             "COMET",
        "yaal":              "YAAL (CU)",
        "ca_yaal":           "YAAL (CA)",
        "swf":               "Simultaneous Words Fraction (%)",
        "efsw":              "Expected Simul. Words Fraction (%)",
        "dsptv":             "Degeneracy Test Value",
        "degenerate_policy": "Likely Degenerate Simultaneous Policy",
    }

    if key in metric_display:
        disp = metric_display[key]
        if is_longform and key in ("yaal", "ca_yaal"):
            # Longform YAAL naming
            if key == "yaal":
                return "LongYAAL (CU)"
            else:
                return "LongYAAL (CA)"
        return disp

    if key.startswith("ca_long_"):
        return "Long" + key[len("ca_long_"):].upper() + " (CA)"
    if key.startswith("long_"):
        return "Long" + key[len("long_"):].upper() + " (CU)"
    if key in ("al", "laal", "ap", "dal", "yaal"):
        return key.upper() + " (CU)"
    if key.startswith("ca_"):
        return key[len("ca_"):].upper() + " (CA)"
    return key.upper()


def dump_scores_tsv(scores: Dict[str, float], output_folder: str, is_longform: bool = False) -> None:
    os.makedirs(output_folder, exist_ok=True)
    with open(
        os.path.join(output_folder, "scores.tsv"), "w", encoding="utf-8"
    ) as f:
        f.write("metric\tvalue\n")
        for k, v in scores.items():
            name = _metric_display_name(k, is_longform)
            if k == "degenerate_policy":
                val_str = "YES" if v == 1.0 else ("NO" if v == 0.0 else "N/A")
            else:
                val_str = f"{v:.4f}"
            f.write(f"{name}\t{val_str}\n")


def format_report(mode_label: str, settings: dict, scores: dict, instance_report: List[str]) -> str:
    """Return a human-readable evaluation report string.

    This mirrors the previous CLI formatting.
    """
    is_longform = "longform" in mode_label.lower()
    lines: list = ["=" * 64, f"OmniSTEval v{__version__}  |  {mode_label}", "=" * 64, ""]

    if settings:
        lines.append("Settings")
        lines.append("-" * 64)
        kw = max(len(k) for k in settings)
        for k, v in settings.items():
            lines.append(f"  {k:<{kw}}  {v}")
        lines.append("")

    if scores:
        lines.append("Scores")
        lines.append("-" * 64)
        rendered = {}
        for k, v in scores.items():
            # Use a simple uppercase name for now; the CLI will map nicer names
            k_lower = k.lower()
            name = _metric_display_name(k_lower, is_longform)
            is_degenerate = k_lower == "degenerate_policy"
            rendered[name] = (
                "YES"
                if is_degenerate and v == 1.0
                else ("NO" if is_degenerate and v == 0.0 else f"{v:.4f}")
            )
        kw = max(len(k) for k in rendered)
        for k, v in rendered.items():
            lines.append(f"  {k:<{kw}}  {v}")
        lines.append("")

        if scores.get("degenerate_policy") == 1.0:
            lines.append("  *** Likely Degenerate Simultaneous Policy ***")
            lines.append("")

    # Append instance report
    lines.append("Instance-level Details")
    lines.append("-" * 64)
    lines.extend(instance_report)

    lines.append("=" * 64)
    return "\n".join(lines)
