"""
VTT Transcript Generator

Generates WebVTT transcripts from speaking events.
Uses state machine for correct timing and utterance segmentation.
Supports S3 upload for transcript storage.

KEY CONCEPTS:
- UtteranceBuffer: accumulates text and timing for one speaking segment
- State Machine: manages transitions between roles (user/assistant)
- Deduplication: removes duplicate start/stop speaking events
- Role Switch: finalizes utterance when speaking role changes
- S3 Storage: uploads generated VTT files to AWS S3
"""

import json
import os
from datetime import datetime, timedelta, timezone
from typing import Dict, Optional

from loguru import logger
from pipecat.frames.frames import (
    BotStartedSpeakingFrame,
    BotStoppedSpeakingFrame,
    UserStartedSpeakingFrame,
    UserStoppedSpeakingFrame,
)
from pipecat.observers.base_observer import BaseObserver

from utils.s3_client import get_s3_client


# Speaking event markers
STARTED_SPEAKING = "__STARTED_SPEAKING__"
STOPPED_SPEAKING = "__STOPPED_SPEAKING__"


class UtteranceBuffer:
    """
    Buffer for one speaking segment of a specific role.
    Accumulates text and manages start/end timing.
    """

    def __init__(self):
        self.start: Optional[timedelta] = None
        self.end: Optional[timedelta] = None
        self.text: str = ""
        self.status: Optional[str] = None  # 'active' or 'closed'

    def start_speaking(self, ts: timedelta):
        """Start of speaking segment: resets buffer, sets start/end/status."""
        self.start = ts
        self.end = ts
        self.text = ""
        self.status = "active"

    def stop_speaking(self, ts: timedelta):
        """End of speaking segment: updates end, changes status to 'closed'."""
        if self.start is not None:
            self.end = ts
            self.status = "closed"

    def add_text(self, text: str, ts: timedelta):
        """Adds text to buffer. If buffer is empty — initializes it."""
        if self.start is None:
            self.start = ts
            self.end = ts
            self.status = "active"
        self.end = ts
        self.text += text
        if self.status == "closed":
            self.status = "active"

    def should_finalize(self) -> bool:
        """Whether buffer should be finalized (has text and defined start)."""
        return bool(self.text.strip() and self.start is not None)

    def finalize(self, end_time: Optional[timedelta] = None):
        """
        Finalizes buffer: returns utterance dict, resets buffer.
        If end == start — adds 10ms for correct VTT.
        """
        if not self.should_finalize():
            return None
        end = end_time if end_time is not None else self.end
        if end is None:
            return None
        if end == self.start:
            end = end + timedelta(milliseconds=10)
        result = {
            "start": self.start,
            "end": end,
            "text": self.text,
        }
        self.reset()
        return result

    def reset(self):
        """Resets buffer to initial state."""
        self.start = None
        self.end = None
        self.text = ""
        self.status = None


