"""
Output writer — serializes elements/chunks to JSONL + manifest.json.

Supports:
- File output (atomic write via .tmp -> rename)
- Stdout output (with stream_end sentinel)
- Both elements mode and chunks mode
"""

from __future__ import annotations

import contextlib
import hashlib
import json
import sys
from datetime import datetime, timezone
from importlib.metadata import PackageNotFoundError as _PNF
from importlib.metadata import version as _pkg_version
from pathlib import Path
from typing import IO, Any

try:
    _EXTRACTOR_VERSION = _pkg_version("coresdk-extractor")
except _PNF:
    _EXTRACTOR_VERSION = "0.0.0+dev"

from .schema import (
    SCHEMA_VERSION,
    BaseElement,
    ChunkRecord,
    ChunkStats,
    ConfidenceStats,
    EnvelopeRecord,
    ManifestRecord,
    ManifestWarning,
    RunConfig,
    SourceInfo,
    StatsRecord,
)


def _sha256_file(path: Path) -> str:
    h = hashlib.sha256()
    with path.open("rb") as f:
        for chunk in iter(lambda: f.read(65536), b""):
            h.update(chunk)
    return h.hexdigest()


def _now_iso() -> str:
    return datetime.now(timezone.utc).isoformat()


def _serialize(obj: BaseElement | EnvelopeRecord | ManifestRecord | dict) -> str:  # type: ignore[return]
    if hasattr(obj, "model_dump"):
        return json.dumps(obj.model_dump(exclude_none=True), ensure_ascii=False)
    return json.dumps(obj, ensure_ascii=False)


