main.py

import logging
import os
import re
import asyncio
import subprocess
import tempfile
import time
from collections import Counter, defaultdict
from pathlib import Path
from typing import DefaultDict, List, Optional, Tuple

import httpx

import google.generativeai as genai
from dotenv import load_dotenv
from fastapi import FastAPI, HTTPException, UploadFile, File, Form
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field, ConfigDict
from transformers import Pipeline, pipeline
from faster_whisper import WhisperModel
import torch

# Load environment variables from a local .env if present.
load_dotenv()

logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO)

app = FastAPI(title="Meeting Summarizer", version="0.1.0")

cors_origins = [
    origin.strip()
    for origin in os.getenv("CORS_ORIGINS", "*").split(",")
    if origin.strip()
]

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

# Simple stopword list to de-prioritize filler words (English + common Korean particles).
DEFAULT_STOPWORDS = {
    "the",
    "a",
    "an",
    "and",
    "or",
    "is",
    "are",
    "was",
    "were",
    "to",
    "in",
    "for",
    "on",
    "of",
    "with",
    "as",
    "at",
    "it",
    "this",
    "that",
    "by",
    "from",
    "be",
    "we",
    "you",
    "they",
    "i",
    "our",
    "your",
    "their",
    "but",
    "if",
    "not",
    "so",
    "do",
    "did",
    "does",
    "have",
    "has",
    "had",
    "been",
    "can",
    "could",
    "should",
    "would",
    "will",
    "about",
    "into",
    "out",
    "up",
    "down",
    "over",
    "under",
    "than",
    "then",
    "there",
    "here",
    "그",
    "그냥",
    "그리고",
    "그래서",
    "이거",
    "이런",
    "저희",
}

HF_MODEL_ID = os.getenv("HF_MODEL_ID", "philschmid/distilbart-cnn-12-6-samsum")
HF_MAX_CHUNK_WORDS = 900

hf_summarizer: Optional[Pipeline] = None
gemini_model: Optional[genai.GenerativeModel] = None
local_whisper_model: Optional[WhisperModel] = None


class SummarizeRequest(BaseModel):
    transcript: str = Field(..., description="STT output of the meeting")
    max_keywords: int = Field(10, ge=3, le=30, description="Number of keywords to keep")


class SummarizeResponse(BaseModel):
    keywords: List[str]
    gemini_summary: Optional[str]



class SttStreamResponse(BaseModel):
    success: bool
    userId: str
    userName: str
    roomId: str
    transcript: Optional[str] = None
    confidence: Optional[float] = None
    timestamp: Optional[int] = None
    error: Optional[str] = None
    detail: Optional[str] = None


class TranscriptEntry(BaseModel):
    room_id: str = Field(..., alias="roomId")
    user_id: str = Field(..., alias="userId")
    user_name: str = Field(..., alias="userName")
    transcript: str
    timestamp: int
    confidence: Optional[float] = None
    model_config = ConfigDict(populate_by_name=True)


class TranscriptListResponse(BaseModel):
    roomId: str
    entries: List[TranscriptEntry]
    fullTranscript: str
    model_config = ConfigDict(populate_by_name=True)


transcripts_store: DefaultDict[str, List[TranscriptEntry]] = defaultdict(list)

SPRING_SERVER_URL = os.getenv("SPRING_SERVER_URL", "http://localhost:8080")


async def notify_spring_server(room_id: str, user_id: str, user_name: str, transcript: str) -> None:
    """STT 결과를 Spring 서버로 WebSocket 브로드캐스트 요청."""
    payload = {
        "type": "stt",
        "roomId": room_id,
        "userId": user_id,
        "userName": user_name,
        "transcript": transcript,
        "timestamp": int(time.time() * 1000),
    }
    try:
        async with httpx.AsyncClient(timeout=5.0) as client:
            response = await client.post(f"{SPRING_SERVER_URL}/api/stt/broadcast", json=payload)
            if response.status_code == 200:
                logger.info("✅ Broadcasted STT to Spring: %s", user_name)
            else:
                logger.warning("⚠️ Spring broadcast failed (%s): %s", response.status_code, response.text)
    except Exception as exc:  # noqa: BLE001
        logger.error("❌ Failed to notify Spring server: %s", exc)


