app.py

import os
import uuid
import asyncio
from datetime import datetime
from typing import List, Optional, Dict, Any
from pathlib import Path
import shutil
import json

from fastapi import FastAPI, File, UploadFile, HTTPException, Depends, BackgroundTasks
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from pydantic import BaseModel, Field
import uvicorn

# Audio processing libraries
import librosa
import numpy as np
from pyannote.audio import Pipeline
from pyannote.audio.pipelines.speaker_verification import PretrainedSpeakerEmbedding
import torch
import torchaudio
from scipy.spatial.distance import cosine

# Database (using SQLite for simplicity - can be replaced with PostgreSQL/MongoDB)
import sqlite3
from contextlib import contextmanager


# Configuration
class Config:
    UPLOAD_DIR = Path("uploads")
    VOICE_LIBRARY_DIR = Path("voice_library")
    DB_PATH = "voice_service.db"
    MAX_FILE_SIZE = 100 * 1024 * 1024  # 100MB
    ALLOWED_EXTENSIONS = {".wav", ".mp3", ".flac", ".m4a", ".ogg"}
    SIMILARITY_THRESHOLD = 0.7


config = Config()

# Ensure directories exist
config.UPLOAD_DIR.mkdir(exist_ok=True)
config.VOICE_LIBRARY_DIR.mkdir(exist_ok=True)


# Pydantic models
class SpeakerSegment(BaseModel):
    speaker_id: str
    start_time: float
    end_time: float
    confidence: float
    text: Optional[str] = None


class DiarizationResult(BaseModel):
    file_id: str
    filename: str
    duration: float
    num_speakers: int
    segments: List[SpeakerSegment]
    processing_time: float
    created_at: datetime


class VoiceProfile(BaseModel):
    profile_id: str
    name: str
    description: Optional[str] = None
    audio_files: List[str]
    embedding: Optional[List[float]] = None
    created_at: datetime
    updated_at: datetime


class UploadResponse(BaseModel):
    file_id: str
    filename: str
    size: int
    duration: float
    status: str


# Database setup
def init_db():
    with sqlite3.connect(config.DB_PATH) as conn:
        conn.execute("""
            CREATE TABLE IF NOT EXISTS audio_files (
                id TEXT PRIMARY KEY,
                filename TEXT NOT NULL,
                file_path TEXT NOT NULL,
                size INTEGER NOT NULL,
                duration REAL NOT NULL,
                status TEXT NOT NULL,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )
        """)

        conn.execute("""
            CREATE TABLE IF NOT EXISTS diarization_results (
                id TEXT PRIMARY KEY,
                file_id TEXT NOT NULL,
                result_data TEXT NOT NULL,
                processing_time REAL NOT NULL,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                FOREIGN KEY (file_id) REFERENCES audio_files (id)
            )
        """)

        conn.execute("""
            CREATE TABLE IF NOT EXISTS voice_profiles (
                id TEXT PRIMARY KEY,
                name TEXT NOT NULL,
                description TEXT,
                embedding TEXT,
                audio_files TEXT,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )
        """)

        conn.commit()


@contextmanager
def get_db():
    conn = sqlite3.connect(config.DB_PATH)
    conn.row_factory = sqlite3.Row
    try:
        yield conn
    finally:
        conn.close()


