call-to-text/old/main.py

import os
import sys
import click
from typing import Optional
from concurrent.futures import ThreadPoolExecutor, as_completed

sys.path.insert(0, os.path.dirname(__file__))

from config import config
from ingestor.email_fetcher import EmailFetcher
from ingestor.parsers.novofon import get_parser
from ingestor.downloader import download_audio
from audio.preprocessor import convert_to_wav, split_stereo_channels
from audio.stt import transcribe
from aligner.aligner import align
from storage.contacts import resolve_contacts, resolve_contact_name
from storage.exporter import export_to_json


def _is_mono_duplicate(wav_path: str, threshold: float = 0.98) -> bool:
    import wave
    import struct

    try:
        with wave.open(wav_path, "rb") as wf:
            if wf.getnchannels() != 2:
                return False
            sampwidth = wf.getsampwidth()
            framerate = wf.getframerate()
            nframes = wf.getnframes()

            total = min(nframes, framerate * 3)
            raw = wf.readframes(total)

            diffs = 0
            non_silent = 0
            for i in range(total):
                offset = i * 2 * sampwidth
                if sampwidth == 2:
                    left = struct.unpack_from("<h", raw, offset)[0]
                    right = struct.unpack_from("<h", raw, offset + 2)[0]
                else:
                    left = struct.unpack_from("<b", raw, offset)[0]
                    right = struct.unpack_from("<b", raw, offset + 1)[0]

                if abs(left) < 50 and abs(right) < 50:
                    continue
                non_silent += 1
                if abs(left - right) > 100:
                    diffs += 1

            if non_silent == 0:
                return True
            return diffs / non_silent < (1 - threshold)
    except Exception:
        return False


def _transcribe_stereo(wav_path: str, meta: dict) -> dict:
    channels = _get_wav_channels(wav_path)

    if channels == 1:
        click.echo(f"  Mono file detected, single transcription.")
        left_segments, left_text = transcribe(
            wav_path,
            model_size=config.WHISPER_MODEL_SIZE,
            device=config.WHISPER_DEVICE,
            compute_type=config.WHISPER_COMPUTE_TYPE,
        )
        right_segments = []
        final_segments = []
        for seg in left_segments:
            final_segments.append({
                "start": seg["start"],
                "end": seg["end"],
                "speaker": "unknown",
                "text": seg["text"].strip(),
            })
        is_stereo = False
    elif _is_mono_duplicate(wav_path):
        click.echo(f"  Stereo with duplicate channels, treating as mono.")
        left_segments, left_text = transcribe(
            wav_path,
            model_size=config.WHISPER_MODEL_SIZE,
            device=config.WHISPER_DEVICE,
            compute_type=config.WHISPER_COMPUTE_TYPE,
        )
        right_segments = []
        final_segments = []
        for seg in left_segments:
            final_segments.append({
                "start": seg["start"],
                "end": seg["end"],
                "speaker": "unknown",
                "text": seg["text"].strip(),
            })
        is_stereo = False
    else:
        click.echo(f"  Real stereo, splitting channels...")
        left_wav, right_wav = split_stereo_channels(wav_path, config.TEMP_DIR)

        click.echo(f"  Transcribing left channel (callee)...")
        left_segments, left_text = transcribe(
            left_wav,
            model_size=config.WHISPER_MODEL_SIZE,
            device=config.WHISPER_DEVICE,
            compute_type=config.WHISPER_COMPUTE_TYPE,
        )

        click.echo(f"  Transcribing right channel (caller)...")
        right_segments, right_text = transcribe(
            right_wav,
            model_size=config.WHISPER_MODEL_SIZE,
            device=config.WHISPER_DEVICE,
            compute_type=config.WHISPER_COMPUTE_TYPE,
        )

        click.echo(f"  Aligning speakers & text...")
        final_segments = align(left_segments, right_segments)
        is_stereo = True

    caller_name = resolve_contact_name(meta.get("caller_number"))
    callee_name = resolve_contact_name(meta.get("callee_number"))
    employee_name = meta.get("employee")

    full_text = ""
    for seg in final_segments:
        spk = seg["speaker"]
        if spk == "caller":
            label = caller_name or "Звонящий"
        elif spk == "callee":
            label = callee_name or (employee_name or "Принимающий")
        else:
            label = "?"
        full_text += f"[{label}]: {seg['text']}\n"

    contact_info = resolve_contacts(
        meta.get("caller_number"),
        meta.get("callee_number"),
    )

    call_data = {
        **meta,
        **contact_info,
        "transcript": full_text,
        "segments": final_segments,
        "status": "processed",
        "raw_metadata": meta,
    }

    try:
        from storage.directus import DirectusClient
        dc = DirectusClient()
        if dc.is_available():
            collection = _choose_collection(meta)
            dc.create_call(call_data, collection=collection)
            click.echo(f"  Saved to Directus ({collection}).")
        dc.close()
    except Exception as e:
        click.echo(f"  Directus save error: {e}")

    to_remove = [wav_path]
    if is_stereo:
        to_remove.extend([left_wav, right_wav])
    for f in to_remove:
        if f and os.path.exists(f):
            os.remove(f)

    return call_data


