processor.py

from __future__ import annotations

import json
import os
import time
from pathlib import Path
from typing import Any, Dict, List

from dotenv import load_dotenv
from groq import Groq
from pydub import AudioSegment
from sentence_transformers import SentenceTransformer

from config import (
    AUDIO_DIR,
    DB_DIR,
    DEFAULT_LANGUAGE,
    EMBED_MODEL,
    GROQ_ASR_MODEL,
    GROQ_CHAT_MODEL,
    HF_TOKEN,
    MAX_CHUNK_SECONDS,
    META_FILE,
    MIN_CHUNK_SECONDS,
    QUERY_LOG_FILE,
    VECTOR_FILE,
)

load_dotenv()

api_key = os.getenv("GROQ_API_KEY")
if not api_key:
    raise ValueError("🚨 GROQ_API_KEY not found! Please check your .env file.")

client = Groq(api_key=api_key)
embedder = SentenceTransformer(EMBED_MODEL)


def _ensure_dirs() -> None:
    DB_DIR.mkdir(parents=True, exist_ok=True)
    AUDIO_DIR.mkdir(parents=True, exist_ok=True)
    for path in (META_FILE, VECTOR_FILE, QUERY_LOG_FILE):
        if not path.exists():
            path.write_text("[]", encoding="utf-8")


def _load_json(path: Path) -> list:
    _ensure_dirs()
    try:
        with path.open("r", encoding="utf-8") as f:
            return json.load(f)
    except Exception:
        return []


def _save_json(path: Path, data: list) -> None:
    _ensure_dirs()
    with path.open("w", encoding="utf-8") as f:
        json.dump(data, f, ensure_ascii=False, indent=2)


def _seconds_to_hhmmss(seconds: float) -> str:
    total = max(0, int(seconds))
    h = total // 3600
    m = (total % 3600) // 60
    s = total % 60
    return f"{h:02d}:{m:02d}:{s:02d}"


def _safe_transcribe(file_path: str, language: str = DEFAULT_LANGUAGE) -> str:
    with open(file_path, "rb") as file:
        transcript = client.audio.transcriptions.create(
            file=(os.path.basename(file_path), file.read()),
            model=GROQ_ASR_MODEL,
            language=language,
            response_format="text",
        )
    return transcript.strip()


def _safe_chat(prompt: str, temperature: float = 0.2) -> str:
    response = client.chat.completions.create(
        model=GROQ_CHAT_MODEL,
        messages=[{"role": "user", "content": prompt}],
        temperature=temperature,
    )
    return response.choices[0].message.content.strip()


def _get_diarization_segments(audio_path: str, duration_sec: float) -> List[Dict[str, Any]]:
    """
    Returns speaker segments. Falls back to a single-speaker segment if pyannote or HF token is unavailable.
    """
    if not HF_TOKEN:
        return [{"speaker": "Speaker 1", "start": 0.0, "end": float(duration_sec)}]

    try:
        from pyannote.audio import Pipeline
    except Exception:
        return [{"speaker": "Speaker 1", "start": 0.0, "end": float(duration_sec)}]

    try:
        pipeline = Pipeline.from_pretrained(
            "pyannote/speaker-diarization-3.1",
            use_auth_token=HF_TOKEN,
        )
        diarization = pipeline(audio_path)
        segments: List[Dict[str, Any]] = []
        speaker_map: Dict[str, str] = {}
        speaker_counter = 1

        for turn, _, speaker_label in diarization.itertracks(yield_label=True):
            mapped = speaker_map.get(speaker_label)
            if mapped is None:
                mapped = f"Speaker {speaker_counter}"
                speaker_map[speaker_label] = mapped
                speaker_counter += 1
            segments.append(
                {
                    "speaker": mapped,
                    "start": float(turn.start),
                    "end": float(turn.end),
                }
            )

        segments = [s for s in segments if s["end"] > s["start"]]
        segments.sort(key=lambda x: (x["start"], x["end"]))
        return segments or [{"speaker": "Speaker 1", "start": 0.0, "end": float(duration_sec)}]
    except Exception:
        return [{"speaker": "Speaker 1", "start": 0.0, "end": float(duration_sec)}]