# Audio processing utilities
class AudioProcessor:
    def __init__(self):
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        try:
            # Initialize speaker diarization pipeline
            self.diarization_pipeline = Pipeline.from_pretrained(
                "pyannote/speaker-diarization-3.1",
                use_auth_token=os.getenv("HUGGINGFACE_TOKEN")  # You'll need to set this
            )

            # Initialize speaker embedding model
            self.embedding_model = PretrainedSpeakerEmbedding(
                "speechbrain/spkrec-ecapa-voxceleb",
                device=self.device
            )
        except Exception as e:
            print(f"Warning: Could not initialize audio models: {e}")
            self.diarization_pipeline = None
            self.embedding_model = None

    def get_audio_info(self, file_path: str) -> Dict[str, Any]:
        """Extract basic audio information"""
        try:
            audio, sr = librosa.load(file_path)
            duration = len(audio) / sr
            return {
                "duration": duration,
                "sample_rate": sr,
                "channels": 1 if len(audio.shape) == 1 else audio.shape[0]
            }
        except Exception as e:
            raise HTTPException(status_code=400, detail=f"Invalid audio file: {str(e)}")

    def perform_diarization(self, file_path: str) -> DiarizationResult:
        """Perform speaker diarization on audio file"""
        if not self.diarization_pipeline:
            raise HTTPException(status_code=503, detail="Diarization model not available")

        start_time = datetime.now()

        try:
            # Load audio
            waveform, sample_rate = torchaudio.load(file_path)

            # Perform diarization
            diarization = self.diarization_pipeline(file_path)

            # Extract segments
            segments = []
            speakers = set()

            for turn, _, speaker in diarization.itertracks(yield_label=True):
                speakers.add(speaker)
                segments.append(SpeakerSegment(
                    speaker_id=speaker,
                    start_time=turn.start,
                    end_time=turn.end,
                    confidence=0.9  # Placeholder confidence
                ))

            audio_info = self.get_audio_info(file_path)
            processing_time = (datetime.now() - start_time).total_seconds()

            return DiarizationResult(
                file_id=str(uuid.uuid4()),
                filename=Path(file_path).name,
                duration=audio_info["duration"],
                num_speakers=len(speakers),
                segments=segments,
                processing_time=processing_time,
                created_at=datetime.now()
            )

        except Exception as e:
            raise HTTPException(status_code=500, detail=f"Diarization failed: {str(e)}")

    def extract_speaker_embedding(self, file_path: str, start_time: float = None, end_time: float = None) -> np.ndarray:
        """Extract speaker embedding from audio segment"""
        if not self.embedding_model:
            raise HTTPException(status_code=503, detail="Embedding model not available")

        try:
            # Load audio
            audio, sr = librosa.load(file_path, sr=16000)

            # Extract segment if specified
            if start_time is not None and end_time is not None:
                start_sample = int(start_time * sr)
                end_sample = int(end_time * sr)
                audio = audio[start_sample:end_sample]

            # Convert to tensor
            audio_tensor = torch.tensor(audio).unsqueeze(0).to(self.device)

            # Extract embedding
            with torch.no_grad():
                embedding = self.embedding_model(audio_tensor)
                return embedding.cpu().numpy().flatten()

        except Exception as e:
            raise HTTPException(status_code=500, detail=f"Embedding extraction failed: {str(e)}")


# Initialize audio processor
audio_processor = AudioProcessor()