def _process_call(audio_path: str, meta: dict) -> dict:
    click.echo(f"  Converting audio...")
    wav_path, channels = convert_to_wav(audio_path, config.TEMP_DIR)
    return _transcribe_stereo(wav_path, meta)


def _get_audio_file(meta, temp_dir: str) -> Optional[str]:
    if meta.audio_data:
        fname = meta.audio_filename or f"call_{meta.call_id}.mp3"
        local_path = os.path.join(temp_dir, fname)
        with open(local_path, "wb") as f:
            f.write(meta.audio_data)
        return local_path
    if meta.audio_url:
        return download_audio(meta.audio_url, temp_dir)
    return None


def _is_already_processed(call_id: str) -> bool:
    from storage.directus import DirectusClient
    dc = DirectusClient()
    try:
        for col in ["calls", "calls_74951284933"]:
            resp = dc._client.get(f"/items/{col}", params={
                "filter[call_id][_eq]": call_id,
                "limit": 1, "fields": "id",
            })
            if resp.status_code == 200 and resp.json().get("data"):
                return True
    except Exception:
        pass
    finally:
        dc.close()
    return False


@click.group()
def cli():
    pass


@cli.command()
@click.option("--limit", default=None, type=int, help="Max emails to process")
@click.option("--parallel", default=False, is_flag=True, help="Parallel download + convert")
def process(limit, parallel):
    cnt = _run_email_batch(limit=limit, parallel=parallel)
    click.echo(f"\nProcessed {cnt} calls.")


