"""Concurrent, crash-recoverable recording and transcription sessions."""

from __future__ import annotations

import json
import logging
import queue
import shutil
import threading
import time
import uuid
import wave
from collections import deque
from dataclasses import asdict, dataclass, is_dataclass, replace
from pathlib import Path
from types import MappingProxyType
from typing import Any, Callable, Mapping

import numpy as np

from blurt_audio import StreamingAudio


_ACTIVE_PATHS: set[Path] = set()
_ACTIVE_LOCK = threading.Lock()


def _atomic_json(path: Path, value: dict[str, Any]) -> None:
    temporary = path.with_suffix(path.suffix + ".tmp")
    temporary.write_text(json.dumps(value, indent=2), encoding="utf-8")
    temporary.replace(path)


def _write_wav(path: Path, samples: np.ndarray) -> None:
    temporary = path.with_suffix(path.suffix + ".tmp")
    with wave.open(str(temporary), "wb") as output:
        output.setnchannels(1)
        output.setsampwidth(2)
        output.setframerate(16000)
        output.writeframes(np.asarray(samples, dtype="<i2").tobytes())
    temporary.replace(path)


def _read_wav(path: Path) -> np.ndarray:
    with wave.open(str(path), "rb") as source:
        if source.getnchannels() != 1 or source.getsampwidth() != 2:
            raise ValueError(f"unsupported recovery WAV format: {path}")
        data = source.readframes(source.getnframes())
        rate = source.getframerate()
    samples = np.frombuffer(data, dtype="<i2").copy()
    if rate != 16000:
        raise ValueError(f"recovery WAV must be 16 kHz: {path}")
    return samples


def _config_dict(config: Any) -> dict[str, Any]:
    if is_dataclass(config):
        return asdict(config)
    if isinstance(config, Mapping):
        return dict(config)
    return dict(vars(config))


def _public_config(config: Any) -> dict[str, Any]:
    blocked = ("key", "token", "secret", "password", "auth", "credential")
    return {
        str(k): v
        for k, v in _config_dict(config).items()
        if not any(word in str(k).lower() for word in blocked)
        and isinstance(v, (str, int, float, bool, type(None)))
    }


@dataclass(frozen=True)
class SessionResult:
    session_id: str
    text: str
    complete: bool
    errors: tuple[str, ...] = ()


@dataclass(frozen=True)
class RecoveryRecord:
    session_id: str
    path: Path
    state: str
    config: Mapping[str, Any]
    done: int
    total: int
    created_at: float = 0.0


@dataclass
class _Job:
    session: "RecordingSession"
    index: int
    path: Path
    config: Any


class TranscriptionScheduler:
    """A shared worker pool whose pending payloads live as WAV files on disk."""

    def __init__(self, request: Callable[[np.ndarray, Any], str], workers: int = 3,
                 queue_size: int = 96):
        if workers < 1:
            raise ValueError("workers must be positive")
        self.request = request
        self._ready: queue.Queue[_Job | None] = queue.Queue(maxsize=max(workers, queue_size))
        self._overflow: deque[_Job] = deque()
        self._lock = threading.Lock()
        self._closed = False
        self._threads = [threading.Thread(target=self._worker, name=f"blurt-upload-{i}", daemon=True)
                         for i in range(workers)]
        for thread in self._threads:
            thread.start()

    def submit(self, session: "RecordingSession", index: int, path: Path,
               config: Any) -> None:
        job = _Job(session, index, path, config)
        with self._lock:
            if self._closed:
                raise RuntimeError("transcription scheduler is shut down")
            try:
                self._ready.put_nowait(job)
            except queue.Full:
                # Only tiny disk references are retained here; audio remains on disk.
                self._overflow.append(job)

    def _refill(self) -> None:
        with self._lock:
            while self._overflow:
                try:
                    self._ready.put_nowait(self._overflow[0])
                except queue.Full:
                    break
                self._overflow.popleft()

    def _worker(self) -> None:
        while True:
            try:
                job = self._ready.get(timeout=0.1)
            except queue.Empty:
                with self._lock:
                    if self._closed and not self._overflow:
                        return
                continue
            text = None
            error = None
            try:
                samples = _read_wav(job.path)
                text = self.request(samples, job.config)
                if not isinstance(text, str) or not text.strip():
                    raise ValueError("transcription returned empty text")
                text = text.strip()
            except Exception as exception:
                error = f"{type(exception).__name__}: {exception}"
            try:
                job.session._job_finished(job.index, text, error)
            except Exception:
                # Session persistence/callback failures must not terminate the pool.
                logging.exception("Could not record transcription result for %s", job.path)
            finally:
                self._ready.task_done()
                self._refill()

    def shutdown(self, wait: bool = True) -> None:
        with self._lock:
            self._closed = True
        if wait:
            for thread in self._threads:
                thread.join()


