# Familiar - Self-Hosted AI Companion Platform
# Copyright (c) 2026 George Scott Foley
# ORCID: 0009-0006-4957-0540
# Email: Georgescottfoley@proton.me
# Licensed under the MIT License - see LICENSE file for details

"""
Video Transcription using OpenAI Whisper.

Whisper is a speech recognition model that runs locally.
Supports multiple languages and provides word-level timestamps.
"""

import asyncio
import json
import logging
import os
import subprocess
import tempfile
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Any, Dict, List, Optional

logger = logging.getLogger(__name__)

# Try to import whisper
try:
    import whisper

    HAS_WHISPER = True
except ImportError:
    HAS_WHISPER = False
    logger.debug("whisper not available")


@dataclass
class TranscriptSegment:
    """A segment of transcribed text with timing."""

    id: int
    start: float  # seconds
    end: float  # seconds
    text: str

    @property
    def start_formatted(self) -> str:
        """Format start time as HH:MM:SS."""
        return self._format_time(self.start)

    @property
    def end_formatted(self) -> str:
        """Format end time as HH:MM:SS."""
        return self._format_time(self.end)

    @property
    def duration(self) -> float:
        """Segment duration in seconds."""
        return self.end - self.start

    def _format_time(self, seconds: float) -> str:
        hours, remainder = divmod(int(seconds), 3600)
        minutes, secs = divmod(remainder, 60)
        if hours:
            return f"{hours}:{minutes:02d}:{secs:02d}"
        return f"{minutes}:{secs:02d}"

    def to_dict(self) -> Dict[str, Any]:
        return asdict(self)

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "TranscriptSegment":
        return cls(**data)


@dataclass
class Transcript:
    """Complete transcript of a video."""

    video_title: str
    language: str
    duration: float
    segments: List[TranscriptSegment]

    @property
    def full_text(self) -> str:
        """Get full transcript text."""
        return " ".join(seg.text for seg in self.segments)

    @property
    def segment_count(self) -> int:
        return len(self.segments)

    def get_text_at(self, timestamp: float) -> Optional[TranscriptSegment]:
        """Get segment at a specific timestamp."""
        for seg in self.segments:
            if seg.start <= timestamp < seg.end:
                return seg
        return None

    def get_text_between(
        self,
        start: float,
        end: float,
    ) -> List[TranscriptSegment]:
        """Get segments between two timestamps."""
        return [seg for seg in self.segments if seg.end > start and seg.start < end]

    def search(self, query: str, max_results: int = 10) -> List[Dict[str, Any]]:
        """
        Search for text in transcript.

        Returns matching segments with context.
        """
        query_lower = query.lower()
        results = []

        for i, seg in enumerate(self.segments):
            if query_lower in seg.text.lower():
                # Get surrounding context
                context_before = ""
                context_after = ""

                if i > 0:
                    context_before = self.segments[i - 1].text
                if i < len(self.segments) - 1:
                    context_after = self.segments[i + 1].text

                results.append(
                    {
                        "segment_id": seg.id,
                        "timestamp": seg.start,
                        "timestamp_formatted": seg.start_formatted,
                        "text": seg.text,
                        "context_before": context_before,
                        "context_after": context_after,
                    }
                )

                if len(results) >= max_results:
                    break

        return results

    def to_srt(self) -> str:
        """Export as SRT subtitle format."""
        lines = []
        for i, seg in enumerate(self.segments, 1):
            start = self._srt_time(seg.start)
            end = self._srt_time(seg.end)
            lines.append(f"{i}")
            lines.append(f"{start} --> {end}")
            lines.append(seg.text)
            lines.append("")
        return "\n".join(lines)

    def _srt_time(self, seconds: float) -> str:
        """Format time for SRT (HH:MM:SS,mmm)."""
        hours, remainder = divmod(int(seconds), 3600)
        minutes, secs = divmod(remainder, 60)
        millis = int((seconds % 1) * 1000)
        return f"{hours:02d}:{minutes:02d}:{secs:02d},{millis:03d}"

    def to_vtt(self) -> str:
        """Export as WebVTT subtitle format."""
        lines = ["WEBVTT", ""]
        for seg in self.segments:
            start = self._vtt_time(seg.start)
            end = self._vtt_time(seg.end)
            lines.append(f"{start} --> {end}")
            lines.append(seg.text)
            lines.append("")
        return "\n".join(lines)

    def _vtt_time(self, seconds: float) -> str:
        """Format time for VTT (HH:MM:SS.mmm)."""
        hours, remainder = divmod(int(seconds), 3600)
        minutes, secs = divmod(remainder, 60)
        millis = int((seconds % 1) * 1000)
        return f"{hours:02d}:{minutes:02d}:{secs:02d}.{millis:03d}"

    def to_dict(self) -> Dict[str, Any]:
        return {
            "video_title": self.video_title,
            "language": self.language,
            "duration": self.duration,
            "segments": [s.to_dict() for s in self.segments],
        }

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "Transcript":
        return cls(
            video_title=data["video_title"],
            language=data["language"],
            duration=data["duration"],
            segments=[TranscriptSegment.from_dict(s) for s in data["segments"]],
        )

    def save(self, path: Path):
        """Save transcript to JSON file."""
        path.write_text(json.dumps(self.to_dict(), indent=2))

    @classmethod
    def from_file(cls, path: Path) -> "Transcript":
        """Load transcript from JSON file."""
        data = json.loads(path.read_text())
        return cls.from_dict(data)