def _normalize_segments(segments: List[Dict[str, Any]], duration_sec: float) -> List[Dict[str, Any]]:
    if not segments:
        return [{"speaker": "Speaker 1", "start": 0.0, "end": float(duration_sec)}]

    normalized = []
    last_end = 0.0
    for seg in segments:
        start = max(0.0, min(float(seg["start"]), duration_sec))
        end = max(0.0, min(float(seg["end"]), duration_sec))
        if end <= start:
            continue
        if start > last_end and start - last_end > 0.75:
            normalized.append({"speaker": "Speaker 1", "start": last_end, "end": start})
        normalized.append({"speaker": seg["speaker"], "start": start, "end": end})
        last_end = max(last_end, end)

    if not normalized:
        return [{"speaker": "Speaker 1", "start": 0.0, "end": float(duration_sec)}]
    return normalized


def _split_segment(start_sec: float, end_sec: float, max_len_sec: int = MAX_CHUNK_SECONDS) -> List[tuple[float, float]]:
    parts = []
    current = start_sec
    while current < end_sec:
        nxt = min(current + max_len_sec, end_sec)
        if nxt - current >= MIN_CHUNK_SECONDS:
            parts.append((current, nxt))
        current = nxt
    return parts


def extract_chunk_facts(chunk_text: str, speaker: str = "Speaker") -> str:
    prompt = f"""
أنت مساعد دقيق جداً. استخرج أهم المعلومات، الأرقام، والأسماء والحقائق من هذا المقطع الصوتي القصير.
تجاهل الحشو والكلام الجانبي، وركز فقط على الجوهر.
إذا كان المقطع لا يحتوي على معلومات مهمة، اكتب: لا توجد معلومات جوهرية.

اكتب الإجابة في شكل نقاط قصيرة ومباشرة.
المتحدث: {speaker}

النص:
{chunk_text}
"""
    return _safe_chat(prompt, temperature=0.1)


def merge_extracted_facts(all_facts: str) -> tuple[str, str]:
    prompt = f"""
أنت خبير في التلخيص المتقدم.
إليك قائمة بالنقاط التي تم استخراجها من تسجيل صوتي.

مهمتك هي التوليف وليس النسخ.
الشروط:
1) اكتب من 5 إلى 7 نقاط رئيسية كحد أقصى.
2) ادمج الحقائق المترابطة في نقطة واحدة كثيفة.
3) احذف التكرار والحشو.
4) استخرج عنواناً دقيقاً يعبر عن جوهر التسجيل بالكامل.

أخرج النتيجة بهذا الشكل فقط:
العنوان: ...
النقاط:
- ...
- ...

النقاط الخام:
{all_facts}
"""
    output = _safe_chat(prompt, temperature=0.2)
    lines = output.splitlines()
    title = "بدون عنوان"
    points = output.strip()
    for i, line in enumerate(lines):
        if line.startswith("العنوان:"):
            title = line.replace("العنوان:", "").strip()
        if line.startswith("النقاط:"):
            points = "\\n".join(lines[i + 1:]).strip()
    return title or "بدون عنوان", points or output.strip()