def extract_keywords(text: str, max_keywords: int) -> List[str]:
    """Return top-N keywords based on frequency, ignoring stopwords."""
    tokens = re.findall(r"[A-Za-z0-9가-힣]+", text.lower())
    filtered = [tok for tok in tokens if tok not in DEFAULT_STOPWORDS and len(tok) > 1]
    counts = Counter(filtered)
    return [word for word, _ in counts.most_common(max_keywords)]


def chunk_text(text: str, chunk_size: int = HF_MAX_CHUNK_WORDS) -> List[str]:
    words = text.split()
    return [" ".join(words[i : i + chunk_size]) for i in range(0, len(words), chunk_size)]


def get_hf_summarizer() -> Optional[Pipeline]:
    global hf_summarizer
    if hf_summarizer is not None:
        return hf_summarizer

    token = os.getenv("HUGGING_FACE_ACCESS_KEY")
    device = 0 if torch.cuda.is_available() else -1
    try:
        kwargs = {"model": HF_MODEL_ID}
        if token:
            kwargs["token"] = token
        if device >= 0:
            kwargs["device"] = device
        hf_summarizer = pipeline("summarization", **kwargs)
        return hf_summarizer
    except Exception as exc:  # noqa: BLE001
        logger.exception("Failed to initialize Hugging Face summarizer: %s", exc)
        return None


def summarize_with_hf(text: str, min_length: int, max_length: int) -> Optional[str]:
    summarizer = get_hf_summarizer()
    if summarizer is None:
        return None

    try:
        chunks = chunk_text(text)
        if not chunks:
            return None
        partial_summaries = []
        for chunk in chunks:
            adjusted_max = max(min_length, min(max_length, len(chunk.split())))
            if len(chunk.split()) < max_length:
                adjusted_max = max(min_length, len(chunk.split()))
            partial = summarizer(
                chunk,
                min_length=min_length,
                max_length=adjusted_max,
                do_sample=False,
            )[0]["summary_text"]
            partial_summaries.append(partial)

        if len(partial_summaries) == 1:
            return partial_summaries[0]

        combined = " ".join(partial_summaries)
        adjusted_max = max(min_length, min(max_length, len(combined.split())))
        if len(combined.split()) < max_length:
            adjusted_max = max(min_length, len(combined.split()))
        final = summarizer(
            combined,
            min_length=min_length,
            max_length=adjusted_max,
            do_sample=False,
        )[0]["summary_text"]
        return final
    except Exception as exc:  # noqa: BLE001
        logger.exception("Hugging Face summarization failed: %s", exc)
        return None


def get_local_whisper_model() -> Optional[WhisperModel]:
    global local_whisper_model
    if local_whisper_model is not None:
        return local_whisper_model

    model_size = os.getenv("WHISPER_LOCAL_MODEL", "small")
    device = os.getenv("WHISPER_LOCAL_DEVICE", "cpu")
    compute_type = os.getenv("WHISPER_LOCAL_COMPUTE", "int8")
    try:
        local_whisper_model = WhisperModel(
            model_size,
            device=device,
            compute_type=compute_type,
        )
        return local_whisper_model
    except Exception as exc:  # noqa: BLE001
        logger.exception("Failed to load faster-whisper model (%s): %s", model_size, exc)
        return None


async def transcribe_with_whisper(
    audio_bytes: bytes,
    filename: Optional[str],
) -> Tuple[Optional[str], Optional[str]]:
    return await transcribe_with_local_whisper(audio_bytes, filename)