class OutputWriter:
    """
    Writes elements/chunks to JSONL output and builds the manifest.

    Usage:
        writer = OutputWriter(source_path, run_config, out_path)
        writer.write_envelope()
        for element in elements:
            writer.write_element(element)
        manifest = writer.finalize()
    """

    def __init__(
        self,
        source_path: Path,
        run_config: RunConfig,
        out_path: Path | None = None,
        extractor_version: str = _EXTRACTOR_VERSION,
        source_sha256: str | None = None,
        include_full_path: bool = False,
    ) -> None:
        self.source_path = source_path
        self.include_full_path = include_full_path
        self.run_config = run_config
        self.out_path = out_path
        self.extractor_version = extractor_version

        # Use provided sha256 or compute it
        self.source_sha256 = (
            source_sha256 if source_sha256 is not None else _sha256_file(source_path)
        )
        self.source_size = source_path.stat().st_size
        self.created_at = _now_iso()

        self._elements_written = 0
        self._chunks_written = 0
        self._warnings: list[ManifestWarning] = []
        self._stats = StatsRecord()
        self._all_confidences: list[float] = []
        self._chunk_token_counts: list[int] = []

        # File handles
        self._tmp_path: Path | None = None
        self._fh: IO[str] | None = None
        self._stdout_mode = out_path is None

        self._source_info: SourceInfo | None = None
        self._finalized = False

    def _open(self, mime_type: str) -> None:
        self._source_info = SourceInfo(
            filename=self.source_path.name,
            path=str(self.source_path) if self.include_full_path else self.source_path.name,
            sha256=self.source_sha256,
            size_bytes=self.source_size,
            media_type=mime_type,
        )
        if self._stdout_mode:
            self._fh = sys.stdout
        else:
            assert self.out_path is not None
            self._tmp_path = self.out_path.with_suffix(self.out_path.suffix + ".tmp")
            self._fh = open(self._tmp_path, "w", encoding="utf-8")

    def begin(self, mime_type: str = "application/octet-stream") -> None:
        self._open(mime_type)
        assert self._source_info is not None
        envelope = EnvelopeRecord(
            extractor_version=self.extractor_version,
            source=self._source_info,
            run_config=self.run_config,
            created_at=self.created_at,
        )
        self._writeline(_serialize(envelope))

    def write_element(self, element: BaseElement) -> None:
        assert self._fh is not None
        self._writeline(_serialize(element))
        self._elements_written += 1
        self._stats.total_elements += 1

        etype = element.element_type
        self._stats.elements_by_type[etype] = self._stats.elements_by_type.get(etype, 0) + 1

        if element.confidence is not None:
            self._all_confidences.append(element.confidence)

        # Tier tracking
        tier_key = str(element.section_path_tier)
        self._stats.section_path_tiers[tier_key] = (
            self._stats.section_path_tiers.get(tier_key, 0) + 1
        )

        # Table stats
        if etype.startswith("table:"):
            self._stats.tables.total += 1
            if etype == "table:continuation":
                self._stats.tables.continuation_segments += 1
            if element.table and element.table.has_header_row:
                self._stats.tables.with_header += 1

    def write_chunk(self, chunk: ChunkRecord) -> None:
        assert self._fh is not None
        self._writeline(_serialize(chunk))
        self._chunks_written += 1
        self._stats.total_chunks += 1
        self._stats.total_tokens += chunk.token_count
        self._chunk_token_counts.append(chunk.token_count)
        self._stats.total_elements += 1
        self._stats.elements_by_type["composite:chunk"] = (
            self._stats.elements_by_type.get("composite:chunk", 0) + 1
        )

    def add_warning(self, code: str, message: str, element_id: str | None = None) -> None:
        self._warnings.append(ManifestWarning(code=code, message=message, element_id=element_id))

    def _writeline(self, line: str) -> None:
        assert self._fh is not None
        self._fh.write(line + "\n")
        if self._stdout_mode:
            self._fh.flush()

    def __enter__(self) -> OutputWriter:
        return self

    def __exit__(self, exc_type: object, exc_val: object, exc_tb: object) -> bool:
        if self._fh is not None and not self._stdout_mode:
            self._fh.close()
            self._fh = None
        if exc_type is not None and not self._finalized:
            if self._tmp_path is not None:
                with contextlib.suppress(FileNotFoundError):
                    self._tmp_path.unlink()
        return False

    def finalize(
        self,
        status: str = "complete",
        parser_chain: list[str] | None = None,
        ner_stats: Any | None = None,
    ) -> ManifestRecord:
        """Flush output, emit stream_end sentinel (stdout only), build and return manifest."""
        # Compute final stats
        if self._all_confidences:
            sorted_conf = sorted(self._all_confidences)
            n = len(sorted_conf)
            below_0_5 = sum(1 for c in sorted_conf if c < 0.5)
            below_0_7 = sum(1 for c in sorted_conf if c < 0.7)
            self._stats.confidence = ConfidenceStats(
                mean=round(sum(sorted_conf) / n, 4),
                p10=round(sorted_conf[max(0, int(n * 0.1))], 4),
                p90=round(sorted_conf[min(n - 1, int(n * 0.9))], 4),
                below_0_5=below_0_5,
                below_0_7=below_0_7,
            )

        if self._chunk_token_counts:
            self._stats.chunks = ChunkStats(
                total=len(self._chunk_token_counts),
                min_tokens=min(self._chunk_token_counts),
                max_tokens=max(self._chunk_token_counts),
                mean_tokens=round(sum(self._chunk_token_counts) / len(self._chunk_token_counts), 1),
                atomic_solo=0,  # TODO: track this
            )

        # NER stats (schema v1.1.0)
        if ner_stats is not None:
            from .schema import NERStatsRecord

            self._stats.ner = NERStatsRecord(
                elements_processed=ner_stats.elements_processed,
                elements_skipped=ner_stats.elements_skipped,
                elements_truncated=ner_stats.elements_truncated,
                inference_errors=ner_stats.inference_errors,
                total_entities=ner_stats.total_entities,
                duration_ms=round(ner_stats.duration_ms, 1),
                model_id=ner_stats.model_id,
                tier=ner_stats.tier,
            )

        # Stream end sentinel (stdout mode only)
        if self._stdout_mode and self._fh:
            sentinel = {
                "type": "stream_end",
                "status": status,
                "total_elements": self._elements_written,
                "total_chunks": self._chunks_written,
                "schema_version": SCHEMA_VERSION,
            }
            self._writeline(json.dumps(sentinel))

        # Close file handle
        if self._fh and not self._stdout_mode:
            self._fh.close()
            self._fh = None
            # Atomic rename
            if self._tmp_path and self.out_path:
                self._tmp_path.rename(self.out_path)
        self._finalized = True

        # Build manifest
        assert self._source_info is not None
        manifest_raw = f"{self.source_sha256}:manifest:{self.created_at}"
        manifest_id = "mf_" + hashlib.sha256(manifest_raw.encode()).hexdigest()[:16]

        manifest = ManifestRecord(
            manifest_id=manifest_id,
            status=status,
            source=self._source_info,
            run_config=self.run_config,
            parser_chain=parser_chain or [],
            stats=self._stats,
            warnings=self._warnings,
            created_at=self.created_at,
            completed_at=_now_iso(),
        )
        return manifest