def process_and_vectorize_audio(audio_path: str) -> Dict[str, Any]:
    """
    End-to-end ingestion:
    - speaker diarization (optional)
    - chunking
    - transcription
    - fact extraction
    - summary synthesis
    - vector indexing
    """
    _ensure_dirs()

    start_time = time.time()
    audio = AudioSegment.from_file(audio_path).set_channels(1).set_frame_rate(16000)
    duration_sec = len(audio) / 1000.0
    base_name = Path(audio_path).stem

    diarized_segments = _normalize_segments(_get_diarization_segments(audio_path, duration_sec), duration_sec)

    full_raw_text_parts: List[str] = []
    all_chunks_facts_parts: List[str] = []
    new_vectors: List[Dict[str, Any]] = []
    chunk_records: List[Dict[str, Any]] = []

    chunk_counter = 0
    turn_counter = 0

    for turn in diarized_segments:
        turn_counter += 1
        speaker = turn["speaker"]

        for start_sec, end_sec in _split_segment(turn["start"], turn["end"], MAX_CHUNK_SECONDS):
            start_ms = int(start_sec * 1000)
            end_ms = int(end_sec * 1000)
            if end_ms <= start_ms:
                continue

            chunk_audio = audio[start_ms:end_ms]
            if len(chunk_audio) < MIN_CHUNK_SECONDS * 1000:
                continue

            chunk_counter += 1
            chunk_filename = f"{base_name}_spk{speaker.replace(' ', '_')}_{start_ms}_{end_ms}.mp3"
            chunk_path = AUDIO_DIR / chunk_filename
            chunk_audio.export(chunk_path, format="mp3")

            try:
                text = _safe_transcribe(str(chunk_path), language=DEFAULT_LANGUAGE)
            except Exception as exc:
                text = ""
                print(f"[WARN] transcription failed for {chunk_filename}: {exc}")

            if not text:
                continue

            full_raw_text_parts.append(text)
            chunk_facts = extract_chunk_facts(text, speaker=speaker)
            all_chunks_facts_parts.append(
                f"- ({speaker} | {_seconds_to_hhmmss(start_sec)} → {_seconds_to_hhmmss(end_sec)}) {chunk_facts}"
            )

            embedding = embedder.encode(text, normalize_embeddings=True).tolist()

            record = {
                "file_link": audio_path,
                "title": base_name,
                "speaker": speaker,
                "segment_index": turn_counter,
                "chunk_index": chunk_counter,
                "chunk_order": chunk_counter,
                "audio_chunk_path": str(chunk_path),
                "start_time": round(start_sec, 3),
                "end_time": round(end_sec, 3),
                "start_time_hms": _seconds_to_hhmmss(start_sec),
                "end_time_hms": _seconds_to_hhmmss(end_sec),
                "turn_start_time": round(turn["start"], 3),
                "turn_end_time": round(turn["end"], 3),
                "turn_start_hms": _seconds_to_hhmmss(turn["start"]),
                "turn_end_hms": _seconds_to_hhmmss(turn["end"]),
                "text_chunk": text,
                "facts": chunk_facts,
                "embedding": embedding,
                "chunk_duration_sec": round(end_sec - start_sec, 3),
            }
            new_vectors.append(record)
            chunk_records.append(record)

    title, final_summary_points = merge_extracted_facts("".join(all_chunks_facts_parts))

    metadata = {
        "file_link": audio_path,
        "title": title,
        "summary": final_summary_points,
        "full_text": " ".join(full_raw_text_parts).strip(),
        "audio_duration_sec": round(duration_sec, 3),
        "audio_duration_hms": _seconds_to_hhmmss(duration_sec),
        "chunk_count": len(chunk_records),
        "speaker_turn_count": len(diarized_segments),
        "speakers": sorted({item["speaker"] for item in chunk_records}) or ["Speaker 1"],
        "processing_time_sec": round(time.time() - start_time, 3),
        "created_at": time.strftime("%Y-%m-%d %H:%M:%S"),
        "chunk_records": [
            {
                "speaker": item["speaker"],
                "start_time": item["start_time"],
                "end_time": item["end_time"],
                "start_time_hms": item["start_time_hms"],
                "end_time_hms": item["end_time_hms"],
                "audio_chunk_path": item["audio_chunk_path"],
                "text_chunk": item["text_chunk"],
            }
            for item in chunk_records
        ],
    }

    all_meta = _load_json(META_FILE)
    all_meta.append(metadata)
    _save_json(META_FILE, all_meta)

    all_vectors = _load_json(VECTOR_FILE)
    all_vectors.extend(new_vectors)
    _save_json(VECTOR_FILE, all_vectors)

    return metadata