class RecordingSession:
    """Owns capture, chunk persistence, ordered results, and one recording lifecycle."""

    def __init__(self, rate: int, channels: int, config: Any, spool_dir: str | Path,
                 scheduler: TranscriptionScheduler, on_update: Callable[["RecordingSession"], None] | None = None,
                 on_complete: Callable[["RecordingSession"], None] | None = None,
                 input_queue_size: int = 64):
        self.session_id = uuid.uuid4().hex
        self.path = Path(spool_dir) / self.session_id
        self.path.mkdir(parents=True, exist_ok=False)
        self.config = config
        self.scheduler = scheduler
        self.on_update = on_update
        self.on_complete = on_complete
        self.state = "recording"
        self.done = 0
        self.total = 0
        self.result: SessionResult | None = None
        self.created_at = time.time()
        self._entries: list[dict[str, Any]] = []
        self._submitted: set[int] = set()
        self._capture_persisted_ok = False
        self._lock = threading.RLock()
        self._condition = threading.Condition(self._lock)
        self._input: queue.Queue[np.ndarray] = queue.Queue(maxsize=input_queue_size)
        self._stop = threading.Event()
        self._producer_done = threading.Event()
        self._overflow_error = threading.Event()
        self._input_statuses: queue.SimpleQueue[str] = queue.SimpleQueue()
        self._audio = StreamingAudio(rate, channels)
        self._last_checkpoint = time.monotonic()
        self._tail_start = 0
        self._manifest_path = self.path / "session.json"
        with _ACTIVE_LOCK:
            _ACTIVE_PATHS.add(self.path.resolve())
        try:
            self._save()
        except Exception:
            self.release_for_retry()
            raise
        self._producer = threading.Thread(target=self._produce, name=f"blurt-capture-{self.session_id[:6]}", daemon=True)
        self._producer.start()

    def feed(self, indata: np.ndarray) -> bool:
        if self.state != "recording":
            return False
        try:
            self._input.put_nowait(np.array(indata, copy=True))
            return True
        except queue.Full:
            # The sound callback must never wait for a lock or filesystem.
            self.state = "capture_failed"
            self._overflow_error.set()
            self._stop.set()
            return False

    def stop(self) -> None:
        with self._lock:
            if self.state == "recording":
                self.state = "finishing"
                self._stop.set()
                self._save()

    def report_input_status(self, status: object) -> None:
        """Record a PortAudio status without blocking its callback thread."""
        if status:
            self._input_statuses.put(str(status))

    def producer_join(self, timeout: float | None = None) -> bool:
        self._producer.join(timeout)
        return not self._producer.is_alive()

    def join(self, timeout: float | None = None) -> SessionResult | None:
        deadline = None if timeout is None else time.monotonic() + timeout
        with self._condition:
            while self.result is None:
                remaining = None if deadline is None else deadline - time.monotonic()
                if remaining is not None and remaining <= 0:
                    return None
                self._condition.wait(remaining)
            return self.result

    def _produce(self) -> None:
        try:
            while True:
                try:
                    block = self._input.get(timeout=0.05)
                except queue.Empty:
                    if self._stop.is_set():
                        break
                    continue
                try:
                    for chunk in self._audio.feed(block):
                        self._seal(chunk)
                    if time.monotonic() - self._last_checkpoint >= 1.0:
                        self._checkpoint()
                finally:
                    self._input.task_done()
            for chunk in self._audio.finish():
                self._seal(chunk)
            self._remove_tail()
            if self._overflow_error.is_set():
                with self._lock:
                    self._entries.append({"index": len(self._entries), "speech": False,
                        "status": "failed", "error": "capture input queue overflowed; audio was lost"})
                    self._save()
            self._persist_input_statuses()
            self._capture_persisted_ok = True
        except Exception as error:
            with self._lock:
                self.state = "capture_failed"
                for entry in self._entries:
                    if (entry.get("speech") and entry.get("status") == "pending"
                            and int(entry["index"]) not in self._submitted):
                        entry["status"] = "failed"
                        entry["error"] = f"capture: {type(error).__name__}: {error}"
                        self.done += 1
                self._entries.append({"index": len(self._entries), "speech": False,
                                      "status": "failed",
                                      "error": f"capture: {type(error).__name__}: {error}"})
                try:
                    self._save()
                except OSError:
                    pass
        finally:
            self._producer_done.set()
            self._maybe_complete()

    def _persist_input_statuses(self) -> None:
        statuses = []
        while not self._input_statuses.empty():
            statuses.append(self._input_statuses.get())
        if statuses:
            with self._lock:
                self._entries.append({"index": len(self._entries), "speech": False,
                    "status": "failed", "error": "input stream status: " + ", ".join(statuses)})
                self._save()

    @staticmethod
    def _chunk_value(chunk: Any, name: str, default: Any = None) -> Any:
        return getattr(chunk, name, default)

    def _seal(self, chunk: Any) -> None:
        samples = np.asarray(self._chunk_value(chunk, "samples"), dtype=np.int16)
        speech = bool(self._chunk_value(chunk, "speech", True))
        with self._lock:
            index = len(self._entries)
            filename = f"chunk-{index:06d}.wav"
            entry = {"index": index, "file": filename, "speech": speech,
                     "start": self._chunk_value(chunk, "start", None),
                     "end": self._chunk_value(chunk, "end", None), "status": "pending"}
            if not speech:
                entry["status"] = "skipped"
                entry["text"] = ""
            self._entries.append(entry)
            self.total += int(speech)
            _write_wav(self.path / filename, samples)
            # Commit the sealed file before removing the previous recovery tail.
            self._save()
            self._remove_tail_locked()
            self._tail_start = int(entry.get("end") or self._tail_start)
            self._save()
        if speech:
            self.scheduler.submit(self, index, self.path / filename, self.config)
            with self._lock:
                self._submitted.add(index)
        self._notify()

    def _checkpoint(self) -> None:
        tail = np.asarray(self._audio.checkpoint(), dtype=np.int16)
        with self._lock:
            if tail.size:
                _write_wav(self.path / "tail.wav", tail)
                ends = [int(e["end"]) for e in self._entries if e.get("end") is not None]
                self._tail_start = max(ends, default=0)
            else:
                self._remove_tail_locked()
            self._last_checkpoint = time.monotonic()
            self._save()

    def _remove_tail_locked(self) -> None:
        (self.path / "tail.wav").unlink(missing_ok=True)

    def _remove_tail(self) -> None:
        with self._lock:
            self._remove_tail_locked()
            self._save()

    def _job_finished(self, index: int, text: str | None, error: str | None) -> None:
        with self._lock:
            entry = self._entries[index]
            entry["status"] = "done" if error is None else "failed"
            if text is not None:
                entry["text"] = text
            if error is not None:
                entry["error"] = error
            self.done += 1
            try:
                self._save()
            except OSError as persistence_error:
                entry.pop("text", None)
                entry["status"] = "failed"
                entry["error"] = (
                    f"could not persist transcription: {type(persistence_error).__name__}: "
                    f"{persistence_error}"
                )
        self._notify()
        self._maybe_complete()

    def _maybe_complete(self) -> None:
        callback = None
        with self._condition:
            if not self._producer_done.is_set() or self.done < self.total or self.result is not None:
                return
            errors = tuple(str(e["error"]) for e in self._entries if e.get("error"))
            has_text = any(entry.get("text") for entry in self._entries)
            parts = []
            gap = False
            for entry in self._entries:
                if entry.get("text"):
                    if gap and has_text:
                        parts.append("[…]")
                    gap = False
                    parts.append(str(entry["text"]))
                elif entry.get("error"):
                    gap = True
            if gap and has_text:
                parts.append("[…]")
            text = " ".join(parts)
            complete = not errors
            self.state = "complete" if complete else "needs_retry"
            self.result = SessionResult(self.session_id, text, complete, errors)
            try:
                self._save()
            except OSError:
                complete = False
                self.state = "needs_retry"
                errors = errors + ("could not persist final session state",)
                self.result = SessionResult(self.session_id, text, False, errors)
            self._condition.notify_all()
            callback = self.on_complete
        if callback:
            try:
                callback(self)
            except Exception:
                logging.exception("Session completion callback failed: %s", self.session_id)
                self.release_for_retry()

    def _notify(self) -> None:
        if self.on_update:
            try:
                self.on_update(self)
            except Exception:
                logging.exception("Session progress callback failed: %s", self.session_id)

    def _save(self) -> None:
        _atomic_json(self._manifest_path, {"version": 1, "session_id": self.session_id,
            "state": self.state, "config": _public_config(self.config), "done": self.done,
            "total": self.total, "created_at": self.created_at, "entries": self._entries,
            "has_tail": (self.path / "tail.wav").exists(), "tail_start": self._tail_start})

    def acknowledge_delivery(self) -> None:
        with self._lock:
            if self.result is None:
                raise RuntimeError("session is not finished")
            if self.result.complete:
                shutil.rmtree(self.path)
            else:
                self.state = "needs_retry"
                self._save()
        with _ACTIVE_LOCK:
            _ACTIVE_PATHS.discard(self.path.resolve())

    def release_for_retry(self) -> None:
        """Make retained data discoverable after delivery is deferred or fails."""
        with _ACTIVE_LOCK:
            _ACTIVE_PATHS.discard(self.path.resolve())

    @classmethod
    def recover(cls, path: str | Path, config: Any, scheduler: TranscriptionScheduler,
                on_update=None, on_complete=None) -> "RecordingSession":
        return _recover_session(Path(path), config, scheduler, on_update, on_complete)

    @classmethod
    def from_wav(cls, wav_path: str | Path, config: Any, spool_dir: str | Path,
                 scheduler: TranscriptionScheduler, on_update=None, on_complete=None) -> "RecordingSession":
        samples = _read_wav(Path(wav_path))
        session = cls(16000, 1, config, spool_dir, scheduler, on_update, on_complete)
        session.feed(samples.reshape(-1, 1).astype(np.float32) / 32768.0)
        session.stop()
        session.producer_join()
        if session._capture_persisted_ok:
            Path(wav_path).unlink(missing_ok=True)
        return session