def _run_email_batch(limit=None, parallel=False):
    fetcher = EmailFetcher(
        config.IMAP_SERVER,
        config.IMAP_PORT,
        config.IMAP_USER,
        config.IMAP_PASSWORD,
        config.IMAP_FOLDER,
    )
    fetcher.connect()
    messages = fetcher.fetch_unseen()
    click.echo(f"Found {len(messages)} unseen messages.")

    parser = get_parser("novofon")

    tasks = []
    seen_call_ids = set()
    for msg in messages[:limit] if limit else messages:
        click.echo(f"\n--- Message: {msg.subject} ---")

        meta = parser.parse_email(msg.subject, msg.body_text, msg.body_html)

        # Fallback: check for audio attachment
        if not meta or not meta.call_id:
            for att in msg._attachments:
                att_meta = parser.parse_attachment(msg.subject, att.filename)
                if att_meta and att_meta.call_id:
                    meta = att_meta
                    meta.audio_data = att.data
                    meta.audio_filename = att.filename
                    click.echo(f"  Using audio attachment: {att.filename}")
                    break

            if not meta or not meta.call_id:
                click.echo("  No Novofon content found, skipping.")
                continue

        meta.source_email_id = msg.message_id
        if not meta.datetime and msg.date_parsed:
            meta.datetime = msg.date_parsed

        if meta.call_id:
            if meta.call_id in seen_call_ids:
                click.echo(f"  Duplicate call_id {meta.call_id} in batch, skipping.")
                continue
            if _is_already_processed(meta.call_id):
                click.echo(f"  Call {meta.call_id} already processed, skipping.")
                fetcher.mark_seen(msg.uid)
                continue
            seen_call_ids.add(meta.call_id)

        tasks.append((msg, meta))

    if not tasks:
        click.echo("No new calls to process.")
        fetcher.disconnect()
        return 0

    click.echo(f"\n=== Phase 1: Download + convert ({len(tasks)} files) ===")

    prepared = []

    if parallel:
        def prep(msg_meta):
            msg, meta = msg_meta
            try:
                audio_path = _get_audio_file(meta, config.TEMP_DIR)
                if not audio_path:
                    return (msg, meta, None, "download failed")
                wav_path, _ = convert_to_wav(audio_path, config.TEMP_DIR)
                if audio_path and os.path.exists(audio_path):
                    os.remove(audio_path)
                return (msg, meta, wav_path, None)
            except Exception as e:
                return (msg, meta, None, str(e))

        with ThreadPoolExecutor(max_workers=config.MAX_WORKERS) as pool:
            futures = [pool.submit(prep, t) for t in tasks]
            for f in as_completed(futures):
                msg, meta, wav_path, err = f.result()
                if err:
                    click.echo(f"  ! Prep failed: {msg.subject} - {err}")
                else:
                    prepared.append((msg, meta, wav_path))
                    click.echo(f"  Prepared: {meta.call_id}")
    else:
        for msg, meta in tasks:
            try:
                click.echo(f"  Preparing {meta.call_id}...")
                audio_path = _get_audio_file(meta, config.TEMP_DIR)
                if not audio_path:
                    click.echo("    Prep failed, skipping.")
                    continue
                wav_path, _ = convert_to_wav(audio_path, config.TEMP_DIR)
                if audio_path and os.path.exists(audio_path):
                    os.remove(audio_path)
                prepared.append((msg, meta, wav_path))
            except Exception as e:
                click.echo(f"    Error: {e}")

    if not prepared:
        click.echo("No files prepared for transcription.")
        fetcher.disconnect()
        return 0

    click.echo(f"\n=== Phase 2: Transcribe + save ({len(prepared)} files, sequential on GPU) ===")

    processed = 0
    for msg, meta, wav_path in prepared:
        try:
            call_data = _transcribe_stereo(wav_path, {
                "call_id": meta.call_id,
                "datetime": meta.datetime,
                "direction": meta.direction,
                "caller_number": meta.caller_number,
                "callee_number": meta.callee_number,
                "employee": meta.employee,
                "duration_sec": meta.duration_sec,
                "audio_url": meta.audio_url,
                "source_email_id": meta.source_email_id,
            })
            fetcher.mark_seen(msg.uid)
            processed += 1
            click.echo(f"  Done! Call ID: {call_data.get('call_id')}")
        except Exception as e:
            click.echo(f"  Error processing {meta.call_id}: {e}")
            if wav_path and os.path.exists(wav_path):
                os.remove(wav_path)

    fetcher.disconnect()
    return processed


@cli.command()
@click.option("--file", required=True, help="Path to audio file")
@click.option("--caller", default=None, help="Caller number")
@click.option("--callee", default=None, help="Callee number")
@click.option("--duration", default=None, type=int, help="Duration in seconds")
def process_file(file, caller, callee, duration):
    click.echo(f"Processing file: {file}")
    meta = {
        "call_id": f"manual_{os.path.splitext(os.path.basename(file))[0]}",
        "caller_number": caller,
        "callee_number": callee,
        "direction": "inbound" if caller else None,
        "duration_sec": duration,
        "audio_url": file,
    }

    call_data = _process_call(file, meta)
    click.echo(f"Done! Call ID: {call_data.get('call_id')}")


@cli.command()
@click.option("--output", default="export.json", help="Output file path")
def export(output):
    export_to_json(output)
    click.echo(f"Exported to {output}")


@cli.command()
@click.option("--max", "max_per_cycle", default=None, type=int, help="Max calls per cycle")
def daemon(max_per_cycle):
    import time
    import schedule

    max_calls = max_per_cycle or config.MAX_PER_CYCLE

    def job():
        click.echo("Checking for new emails...")
        try:
            cnt = _run_email_batch(limit=max_calls)
            click.echo(f"Processed {cnt} calls this cycle.")
        except Exception as e:
            click.echo(f"Error: {e}")

    schedule.every(5).minutes.do(job)

    click.echo(f"Daemon started. Checking every 5 minutes, max {max_calls} calls/cycle.")
    while True:
        schedule.run_pending()
        time.sleep(30)


def _get_wav_channels(wav_path: str) -> int:
    import wave
    try:
        with wave.open(wav_path, "rb") as wf:
            return wf.getnchannels()
    except Exception:
        return 2


def _choose_collection(meta: dict) -> str:
    caller = meta.get("caller_number") or ""
    callee = meta.get("callee_number") or ""
    for num in [caller, callee]:
        cleaned = num.replace("+", "").replace("-", "").replace(" ", "")
        if "74951284933" in cleaned:
            return "calls_74951284933"
    return "calls"


if __name__ == "__main__":
    cli()