# License: Apache-2.0 License
# Created by: Patrick Lumbantobing, VertoX-AI
# Copyright (c) 2026 VertoX-AI. All rights reserved.
#
# This work is licensed under the Apache-2.0 License.
# To view a copy of this license, visit
# https://www.apache.org/licenses/LICENSE-2.0

"""Persist and verify backend-neutral full-word evidence with explicit provenance.

Native parsing remains in the existing backend adapters. This module accepts only
canonical alignment capsules and immutable file/member references, never native words.
"""

from __future__ import annotations

import hashlib
import json
import zipfile
from collections.abc import Mapping
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Any

from speech_text_data_aligners.core.errors import ConfigurationError
from speech_text_data_aligners.core.models import (
    AlignmentProvenance,
    AlignmentRequest,
    AlignmentResult,
    Diagnostic,
    DiagnosticSeverity,
    JsonValue,
    LexicalWord,
    WordAlignment,
)
from speech_text_data_aligners.core.validation import validate_alignment_result
from speech_text_data_aligners.runtime.receipts import sha256_file


@dataclass(frozen=True, slots=True)
class EvidenceCandidate:
    """Locate one previously audited canonical full-word capsule.

    Parameters
    ----------
        path: Local canonical JSON file to verify.
        sha256: Expected capsule SHA-256; missing/mismatched hashes reject reuse.

    """

    path: Path
    sha256: str


def artifact_sha256(reference: str) -> str:
    """Hash a file or one exact ``archive.zip::member`` without extracting it.

    Parameters
    ----------
        reference: Existing local path, optionally with an explicit ZIP member.

    Returns
    -------
        SHA-256 of the referenced bytes.

    Raises
    ------
        ConfigurationError: If the source/member is missing, duplicated, or unreadable.

    Side Effects:
        Reads source bytes only, including retained WAVs when verifying prior audio.

    """
    if "::" not in reference:
        return sha256_file(Path(reference))
    archive, member = reference.split("::", 1)
    try:
        with zipfile.ZipFile(archive) as source:
            if source.namelist().count(member) != 1:
                raise ConfigurationError(
                    "archive evidence member is missing or ambiguous"
                )
            with source.open(member) as stream:
                digest = hashlib.sha256()
                for chunk in iter(lambda: stream.read(1024 * 1024), b""):
                    digest.update(chunk)
                return digest.hexdigest()
    except (OSError, ValueError, zipfile.BadZipFile) as error:
        raise ConfigurationError(
            f"cannot hash archived evidence: {reference}"
        ) from error


def publish_json(path: Path, payload: object, *, jsonl: bool = False) -> None:
    """Publish a small deterministic artifact without replacing an existing path.

    Parameters
    ----------
        path: New destination, whose parent is created as needed.
        payload: JSON-safe value, or a list/tuple of JSON-safe rows for JSONL.
        jsonl: Write one compact JSON value per row when true.

    Raises
    ------
        ConfigurationError: If serialization fails or publication would overwrite data.

    Side Effects:
        Creates the destination after completing serialization with NaN forbidden.

    """
    try:
        if jsonl:
            if not isinstance(payload, (list, tuple)):
                raise ValueError("JSONL requires a row sequence")
            text = "".join(
                json.dumps(row, sort_keys=True, ensure_ascii=False, allow_nan=False)
                + "\n"
                for row in payload
            )
        else:
            text = (
                json.dumps(
                    payload,
                    sort_keys=True,
                    indent=2,
                    ensure_ascii=False,
                    allow_nan=False,
                )
                + "\n"
            )
        path.parent.mkdir(parents=True, exist_ok=True)
        with path.open("x", encoding="utf-8", newline="\n") as stream:
            stream.write(text)
    except (OSError, ValueError, TypeError) as error:
        raise ConfigurationError(f"cannot publish evidence: {path}: {error}") from error


def _freeze(value: Any) -> JsonValue:
    """Restore immutable attribute tuples from JSON ``value``; reject other types."""
    if isinstance(value, list):
        return tuple(_freeze(item) for item in value)
    if value is None or isinstance(value, (str, int, float, bool)):
        return value
    raise ConfigurationError("alignment attributes must contain JSON scalars/tuples")