async def transcribe_with_local_whisper(
    audio_bytes: bytes,
    filename: Optional[str],
) -> Tuple[Optional[str], Optional[str]]:
    def convert_audio_to_wav(src_path: str) -> Tuple[str, bool]:
        with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as dst:
            dst_path = dst.name
        try:
            subprocess.run(
                [
                    "ffmpeg",
                    "-y",
                    "-loglevel",
                    "error",
                    "-i",
                    src_path,
                    "-ar",
                    "16000",
                    "-ac",
                    "1",
                    "-c:a",
                    "pcm_s16le",
                    "-f",
                    "wav",
                    dst_path,
                ],
                check=True,
                stdout=subprocess.DEVNULL,
                stderr=subprocess.PIPE,
            )
            return dst_path, True
        except subprocess.CalledProcessError as exc:  # noqa: BLE001
            error_msg = exc.stderr.decode(errors="ignore") if exc.stderr else str(exc)
            logger.warning(
                "FFmpeg conversion failed (%s); using original audio.",
                error_msg,
            )
            try:
                os.remove(dst_path)
            except OSError:
                pass
            return src_path, False
        except Exception as exc:  # noqa: BLE001
            logger.warning("FFmpeg failed (%s); using original audio.", exc)
            try:
                os.remove(dst_path)
            except OSError:
                pass
            return src_path, False
            try:
                os.remove(dst_path)
            except OSError:
                pass
            return src_path, False

    model = get_local_whisper_model()
    if model is None:
        return None, "local_whisper_unavailable"

    suffix = Path(filename or "").suffix or ".webm"
    with tempfile.NamedTemporaryFile(delete=True, suffix=suffix) as tmp:
        tmp.write(audio_bytes)
        tmp.flush()
        tmp.seek(0)
        converted_path, converted = await asyncio.to_thread(convert_audio_to_wav, tmp.name)
        try:
            segments, _ = await asyncio.to_thread(
                model.transcribe,
                converted_path,
                language=os.getenv("WHISPER_LANGUAGE", "ko"),
            )
            text = " ".join(segment.text.strip() for segment in segments if segment.text)
            if text:
                return text, None
            if not converted:
                return None, "local_whisper_empty"
            return None, "local_whisper_empty_after_conversion"
        except Exception as exc:  # noqa: BLE001
            logger.exception("Local Whisper transcription failed: %s", exc)
            return None, str(exc)
        finally:
            if converted_path != tmp.name and Path(converted_path).exists():
                Path(converted_path).unlink(missing_ok=True)

    return None, "local_whisper_failed"


def summarize_with_gemini(transcript: str, keywords: List[str]) -> Optional[str]:

    if len(transcript.strip()) < 20 or len(transcript.split()) < 3:
        return "요약할 만큼 충분한 회의 내용이 없습니다. 조금 더 긴 발화를 입력해주세요."

    api_key = os.getenv("GEMINI_API_KEY")
    if not api_key:
        logger.error("GEMINI_API_KEY not configured.")
        return None
    logger.info("GEMINI_API_KEY found (length: %d)", len(api_key))

    try:
        genai.configure(api_key=api_key)
        logger.info("Gemini configured.")
    except Exception as exc:  # noqa: BLE001
        logger.exception("Failed to configure Gemini client: %s", exc)
        return None

    trimmed = transcript[:5000]
    prompt = f"""
다음 회의 내용을 한국어로 핵심만 간단하게 요약해줘.
주요 논의 내용만 간결하게 작성하고, 마크다운 강조 표시(**, -, # 등)는 사용하지 마.
제목이나 섹션 구분 없이 본문만 작성해줘.

[회의 내용]
{trimmed}

[핵심 키워드]
{', '.join(keywords)}
"""

    requested = os.getenv("GEMINI_MODEL") or os.getenv("gemini_model")
    candidates = [
        requested or "gemini-2.5-flash",
        "gemini-2.0-flash",
        "gemini-flash-latest",
    ]
    tried: List[str] = []

    for model_id in candidates:
        if not model_id or model_id in tried:
            continue
        tried.append(model_id)
        try:
            model = genai.GenerativeModel(model_id)
            logger.info("Generating content with model %s...", model_id)
            response = model.generate_content(
                prompt,
                generation_config={"temperature": 0.3, "max_output_tokens": 500},
            )
            logger.info("Content generated.")
            text_resp = getattr(response, "text", None)
            if text_resp:
                logger.info("Gemini summary succeeded with model %s", model_id)
                return text_resp.strip()
        except Exception as exc:  # noqa: BLE001
            error_msg = str(exc)
            if "429" in error_msg or "quota" in error_msg.lower() or "ResourceExhausted" in error_msg:
                logger.error("Gemini API quota exceeded: %s", exc)
                return "⚠️ Gemini API 할당량이 초과되었습니다. 잠시 후 다시 시도해주세요."
            logger.warning("Gemini model %s failed (%s)", model_id, exc)

    logger.error("Gemini summarization failed; tried %s", tried)
    return None

def build_full_transcript(entries: List[TranscriptEntry]) -> str:
    return "\n".join(f"{entry.user_name}: {entry.transcript}" for entry in entries)