# Whisper model sizes and their requirements
WHISPER_MODELS = {
    "tiny": {"size": "39M", "vram": "~1GB", "speed": "fastest"},
    "base": {"size": "74M", "vram": "~1GB", "speed": "fast"},
    "small": {"size": "244M", "vram": "~2GB", "speed": "moderate"},
    "medium": {"size": "769M", "vram": "~5GB", "speed": "slow"},
    "large": {"size": "1550M", "vram": "~10GB", "speed": "slowest"},
    "large-v2": {"size": "1550M", "vram": "~10GB", "speed": "slowest"},
    "large-v3": {"size": "1550M", "vram": "~10GB", "speed": "slowest"},
}


class Transcriber:
    """
    Video transcription using OpenAI Whisper.

    Whisper runs entirely locally, providing:
    - Speech-to-text transcription
    - Word-level timestamps
    - Multi-language support
    - Translation to English
    """

    def __init__(self, config):
        self.config = config
        self.model_name = config.whisper_model
        self._model = None

    def _load_model(self):
        """Load Whisper model (lazy loading)."""
        global HAS_WHISPER
        if self._model is None:
            if not HAS_WHISPER:
                from familiar.core.deps import WHISPER_PACKAGES, ensure_packages

                ok, _ = ensure_packages(WHISPER_PACKAGES)
                if ok:
                    HAS_WHISPER = True
                else:
                    raise ImportError("Whisper transcription not available. Check server logs.")

            logger.info(f"Loading Whisper model: {self.model_name}")
            self._model = whisper.load_model(self.model_name)
            logger.info("Whisper model loaded")

        return self._model

    async def transcribe(
        self,
        video_path: str,
        video_title: str = "Unknown",
        language: Optional[str] = None,
        translate: bool = False,
    ) -> Transcript:
        """
        Transcribe a video file.

        Args:
            video_path: Path to video file
            video_title: Title for the transcript
            language: Language code (auto-detect if None)
            translate: Translate to English

        Returns:
            Transcript object
        """
        # Extract audio if needed
        audio_path = await self._extract_audio(video_path)

        try:
            # Run transcription
            result = await self._run_whisper(
                audio_path,
                language=language,
                translate=translate,
            )

            # Build transcript
            segments = []
            for i, seg in enumerate(result["segments"]):
                segments.append(
                    TranscriptSegment(
                        id=i,
                        start=seg["start"],
                        end=seg["end"],
                        text=seg["text"].strip(),
                    )
                )

            # Calculate duration
            duration = segments[-1].end if segments else 0

            return Transcript(
                video_title=video_title,
                language=result.get("language", "en"),
                duration=duration,
                segments=segments,
            )

        finally:
            # Cleanup temp audio file
            if audio_path != video_path and os.path.exists(audio_path):
                os.remove(audio_path)

    async def _extract_audio(self, video_path: str) -> str:
        """Extract audio from video using ffmpeg."""
        # Check if already audio
        if video_path.endswith((".mp3", ".wav", ".m4a", ".flac", ".ogg")):
            return video_path

        # Extract audio with ffmpeg
        fd, audio_path = tempfile.mkstemp(suffix=".wav")
        os.close(fd)

        cmd = [
            "ffmpeg",
            "-i",
            video_path,
            "-vn",  # No video
            "-acodec",
            "pcm_s16le",
            "-ar",
            "16000",  # 16kHz for Whisper
            "-ac",
            "1",  # Mono
            "-y",  # Overwrite
            audio_path,
        ]

        loop = asyncio.get_event_loop()

        def run():
            result = subprocess.run(
                cmd,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True,
            )
            if result.returncode != 0 and result.stderr:
                logger.debug(f"ffmpeg stderr: {result.stderr[-400:]}")
            return result.returncode == 0

        success = await loop.run_in_executor(None, run)

        if not success:
            # Try without ffmpeg (Whisper can handle some formats)
            return video_path

        return audio_path

    async def _run_whisper(
        self,
        audio_path: str,
        language: Optional[str] = None,
        translate: bool = False,
    ) -> Dict[str, Any]:
        """Run Whisper transcription."""
        loop = asyncio.get_event_loop()

        def transcribe():
            model = self._load_model()

            options = {
                "verbose": False,
            }

            if language:
                options["language"] = language

            if translate:
                options["task"] = "translate"

            return model.transcribe(audio_path, **options)

        return await loop.run_in_executor(None, transcribe)

    async def transcribe_url(
        self,
        url: str,
        video_title: str = "Unknown",
    ) -> Transcript:
        """
        Transcribe a video from URL.

        Downloads to temp file first.
        """
        # This would use VideoExtractor to download
        # For now, raise not implemented
        raise NotImplementedError(
            "Direct URL transcription not implemented. Please download the video first."
        )

    def get_model_info(self) -> Dict[str, Any]:
        """Get information about the current model."""
        info = WHISPER_MODELS.get(self.model_name, {})
        return {
            "name": self.model_name,
            "loaded": self._model is not None,
            **info,
        }

    @staticmethod
    def available_models() -> List[str]:
        """List available Whisper models."""
        return list(WHISPER_MODELS.keys())


class WhisperCLI:
    """
    Alternative transcription using whisper CLI.

    Useful if the Python library has issues or for
    running in a subprocess.
    """

    @staticmethod
    async def transcribe(
        audio_path: str,
        model: str = "base",
        output_format: str = "json",
    ) -> Dict[str, Any]:
        """Run whisper CLI."""
        cmd = [
            "whisper",
            audio_path,
            "--model",
            model,
            "--output_format",
            output_format,
            "--verbose",
            "False",
        ]

        loop = asyncio.get_event_loop()

        def run():
            result = subprocess.run(
                cmd,
                capture_output=True,
                text=True,
            )
            return {
                "returncode": result.returncode,
                "stdout": result.stdout,
                "stderr": result.stderr,
            }

        result = await loop.run_in_executor(None, run)

        if result["returncode"] != 0:
            raise Exception(f"Whisper CLI failed: {result['stderr']}")

        # Parse output JSON
        output_file = audio_path.rsplit(".", 1)[0] + ".json"
        if os.path.exists(output_file):
            with open(output_file) as f:
                return json.load(f)

        return {}