def _result(payload: dict[str, Any]) -> AlignmentResult:
    """Decode a canonical ``payload`` result, preserving provenance and diagnostics.

    Returns:
        Typed immutable alignment; validation is performed by the caller.

    Raises:
        KeyError, TypeError, ValueError: If the canonical schema is malformed.
        ConfigurationError: If attribute serialization is invalid.

    """
    provenance = dict(payload["provenance"])
    provenance["attributes"] = tuple(
        (key, _freeze(value)) for key, value in provenance.get("attributes", [])
    )
    provenance["artifacts"] = tuple(provenance.get("artifacts", []))
    words = tuple(
        WordAlignment(
            LexicalWord(**word["lexical_word"]),
            word["start_s"],
            word["end_s"],
            word.get("confidence"),
            tuple((key, _freeze(value)) for key, value in word.get("attributes", [])),
        )
        for word in payload["words"]
    )
    diagnostics = tuple(
        Diagnostic(
            d["code"],
            DiagnosticSeverity(d["severity"]),
            d["message"],
            d["stage"],
            tuple((k, _freeze(v)) for k, v in d.get("details", [])),
        )
        for d in payload.get("diagnostics", [])
    )
    return AlignmentResult(
        payload["request_id"],
        words,
        AlignmentProvenance(**provenance),
        diagnostics,
        payload.get("runtime_s", 0.0),
    )


def save_alignment_evidence(
    path: Path,
    request: AlignmentRequest,
    result: AlignmentResult,
    *,
    audio_sha256: str,
    transcript_sha256: str,
    source_artifacts: Mapping[str, str],
) -> EvidenceCandidate:
    """Persist a validated reference-text result as a reusable full-word capsule.

    Parameters
    ----------
        path: New capsule destination.
        request: Original source request to validate against the result.
        result: Full neutral result returned/reconstructed by a repository adapter.
        audio_sha256: Previously verified exact WAV digest.
        transcript_sha256: Previously verified exact transcript digest.
        source_artifacts: Immutable reference-to-SHA mapping proving source/provenance.

    Returns
    -------
        Hashed candidate; reuse independently verifies its contents and references.

    Raises
    ------
        SpeechTextAlignerError: If alignment validation or publication fails.

    Side Effects:
        Publishes only canonical words and provenance; no WAV/model bytes are copied.

    """
    validate_alignment_result(request, result)
    publish_json(
        path,
        dict(
            schema_version="w1t.alignment_evidence.v1",
            audio_id=request.request_id,
            audio_sha256=audio_sha256,
            original_en_transcript_sha256=transcript_sha256,
            reference_text_authoritative=True,
            result=asdict(result),
            source_artifacts=dict(source_artifacts),
        ),
    )
    return EvidenceCandidate(path, sha256_file(path))


def verify_alignment_evidence(
    candidate: EvidenceCandidate,
    request: AlignmentRequest,
    *,
    audio_sha256: str,
    transcript_sha256: str,
    expected_identity: Mapping[str, str],
) -> AlignmentResult:
    """Require exact source, backend, model/configuration, and full-word evidence.

    Parameters
    ----------
        candidate: Canonical capsule and its expected immutable hash.
        request: Current original transcript/audio request.
        audio_sha256: Verified current WAV digest.
        transcript_sha256: Verified current exact transcript digest.
        expected_identity: Required provenance fields, including backend_name.

    Returns
    -------
        Validated full result, unchanged.

    Raises
    ------
        SpeechTextAlignerError: If a required proof is absent/mismatched or times fail.

    Side Effects:
        Hashes the capsule and every cited source artifact; no model is invoked.

    """
    if not candidate.sha256 or sha256_file(candidate.path) != candidate.sha256:
        raise ConfigurationError("capsule artifact SHA-256 mismatch")
    try:
        value = json.loads(candidate.path.read_text(encoding="utf-8"))
        if (
            value["schema_version"] != "w1t.alignment_evidence.v1"
            or value["audio_id"] != request.request_id
        ):
            raise ConfigurationError("capsule schema/audio identity mismatch")
        if value["audio_sha256"] != audio_sha256:
            raise ConfigurationError("audio SHA-256 mismatch")
        if value["original_en_transcript_sha256"] != transcript_sha256:
            raise ConfigurationError("transcript SHA-256 mismatch")
        if value["reference_text_authoritative"] is not True:
            raise ConfigurationError("reference-text authority is not proved")
        result = _result(value["result"])
        provenance = asdict(result.provenance)
        required = (
            "backend_name",
            "backend_version",
            "model_id",
            "config_sha256",
            "python_version",
            "device",
        )
        if any(not provenance[k] for k in required) or not (
            provenance["model_revision"] or provenance["model_sha256"]
        ):
            raise ConfigurationError(
                "insufficient backend/model/configuration provenance"
            )
        if not expected_identity.get("backend_name") or any(
            provenance.get(k) != v for k, v in expected_identity.items()
        ):
            raise ConfigurationError("backend identity mismatch")
        sources = value["source_artifacts"]
        if not isinstance(sources, dict) or not sources:
            raise ConfigurationError("source artifact proof is missing")
        for reference, digest in sources.items():
            if artifact_sha256(reference) != digest:
                raise ConfigurationError("source artifact SHA-256 mismatch")
        return validate_alignment_result(request, result)
    except (
        OSError,
        UnicodeError,
        ValueError,
        KeyError,
        TypeError,
        AttributeError,
    ) as error:
        raise ConfigurationError(f"malformed full-word evidence: {error}") from error