class TranscriptCollector:
    """
    State machine for generating transcripts from speaking events.

    WORKING LOGIC:
    - Each role (user/assistant) has separate UtteranceBuffer
    - Utterance finalization happens on role change (role switch)
    - Interrupts handled correctly: end of previous = start of next
    """

    def __init__(self):
        self.session_start = None
        self.messages = []  # Finalized utterances for VTT
        self._buffers: Dict[str, UtteranceBuffer] = {
            "user": UtteranceBuffer(),
            "assistant": UtteranceBuffer(),
        }
        self._active_role = None
        self.logs = []  # Event history for debugging
        self._last_started_speaking: Dict[str, Optional[timedelta]] = {
            "user": None,
            "assistant": None,
        }

    def start_session(self, session_start_iso):
        """
        Initializes session start. Clears all buffers.
        Bot starts speaking at 00:00:00.000 (first speaker).
        """
        logger.debug(f"[TranscriptCollector] start_session: {session_start_iso}")
        self.session_start = datetime.fromisoformat(
            session_start_iso.replace("Z", "+00:00")
        )
        self.messages = []
        self._active_role = None
        self._buffers = {"user": UtteranceBuffer(), "assistant": UtteranceBuffer()}
        self.logs = []
        self._last_started_speaking = {"user": None, "assistant": None}
        # Bot starts speaking immediately at 00:00:00.000
        self._buffers["assistant"].start_speaking(timedelta(seconds=0))
        self._last_started_speaking["assistant"] = timedelta(seconds=0)

    def process_utterance(self, role, text, timestamp_iso: Optional[str] = None):
        """
        State machine for processing events:
        - STARTED_SPEAKING: finalizes previous utterance, opens new buffer
        - STOPPED_SPEAKING: closes buffer (doesn't finalize immediately)
        - Text: adds to buffer, updates end
        """
        if timestamp_iso is None:
            timestamp_iso = datetime.now(timezone.utc).isoformat()

        if self.session_start is None:
            # Pre-recording: use 00:00:00.000 as baseline for VTT timeline
            rel_ts = timedelta(seconds=0)
        else:
            # Normal processing with real relative time
            ts = datetime.fromisoformat(timestamp_iso.replace("Z", "+00:00"))
            rel_ts = ts - self.session_start
        rel_time_str = self._format_td(rel_ts)
        self.logs.append(
            {
                "event": "process_utterance",
                "role": role,
                "text": text,
                "timestamp": timestamp_iso,
                "rel_time": rel_time_str,
            }
        )
        buf = self._buffers[role]
        other_role = "assistant" if role == "user" else "user"

        # STARTED_SPEAKING
        if text == STARTED_SPEAKING:
            # Finalize previous utterance of this role
            finalized = buf.finalize()
            if finalized:
                finalized["role"] = role
                self.messages.append(finalized)
            # Finalize utterance of other role (role switch)
            finalized_other = self._buffers[other_role].finalize()
            if finalized_other:
                finalized_other["role"] = other_role
                self.messages.append(finalized_other)
            buf.start_speaking(rel_ts)
            self._last_started_speaking[role] = rel_ts
            return

        # STOPPED_SPEAKING
        if text == STOPPED_SPEAKING:
            buf.stop_speaking(rel_ts)
            return

        # TEXT
        buf.add_text(text, rel_ts)

    def finalize(self):
        """Finalizes all buffers with text for both roles."""
        for role in ["user", "assistant"]:
            finalized = self._buffers[role].finalize()
            if finalized:
                finalized["role"] = role
                self.messages.append(finalized)
        self._active_role = None
        self._buffers = {"user": UtteranceBuffer(), "assistant": UtteranceBuffer()}

    def get_messages(self) -> list:
        """
        Get transcript messages in a serializable format for database storage.

        Returns list of dicts with role and content (text).
        """
        result = []
        sorted_messages = sorted(
            self.messages,
            key=lambda m: m["start"].total_seconds() if m.get("start") else 0,
        )
        for m in sorted_messages:
            result.append({
                "role": m.get("role", "unknown"),
                "content": m.get("text", ""),
            })
        return result

    def to_vtt(self):
        """Generates VTT file from messages (sorted by time) and DEBUG LOGS."""
        logger.debug(f"[TranscriptCollector] to_vtt called. messages={self.messages}")
        sorted_messages = sorted(
            self.messages,
            key=lambda m: m["start"].total_seconds() if m["start"] is not None else 0,
        )
        vtt = "WEBVTT\n\n"
        for i, m in enumerate(sorted_messages):
            start = self._format_td(m["start"])
            end = self._format_td(m["end"])
            vtt += f"transcript:{i + 1}\n{start} --> {end}\n<v>{m['role']}:</v>{m['text']}\n\n"
        # Add event log to end of VTT
        # vtt += "NOTE\n"
        # vtt += "==== DEBUG LOGS (JSON) ====\n"
        # vtt += json.dumps(self.logs, ensure_ascii=False, indent=2)
        # vtt += "\n"
        return vtt

    def save_to_s3(
        self,
        key: str,
        vtt_content: Optional[str] = None,
        save_logs_separately: bool = True,
    ) -> bool:
        """
        Saves transcript VTT content to S3.

        Args:
            key: S3 object key (e.g., 'transcripts/session_123.vtt')
            vtt_content: VTT content to upload. If None, generates from current messages.
            save_logs_separately: If True, saves logs as separate JSON file

        Returns:
            True if upload succeeded, False otherwise
        """
        if vtt_content is None:
            vtt_content = self.to_vtt()

        try:
            s3_client = get_s3_client()

            # Upload VTT file
            if not s3_client.upload_transcript(key, vtt_content):
                return False

            # Save logs separately if requested
            if save_logs_separately and self.logs:
                base_key = key.replace(".vtt", "") if key.endswith(".vtt") else key
                logs_key = f"{base_key}_logs.json"
                logs_content = json.dumps(self.logs, ensure_ascii=False, indent=2)
                s3_client.upload_json(logs_key, logs_content)

            return True
        except Exception as e:
            logger.error(f"[TranscriptCollector] Error uploading to S3: {e}")
            return False

    def save_to_file(self, file_path: str, vtt_content: Optional[str] = None) -> bool:
        """
        Saves transcript VTT content to a local file.

        Args:
            file_path: Local file path to save the transcript
            vtt_content: VTT content to save. If None, generates from current messages.

        Returns:
            True if save succeeded, False otherwise
        """
        if vtt_content is None:
            vtt_content = self.to_vtt()

        try:
            os.makedirs(os.path.dirname(file_path), exist_ok=True)

            with open(file_path, "w", encoding="utf-8") as f:
                f.write(vtt_content)

            logger.debug(f"[TranscriptCollector] Saved transcript to {file_path}")
            return True
        except Exception as e:
            logger.error(f"[TranscriptCollector] Error saving to file: {e}")
            return False

    def _format_td(self, td: timedelta) -> str:
        """Formats timedelta to VTT string HH:MM:SS.mmm"""
        total_seconds = int(td.total_seconds())
        ms = int(td.microseconds / 1000)
        hours = total_seconds // 3600
        minutes = (total_seconds % 3600) // 60
        seconds = total_seconds % 60
        return f"{hours:02}:{minutes:02}:{seconds:02}.{ms:03}"


