from __future__ import annotations

import secrets
import threading
import time
import uuid
from collections import deque
from dataclasses import dataclass, field


@dataclass
class DuplexSession:
    session_id: str
    resume_token: str
    thread_id: str
    sample_rate: int
    channels: int
    codec: str
    run_id: str = ""
    created_at: float = field(default_factory=time.time)
    updated_at: float = field(default_factory=time.time)
    last_seq: int = 0
    bytes_received: int = 0
    partial_text: str = ""
    final_text: str = ""
    speech_active: bool = False
    last_audio_ts: float = field(default_factory=time.time)
    speech_start_ts: float = 0.0
    vad_silence_ms: int = 700
    max_utterance_ms: int = 12000
    forced_finalize_ms: int = 1500
    punctuation_stabilization_ms: int = 350
    control_events: deque[float] = field(default_factory=deque)


class DuplexSessionController:
    def __init__(
        self,
        max_sessions: int = 64,
        idle_ttl_s: float = 900.0,
        max_frame_bytes: int = 256 * 1024,
        max_session_bytes: int = 64 * 1024 * 1024,
        max_control_events_per_min: int = 180,
    ) -> None:
        self._max_sessions = max_sessions
        self._idle_ttl_s = idle_ttl_s
        self.max_frame_bytes = max_frame_bytes
        self.max_session_bytes = max_session_bytes
        self.max_control_events_per_min = max_control_events_per_min
        self._lock = threading.Lock()
        self._sessions: dict[str, DuplexSession] = {}

    def active_count(self) -> int:
        with self._lock:
            self._evict_stale()
            return len(self._sessions)

    def end_sessions_for_thread(self, thread_id: str) -> int:
        with self._lock:
            self._evict_stale()
            doomed = [sid for sid, sess in self._sessions.items() if sess.thread_id == thread_id]
            for sid in doomed:
                self._sessions.pop(sid, None)
            return len(doomed)

    def session_age_seconds(self, session_id: str) -> float:
        with self._lock:
            sess = self._sessions.get(session_id)
            if not sess:
                return 0.0
            return max(0.0, time.time() - float(sess.created_at))

    def idle_age_seconds(self, session_id: str) -> float:
        with self._lock:
            sess = self._sessions.get(session_id)
            if not sess:
                return 0.0
            return max(0.0, time.time() - float(sess.updated_at))

    def _evict_stale(self) -> None:
        now = time.time()
        stale = [
            sid
            for sid, sess in self._sessions.items()
            if (now - sess.updated_at) > self._idle_ttl_s
        ]
        for sid in stale:
            self._sessions.pop(sid, None)
        while len(self._sessions) > self._max_sessions:
            oldest = min(self._sessions.values(), key=lambda s: s.updated_at, default=None)
            if oldest is None:
                break
            self._sessions.pop(oldest.session_id, None)

    def start_session(
        self,
        thread_id: str,
        sample_rate: int,
        channels: int,
        codec: str,
        run_id: str = "",
        vad_silence_ms: int = 700,
        max_utterance_ms: int = 12000,
        forced_finalize_ms: int = 1500,
        punctuation_stabilization_ms: int = 350,
    ) -> DuplexSession:
        with self._lock:
            self._evict_stale()
            sid = str(uuid.uuid4())
            token = secrets.token_urlsafe(24)
            sess = DuplexSession(
                session_id=sid,
                resume_token=token,
                thread_id=thread_id,
                sample_rate=sample_rate,
                channels=channels,
                codec=codec,
                run_id=run_id,
                vad_silence_ms=max(200, min(2500, int(vad_silence_ms))),
                max_utterance_ms=max(1000, min(60000, int(max_utterance_ms))),
                forced_finalize_ms=max(200, min(5000, int(forced_finalize_ms))),
                punctuation_stabilization_ms=max(0, min(3000, int(punctuation_stabilization_ms))),
            )
            self._sessions[sid] = sess
            return sess

    def resume_session(self, session_id: str, resume_token: str) -> DuplexSession | None:
        with self._lock:
            self._evict_stale()
            sess = self._sessions.get(session_id)
            if not sess:
                return None
            if secrets.compare_digest(sess.resume_token, resume_token):
                sess.updated_at = time.time()
                return sess
            return None

    def get(self, session_id: str) -> DuplexSession | None:
        with self._lock:
            sess = self._sessions.get(session_id)
            if not sess:
                return None
            sess.updated_at = time.time()
            return sess

    def validate_and_advance_seq(self, sess: DuplexSession, seq: int | None) -> bool:
        if seq is None:
            return True
        with self._lock:
            cur = self._sessions.get(sess.session_id)
            if not cur:
                return False
            if seq <= cur.last_seq:
                return False
            cur.last_seq = seq
            cur.updated_at = time.time()
            return True

    def count_control_event(self, sess: DuplexSession) -> bool:
        now = time.time()
        with self._lock:
            cur = self._sessions.get(sess.session_id)
            if not cur:
                return False
            while cur.control_events and (now - cur.control_events[0]) > 60.0:
                cur.control_events.popleft()
            if len(cur.control_events) >= self.max_control_events_per_min:
                return False
            cur.control_events.append(now)
            cur.updated_at = now
            return True

    def add_audio_bytes(self, sess: DuplexSession, nbytes: int) -> tuple[bool, str | None]:
        with self._lock:
            cur = self._sessions.get(sess.session_id)
            if not cur:
                return False, "unknown_session"
            if nbytes > self.max_frame_bytes:
                return False, "frame_too_large"
            if cur.bytes_received + nbytes > self.max_session_bytes:
                return False, "session_limit_exceeded"
            cur.bytes_received += nbytes
            cur.updated_at = time.time()
            return True, None

    def append_partial_text(self, sess: DuplexSession, text: str) -> None:
        with self._lock:
            cur = self._sessions.get(sess.session_id)
            if not cur:
                return
            incoming = str(text or "").strip()
            if not incoming:
                return
            existing = str(cur.partial_text or "").strip()
            if not existing:
                merged = incoming
            elif incoming.startswith(existing):
                # Most STT engines send cumulative partials; keep newest snapshot.
                merged = incoming
            elif existing.startswith(incoming):
                # If the recognizer briefly regresses to an older prefix, keep the
                # longer hypothesis to avoid thrashing.
                merged = existing
            else:
                # Recognizer rewrites can alter earlier words. If both strings share a
                # substantial common prefix, treat incoming as the newest snapshot
                # instead of concatenating and exploding prompt length.
                common_prefix = 0
                max_common = min(len(existing), len(incoming))
                while common_prefix < max_common and existing[common_prefix] == incoming[common_prefix]:
                    common_prefix += 1
                min_common_for_replace = max(24, int(0.5 * max_common))
                if common_prefix >= min_common_for_replace:
                    merged = incoming
                else:
                    # True append fallback for non-cumulative chunking.
                    overlap = 0
                    max_overlap = min(len(existing), len(incoming))
                    for i in range(max_overlap, 0, -1):
                        if existing.endswith(incoming[:i]):
                            overlap = i
                            break
                    merged = existing + incoming[overlap:]
            cur.partial_text = merged.strip()
            cur.updated_at = time.time()

    def should_finalize_by_timeout(self, sess: DuplexSession, now_ts: float | None = None) -> bool:
        now = now_ts if now_ts is not None else time.time()
        with self._lock:
            cur = self._sessions.get(sess.session_id)
            if not cur:
                return False
            if not cur.speech_active:
                return False
            since_last_ms = (now - cur.last_audio_ts) * 1000.0
            since_start_ms = (now - cur.speech_start_ts) * 1000.0 if cur.speech_start_ts else 0.0
            if since_last_ms >= cur.vad_silence_ms:
                return True
            if since_start_ms >= cur.max_utterance_ms:
                return True
            return False

    def mark_audio_activity(self, sess: DuplexSession, rms: float | None = None) -> tuple[bool, bool]:
        """Returns (speech_started, speech_ended_by_silence)."""
        now = time.time()
        with self._lock:
            cur = self._sessions.get(sess.session_id)
            if not cur:
                return False, False
            speech_started = False
            speech_ended = False
            level = float(rms or 0.0)
            active = level >= 120.0
            if active:
                if not cur.speech_active:
                    cur.speech_active = True
                    cur.speech_start_ts = now
                    speech_started = True
                cur.last_audio_ts = now
            elif cur.speech_active:
                silence_ms = (now - cur.last_audio_ts) * 1000.0
                if silence_ms >= cur.vad_silence_ms:
                    cur.speech_active = False
                    speech_ended = True
            cur.updated_at = now
            return speech_started, speech_ended

    def finalize_text(self, sess: DuplexSession) -> str:
        with self._lock:
            cur = self._sessions.get(sess.session_id)
            if not cur:
                return ""
            latest_snapshot = str(cur.partial_text or "").strip()
            previous_snapshot = str(cur.final_text or "").strip()
            emitted = latest_snapshot
            if latest_snapshot and previous_snapshot:
                if latest_snapshot.startswith(previous_snapshot):
                    emitted = latest_snapshot[len(previous_snapshot):].strip()
                else:
                    # Handle recognizer rewrites while still preventing replay of
                    # the full cumulative transcript across utterances.
                    common_prefix = 0
                    max_common = min(len(latest_snapshot), len(previous_snapshot))
                    while (
                        common_prefix < max_common
                        and latest_snapshot[common_prefix] == previous_snapshot[common_prefix]
                    ):
                        common_prefix += 1
                    if common_prefix >= max(24, int(0.5 * max_common)):
                        emitted = latest_snapshot[common_prefix:].strip()
            cur.final_text = latest_snapshot
            cur.partial_text = ""
            cur.updated_at = time.time()
            return emitted

    def end_session(self, session_id: str) -> None:
        with self._lock:
            self._sessions.pop(session_id, None)