@app.post("/stt/stream", response_model=SttStreamResponse)
async def stt_stream(
    audio: UploadFile = File(...),
    userId: str = Form(...),
    userName: str = Form(...),
    roomId: str = Form(...),
    timestamp: Optional[int] = Form(None),
) -> SttStreamResponse:
    ts = int(timestamp) if timestamp is not None else int(time.time() * 1000)
    logger.info(
        "📤 STT 전송 중: %s",
        {
            "userId": userId,
            "userName": userName,
            "roomId": roomId,
            "timestamp": ts,
            "contentType": audio.content_type,
        },
    )
    audio_bytes = await audio.read()
    if not audio_bytes:
        return SttStreamResponse(
            success=False,
            userId=userId,
            userName=userName,
            roomId=roomId,
            transcript="[none]",
            confidence=None,
            timestamp=ts,
            error="empty_audio",
            detail="No audio data received.",
        )

    allowed_bases = {
        "audio/webm",
        "audio/opus",
        "audio/ogg",
        "audio/wav",
        "audio/mpeg",
        "audio/mp3",
    }
    content_type = (audio.content_type or "").split(";")[0].strip().lower()
    if content_type and content_type not in allowed_bases:
        return SttStreamResponse(
            success=False,
            userId=userId,
            userName=userName,
            roomId=roomId,
            transcript="[none]",
            confidence=None,
            timestamp=ts,
            error="unsupported_audio",
            detail=f"Unsupported audio content type: {audio.content_type}",
        )

    transcript, error_detail = await transcribe_with_whisper(audio_bytes, audio.filename)
    if not transcript:
        logger.warning(
            "STT failed for user %s in room %s: %s",
            userName,
            roomId,
            error_detail or "unknown error",
        )
        return SttStreamResponse(
            success=False,
            userId=userId,
            userName=userName,
            roomId=roomId,
            transcript="[none]",
            confidence=None,
            timestamp=ts,
            error="stt_failed",
            detail=error_detail or "STT processing failed.",
        )

    entry = TranscriptEntry(
        room_id=roomId,
        user_id=userId,
        user_name=userName,
        transcript=transcript,
        timestamp=ts,
        confidence=None,
    )
    transcripts_store[roomId].append(entry)

    logger.info(
        "✅ STT 응답: %s",
        {
            "success": True,
            "userId": userId,
            "userName": userName,
            "roomId": roomId,
            "transcript": transcript,
            "confidence": entry.confidence,
            "timestamp": ts,
        },
    )
    logger.info("🗣️ %s: %s", userName, transcript)
    await notify_spring_server(roomId, userId, userName, transcript)

    return SttStreamResponse(
        success=True,
        userId=userId,
        userName=userName,
        roomId=roomId,
        transcript=transcript,
        confidence=entry.confidence,
        timestamp=ts,
    )


@app.get("/stt/transcript/{room_id}", response_model=TranscriptListResponse)
async def get_room_transcript(room_id: str) -> TranscriptListResponse:
    entries = transcripts_store.get(room_id, [])
    full_transcript = build_full_transcript(entries)
    return TranscriptListResponse(roomId=room_id, entries=entries, fullTranscript=full_transcript)


@app.delete("/stt/transcript/{room_id}")
async def delete_room_transcript(room_id: str) -> dict:
    deleted = room_id in transcripts_store
    transcripts_store.pop(room_id, None)
    return {"success": True, "roomId": room_id, "deleted": deleted}


@app.get("/health")
async def health() -> dict:
    return {"status": "ok"}


@app.post("/summaries", response_model=SummarizeResponse)
async def create_summary(payload: SummarizeRequest) -> SummarizeResponse:
    logger.info("Received summary request with transcript length: %d", len(payload.transcript))
    if not payload.transcript.strip():
        raise HTTPException(status_code=400, detail="Transcript is empty.")

    logger.info("Extracting keywords...")
    keywords = extract_keywords(payload.transcript, payload.max_keywords)
    logger.info("Keywords extracted: %s", keywords)
    
    logger.info("Calling summarize_with_gemini...")
    try:
        gemini_summary = await asyncio.wait_for(
            asyncio.to_thread(summarize_with_gemini, payload.transcript, keywords),
            timeout=30.0
        )
        logger.info("summarize_with_gemini returned.")
    except asyncio.TimeoutError:
        logger.error("Gemini summarization timed out after 30s.")
        gemini_summary = "요약 생성 시간이 초과되었습니다. (Gemini API Timeout)"
    except Exception as exc:
        logger.error("Error during summarization: %s", exc)
        gemini_summary = None

    return SummarizeResponse(
        keywords=keywords,
        gemini_summary=gemini_summary,
    )



if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="0.0.0.0", port=int(os.getenv("PORT", 8000)))