# TranscriptObserver deduplication logic:
#
# Deduplication of start/stop speaking events:
# Some voice transports may send duplicate events (e.g., multiple times in a row user_started_speaking without user_stopped_speaking between them).
# To avoid unnecessary state changes and incorrect utterances, we keep the last speaking status for each role (user/assistant).
# If a start event comes in while is_speaking==True — we ignore it. If a stop event comes in while is_speaking==False — we ignore it.
# This ensures that each "mute→speaking" and "speaking→mute" transition is handled only once.


class TranscriptObserver(BaseObserver):
    """
    Observer for processing Pipecat frames with speaking events.
    Performs deduplication of start/stop speaking events.
    """

    def __init__(self, transcript_collector):
        super().__init__()
        self._name = "TranscriptObserver"
        self.transcript_collector = transcript_collector
        # By default bot speaks from session start
        self._speaking_state = {"user": False, "assistant": True}

    async def on_push_frame(self, data):
        try:
            frame = getattr(data, "frame", None)
            if frame is None:
                return
            now = datetime.now(timezone.utc)
            now_iso = now.isoformat()
            rel_time_str = None
            if self.transcript_collector.session_start:
                rel_time_str = self.transcript_collector._format_td(
                    now - self.transcript_collector.session_start
                )

            # Deduplication for each role
            if isinstance(frame, BotStartedSpeakingFrame):
                if self._speaking_state["assistant"]:
                    logger.debug(
                        "[TranscriptObserver] Duplicate bot_started_speaking ignored"
                    )
                    return
                self._speaking_state["assistant"] = True
                logger.debug("[TranscriptObserver] 🤖 Bot started speaking (frame)")
                self.transcript_collector.logs.append(
                    {
                        "event": "bot_started_speaking",
                        "role": "assistant",
                        "timestamp": now_iso,
                        "rel_time": rel_time_str,
                    }
                )
                self.transcript_collector.process_utterance(
                    "assistant", STARTED_SPEAKING, now_iso
                )
            elif isinstance(frame, BotStoppedSpeakingFrame):
                if not self._speaking_state["assistant"]:
                    logger.debug(
                        "[TranscriptObserver] Duplicate bot_stopped_speaking ignored"
                    )
                    return
                self._speaking_state["assistant"] = False
                logger.debug("[TranscriptObserver] 🤖 Bot stopped speaking (frame)")
                self.transcript_collector.logs.append(
                    {
                        "event": "bot_stopped_speaking",
                        "role": "assistant",
                        "timestamp": now_iso,
                        "rel_time": rel_time_str,
                    }
                )
                self.transcript_collector.process_utterance(
                    "assistant", STOPPED_SPEAKING, now_iso
                )
            elif isinstance(frame, UserStartedSpeakingFrame):
                if self._speaking_state["user"]:
                    logger.debug(
                        "[TranscriptObserver] Duplicate user_started_speaking ignored"
                    )
                    return
                self._speaking_state["user"] = True
                logger.debug("[TranscriptObserver] 👤 User started speaking (frame)")
                self.transcript_collector.logs.append(
                    {
                        "event": "user_started_speaking",
                        "role": "user",
                        "timestamp": now_iso,
                        "rel_time": rel_time_str,
                    }
                )
                self.transcript_collector.process_utterance(
                    "user", STARTED_SPEAKING, now_iso
                )
            elif isinstance(frame, UserStoppedSpeakingFrame):
                if not self._speaking_state["user"]:
                    logger.debug(
                        "[TranscriptObserver] Duplicate user_stopped_speaking ignored"
                    )
                    return
                self._speaking_state["user"] = False
                logger.debug("[TranscriptObserver] 👤 User stopped speaking (frame)")
                self.transcript_collector.logs.append(
                    {
                        "event": "user_stopped_speaking",
                        "role": "user",
                        "timestamp": now_iso,
                        "rel_time": rel_time_str,
                    }
                )
                self.transcript_collector.process_utterance(
                    "user", STOPPED_SPEAKING, now_iso
                )
        except Exception as e:
            logger.debug(f"[TranscriptObserver] Error processing frame: {e}")
            return