modules/models.py

"""
Pydantic data models for the meeting intelligence pipeline.
"""
from __future__ import annotations

import uuid
from datetime import datetime
from typing import Dict, List, Optional

from pydantic import BaseModel, Field


# ---------------------------------------------------------------------------
# Transcription
# ---------------------------------------------------------------------------
class TranscriptionSegment(BaseModel):
    """A single segment returned by Whisper."""
    start: float
    end: float
    text: str


# ---------------------------------------------------------------------------
# Diarization
# ---------------------------------------------------------------------------
class DiarizationSegment(BaseModel):
    """A single speaker turn returned by pyannote."""
    start: float
    end: float
    speaker: str


# ---------------------------------------------------------------------------
# Aligned (merged) segment
# ---------------------------------------------------------------------------
class AlignedSegment(BaseModel):
    """Transcript segment associated with a speaker."""
    start: float
    end: float
    speaker: str
    text: str


# ---------------------------------------------------------------------------
# Meeting Intelligence
# ---------------------------------------------------------------------------
class ActionItem(BaseModel):
    task: str
    owner: str = "Unassigned"
    deadline: str = "Not specified"
    priority: str = "Not specified"


class MeetingIntelligence(BaseModel):
    executive_summary: str = ""
    key_points: List[str] = Field(default_factory=list)
    pain_points: List[str] = Field(default_factory=list)
    decisions: List[str] = Field(default_factory=list)
    action_items: List[ActionItem] = Field(default_factory=list)
    recommendations: List[str] = Field(default_factory=list)
    follow_ups: List[str] = Field(default_factory=list)
    open_questions: List[str] = Field(default_factory=list)


# ---------------------------------------------------------------------------
# Meeting Record (top-level entity)
# ---------------------------------------------------------------------------
class MeetingRecord(BaseModel):
    meeting_id: str = Field(default_factory=lambda: str(uuid.uuid4()))
    filename: str = ""
    created_at: datetime = Field(default_factory=datetime.now)
    duration: float = 0.0
    language: str = "en"
    segments: List[AlignedSegment] = Field(default_factory=list)
    intelligence: Optional[MeetingIntelligence] = None
    speaker_map: Dict[str, str] = Field(default_factory=dict)

    # -- helpers ----------------------------------------------------------

    def get_speaker_label(self, original_label: str) -> str:
        """Return the user-assigned name for a speaker, or the original."""
        return self.speaker_map.get(original_label, original_label)

    def get_formatted_transcript(self) -> str:
        """Full transcript with timestamps and speaker labels."""
        lines: list[str] = []
        for seg in self.segments:
            ts = f"[{format_time(seg.start)} - {format_time(seg.end)}]"
            speaker = self.get_speaker_label(seg.speaker)
            lines.append(f"{ts}\n{speaker}:\n{seg.text}\n")
        return "\n".join(lines)

    def get_plain_transcript(self) -> str:
        """Plain speaker: text transcript for LLM consumption."""
        lines: list[str] = []
        for seg in self.segments:
            speaker = self.get_speaker_label(seg.speaker)
            lines.append(f"{speaker}: {seg.text}")
        return "\n".join(lines)


# ---------------------------------------------------------------------------
# Utilities
# ---------------------------------------------------------------------------
def format_time(seconds: float) -> str:
    """Convert seconds to HH:MM:SS."""
    h = int(seconds // 3600)
    m = int((seconds % 3600) // 60)
    s = int(seconds % 60)
    return f"{h:02d}:{m:02d}:{s:02d}"