# FastAPI app
app = FastAPI(
    title="Voice Processing Service",
    description="Backend service for voice file upload, speaker diarization, and voice library management",
    version="1.0.0"
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


# Initialize database on startup
@app.on_event("startup")
async def startup_event():
    init_db()


# Utility functions
def validate_audio_file(file: UploadFile) -> bool:
    """Validate uploaded audio file"""
    if not file.filename:
        return False

    file_ext = Path(file.filename).suffix.lower()
    return file_ext in config.ALLOWED_EXTENSIONS


async def save_uploaded_file(file: UploadFile) -> tuple[str, str]:
    """Save uploaded file and return file_id and file_path"""
    file_id = str(uuid.uuid4())
    file_ext = Path(file.filename).suffix
    file_path = config.UPLOAD_DIR / f"{file_id}{file_ext}"

    with open(file_path, "wb") as buffer:
        shutil.copyfileobj(file.file, buffer)

    return file_id, str(file_path)


# API Endpoints

@app.post("/upload", response_model=UploadResponse)
async def upload_audio_file(file: UploadFile = File(...)):
    """Upload audio file for processing"""

    # Validate file
    if not validate_audio_file(file):
        raise HTTPException(status_code=400, detail="Invalid file type")

    if file.size > config.MAX_FILE_SIZE:
        raise HTTPException(status_code=400, detail="File too large")

    try:
        # Save file
        file_id, file_path = await save_uploaded_file(file)

        # Get audio info
        audio_info = audio_processor.get_audio_info(file_path)

        # Save to database
        with get_db() as conn:
            conn.execute("""
                INSERT INTO audio_files (id, filename, file_path, size, duration, status)
                VALUES (?, ?, ?, ?, ?, ?)
            """, (file_id, file.filename, file_path, file.size, audio_info["duration"], "uploaded"))
            conn.commit()

        return UploadResponse(
            file_id=file_id,
            filename=file.filename,
            size=file.size,
            duration=audio_info["duration"],
            status="uploaded"
        )

    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Upload failed: {str(e)}")


@app.post("/diarize/{file_id}")
async def perform_diarization(file_id: str, background_tasks: BackgroundTasks):
    """Perform speaker diarization on uploaded file"""

    # Get file info from database
    with get_db() as conn:
        row = conn.execute("SELECT * FROM audio_files WHERE id = ?", (file_id,)).fetchone()
        if not row:
            raise HTTPException(status_code=404, detail="File not found")

    file_path = row["file_path"]

    # Check if file exists
    if not Path(file_path).exists():
        raise HTTPException(status_code=404, detail="Audio file not found on disk")

    def process_diarization():
        try:
            # Perform diarization
            result = audio_processor.perform_diarization(file_path)

            # Save result to database
            with get_db() as conn:
                conn.execute("""
                    INSERT INTO diarization_results (id, file_id, result_data, processing_time)
                    VALUES (?, ?, ?, ?)
                """, (result.file_id, file_id, result.json(), result.processing_time))

                # Update file status
                conn.execute("""
                    UPDATE audio_files SET status = 'processed' WHERE id = ?
                """, (file_id,))

                conn.commit()

        except Exception as e:
            # Update status to failed
            with get_db() as conn:
                conn.execute("""
                    UPDATE audio_files SET status = 'failed' WHERE id = ?
                """, (file_id,))
                conn.commit()
            print(f"Diarization failed for {file_id}: {e}")

    # Start background processing
    background_tasks.add_task(process_diarization)

    # Update status to processing
    with get_db() as conn:
        conn.execute("UPDATE audio_files SET status = 'processing' WHERE id = ?", (file_id,))
        conn.commit()

    return {"message": "Diarization started", "file_id": file_id, "status": "processing"}


@app.get("/diarization/{file_id}", response_model=DiarizationResult)
async def get_diarization_result(file_id: str):
    """Get diarization results for a file"""

    with get_db() as conn:
        row = conn.execute("""
            SELECT * FROM diarization_results WHERE file_id = ?
        """, (file_id,)).fetchone()

        if not row:
            raise HTTPException(status_code=404, detail="Diarization result not found")

    result_data = json.loads(row["result_data"])
    return DiarizationResult(**result_data)


@app.post("/voice-library/profiles", response_model=VoiceProfile)
async def create_voice_profile(
        name: str,
        description: str = None,
        audio_files: List[UploadFile] = File(...)
):
    """Create a new voice profile with audio samples"""

    profile_id = str(uuid.uuid4())
    saved_files = []
    embeddings = []

    try:
        # Process each audio file
        for file in audio_files:
            if not validate_audio_file(file):
                continue

            # Save file to voice library
            file_ext = Path(file.filename).suffix
            file_path = config.VOICE_LIBRARY_DIR / f"{profile_id}_{len(saved_files)}{file_ext}"

            with open(file_path, "wb") as buffer:
                shutil.copyfileobj(file.file, buffer)

            saved_files.append(str(file_path))

            # Extract embedding
            embedding = audio_processor.extract_speaker_embedding(str(file_path))
            embeddings.append(embedding)

        # Average embeddings if multiple files
        if embeddings:
            avg_embedding = np.mean(embeddings, axis=0).tolist()
        else:
            avg_embedding = None

        # Save to database
        with get_db() as conn:
            conn.execute("""
                INSERT INTO voice_profiles (id, name, description, embedding, audio_files)
                VALUES (?, ?, ?, ?, ?)
            """, (profile_id, name, description, json.dumps(avg_embedding), json.dumps(saved_files)))
            conn.commit()

        return VoiceProfile(
            profile_id=profile_id,
            name=name,
            description=description,
            audio_files=saved_files,
            embedding=avg_embedding,
            created_at=datetime.now(),
            updated_at=datetime.now()
        )

    except Exception as e:
        # Cleanup on error
        for file_path in saved_files:
            Path(file_path).unlink(missing_ok=True)
        raise HTTPException(status_code=500, detail=f"Profile creation failed: {str(e)}")


@app.get("/voice-library/profiles")
async def list_voice_profiles():
    """List all voice profiles"""

    with get_db() as conn:
        rows = conn.execute("SELECT * FROM voice_profiles ORDER BY created_at DESC").fetchall()

    profiles = []
    for row in rows:
        profiles.append({
            "profile_id": row["id"],
            "name": row["name"],
            "description": row["description"],
            "audio_files": json.loads(row["audio_files"]) if row["audio_files"] else [],
            "created_at": row["created_at"],
            "updated_at": row["updated_at"]
        })

    return {"profiles": profiles}


@app.post("/voice-library/identify/{file_id}")
async def identify_speakers(file_id: str):
    """Identify speakers in diarized audio using voice library"""

    # Get diarization results
    with get_db() as conn:
        diarization_row = conn.execute("""
            SELECT * FROM diarization_results WHERE file_id = ?
        """, (file_id,)).fetchone()

        if not diarization_row:
            raise HTTPException(status_code=404, detail="Diarization result not found")

        file_row = conn.execute("""
            SELECT * FROM audio_files WHERE id = ?
        """, (file_id,)).fetchone()

        profiles_rows = conn.execute("SELECT * FROM voice_profiles").fetchall()

    diarization_data = json.loads(diarization_row["result_data"])
    file_path = file_row["file_path"]

    # Load voice profiles
    voice_profiles = {}
    for row in profiles_rows:
        if row["embedding"]:
            voice_profiles[row["id"]] = {
                "name": row["name"],
                "embedding": np.array(json.loads(row["embedding"]))
            }

    # Identify speakers for each segment
    identified_segments = []

    for segment in diarization_data["segments"]:
        # Extract embedding for this segment
        segment_embedding = audio_processor.extract_speaker_embedding(
            file_path,
            segment["start_time"],
            segment["end_time"]
        )

        # Find best match
        best_match = None
        best_similarity = 0

        for profile_id, profile_data in voice_profiles.items():
            similarity = 1 - cosine(segment_embedding, profile_data["embedding"])
            if similarity > best_similarity and similarity > config.SIMILARITY_THRESHOLD:
                best_similarity = similarity
                best_match = {
                    "profile_id": profile_id,
                    "name": profile_data["name"],
                    "similarity": similarity
                }

        identified_segments.append({
            **segment,
            "identified_speaker": best_match
        })

    return {
        "file_id": file_id,
        "segments": identified_segments
    }


@app.get("/files")
async def list_files():
    """List all uploaded files"""

    with get_db() as conn:
        rows = conn.execute("""
            SELECT af.*, dr.id as diarization_id 
            FROM audio_files af 
            LEFT JOIN diarization_results dr ON af.id = dr.file_id
            ORDER BY af.created_at DESC
        """).fetchall()

    files = []
    for row in rows:
        files.append({
            "file_id": row["id"],
            "filename": row["filename"],
            "size": row["size"],
            "duration": row["duration"],
            "status": row["status"],
            "has_diarization": bool(row["diarization_id"]),
            "created_at": row["created_at"]
        })

    return {"files": files}


@app.delete("/files/{file_id}")
async def delete_file(file_id: str):
    """Delete uploaded file and associated data"""

    with get_db() as conn:
        # Get file info
        row = conn.execute("SELECT * FROM audio_files WHERE id = ?", (file_id,)).fetchone()
        if not row:
            raise HTTPException(status_code=404, detail="File not found")

        # Delete file from disk
        file_path = Path(row["file_path"])
        file_path.unlink(missing_ok=True)

        # Delete from database
        conn.execute("DELETE FROM diarization_results WHERE file_id = ?", (file_id,))
        conn.execute("DELETE FROM audio_files WHERE id = ?", (file_id,))
        conn.commit()

    return {"message": "File deleted successfully"}


@app.get("/health")
async def health_check():
    """Health check endpoint"""
    return {
        "status": "healthy",
        "timestamp": datetime.now().isoformat(),
        "models_loaded": {
            "diarization": audio_processor.diarization_pipeline is not None,
            "embedding": audio_processor.embedding_model is not None
        }
    }


if __name__ == "__main__":
    uvicorn.run(
        "main:app",
        host="0.0.0.0",
        port=8000,
        reload=True,
        log_level="info"
    )