def _recover_session(path: Path, config: Any, scheduler: TranscriptionScheduler,
                     on_update=None, on_complete=None) -> RecordingSession:
    data = json.loads((path / "session.json").read_text(encoding="utf-8"))
    session = object.__new__(RecordingSession)
    session.session_id, session.path = data["session_id"], path
    saved = data.get("config", {})
    if is_dataclass(config):
        fields = getattr(config, "__dataclass_fields__", {})
        session.config = replace(config, **{k: v for k, v in saved.items() if k in fields})
    else:
        merged = dict(saved); merged.update(_config_dict(config))
        session.config = MappingProxyType(merged)
    session.scheduler, session.on_update, session.on_complete = scheduler, on_update, on_complete
    session.state, session.done, session.total = "finishing", 0, 0
    session.result = None
    session.created_at = float(data.get("created_at", path.stat().st_mtime))
    session._entries = data.get("entries", [])
    session._submitted = set()
    session._capture_persisted_ok = True
    session._lock = threading.RLock(); session._condition = threading.Condition(session._lock)
    session._input = queue.Queue(); session._stop = threading.Event(); session._stop.set()
    session._producer_done = threading.Event(); session._producer_done.set()
    session._overflow_error = threading.Event()
    session._input_statuses = queue.SimpleQueue()
    session._audio = None; session._last_checkpoint = time.monotonic()
    session._tail_start = int(data.get("tail_start", 0))
    session._manifest_path = path / "session.json"
    referenced = {entry.get("file") for entry in session._entries}
    orphans = [candidate for candidate in sorted(path.glob("chunk-*.wav"))
               if candidate.name not in referenced]
    sealed_end = max((int(e.get("end") or 0) for e in session._entries), default=0)
    for orphan in orphans:
        samples = _read_wav(orphan)
        index = len(session._entries)
        speech = bool(samples.size and np.max(np.abs(samples.astype(np.int32))) > 1)
        session._entries.append({"index": index, "file": orphan.name, "speech": speech,
                                 "status": "pending" if speech else "skipped",
                                 "start": sealed_end, "end": sealed_end + len(samples),
                                 "recovered_orphan": True})
        sealed_end += len(samples)
    tail = path / "tail.wav"
    if tail.exists():
        tail_samples = _read_wav(tail)
        overlap = max(0, sealed_end - session._tail_start)
        if overlap:
            tail_samples = tail_samples[min(overlap, len(tail_samples)):]
        index = len(session._entries)
        target = path / f"chunk-{index:06d}.wav"
        if len(tail_samples):
            _write_wav(target, tail_samples)
            speech = bool(np.max(np.abs(tail_samples.astype(np.int32))) > 1)
            session._entries.append({"index": index, "file": target.name, "speech": speech,
                                     "status": "pending" if speech else "skipped",
                                     "start": sealed_end, "end": sealed_end + len(tail_samples),
                                     "recovered_tail": True})
    pending = []
    for entry in session._entries:
        if entry.get("speech") and entry.get("status") != "done" and entry.get("file") and (path / entry["file"]).exists():
            entry["status"] = "pending"; entry.pop("error", None)
            pending.append(entry)
        elif entry.get("speech") and entry.get("status") != "done":
            entry["status"] = "failed"
            entry["error"] = "recovery audio file is missing"
            session.done += 1
        elif entry.get("speech"):
            session.done += 1
    session.total = sum(bool(e.get("speech")) for e in session._entries)
    with _ACTIVE_LOCK:
        _ACTIVE_PATHS.add(path.resolve())
    # Commit adopted orphan/suffix entries before deleting the original tail.
    try:
        session._save()
        tail.unlink(missing_ok=True)
        session._save()
    except Exception:
        session.release_for_retry()
        raise

    def start_recovery() -> None:
        for entry in pending:
            try:
                scheduler.submit(session, entry["index"], path / entry["file"], session.config)
                session._submitted.add(int(entry["index"]))
            except Exception as error:
                session._job_finished(entry["index"], None, f"Could not schedule retry: {error}")
        session._maybe_complete()

    session._producer = threading.Thread(target=start_recovery, name="BlurtRecovery", daemon=True)
    session._producer.start()
    return session


class RecoveryStore:
    def __init__(self, spool_dir: str | Path):
        self.spool_dir = Path(spool_dir)

    def discover(self) -> list[RecoveryRecord]:
        records = []
        if not self.spool_dir.exists():
            return records
        for manifest in sorted(self.spool_dir.glob("*/session.json")):
            try:
                with _ACTIVE_LOCK:
                    if manifest.parent.resolve() in _ACTIVE_PATHS:
                        continue
                data = json.loads(manifest.read_text(encoding="utf-8"))
                records.append(RecoveryRecord(data["session_id"], manifest.parent,
                    data.get("state", "interrupted"), MappingProxyType(data.get("config", {})),
                    int(data.get("done", 0)), int(data.get("total", 0)),
                    float(data.get("created_at", manifest.stat().st_mtime))))
            except (OSError, ValueError, KeyError):
                continue
        return sorted(records, key=lambda record: record.created_at)

    def resume(self, record: RecoveryRecord | str, config: Any,
               scheduler: TranscriptionScheduler, on_update=None, on_complete=None) -> RecordingSession:
        if isinstance(record, str):
            path = self.spool_dir / record
        else:
            path = record.path
        return RecordingSession.recover(path, config, scheduler, on_update, on_complete)