data/main.py

#!/usr/bin/env python3
"""
Audio Transcription Pipeline CLI.
Process audio files through transcription, audience classification,
diarization, summarization, and ASCII spectrogram visualization.

Usage:
    python main.py transcribe <audio_path> [--output json]
    python main.py summarize <audio_path> [--output json]
    python main.py audience <audio_path> [--output json]
    python main.py ascii-viz <audio_path> [--output file]
    python main.py stream <device> [--output json]
    python main.py all <audio_path> [--output json] [--ascii]
"""

import argparse
import json
import logging
import os
import sys
from typing import Dict, Any

logging.basicConfig(
    level=logging.INFO,
    format="%(levelname)s:%(name)s:%(message)s",
)

logger = logging.getLogger(__name__)

# Disable verbose logging from libraries
logging.getLogger("faster_whisper").setLevel(logging.WARNING)
logging.getLogger("transformers").setLevel(logging.WARNING)
logging.getLogger("librosa").setLevel(logging.WARNING)
logging.getLogger("pipeline").setLevel(logging.INFO)


def _import_pipeline():
    """Import pipeline modules (lazy to avoid slow startup for help)."""
    from pipeline.transcriber import Transcriber
    from pipeline.audience_classifier import AudienceResponseClassifier
    from pipeline.diarizer import Diarizer
    from pipeline.summarizer import MeetingSummarizer
    from pipeline.ascii_spectrogram import AsciiSpectrogram
    from pipeline.orchestrator import AudioPipeline

    return {
        "Transcriber": Transcriber,
        "AudienceResponseClassifier": AudienceResponseClassifier,
        "Diarizer": Diarizer,
        "MeetingSummarizer": MeetingSummarizer,
        "AsciiSpectrogram": AsciiSpectrogram,
        "AudioPipeline": AudioPipeline,
    }


def _output_result(result: Dict[str, Any], output_format: str, output_path: str = None):
    """Output result in specified format."""
    # Convert non-serializable items
    clean = _clean_for_json(result)

    if output_format == "json":
        output = json.dumps(clean, indent=2, ensure_ascii=False, default=str)
        if output_path:
            with open(output_path, "w") as f:
                f.write(output)
            logger.info(f"Output written to {output_path}")
        else:
            print(output)
    else:
        # Text format
        _print_text_result(clean, output_path)


def _clean_for_json(obj):
    """Recursively clean objects for JSON serialization."""
    if isinstance(obj, dict):
        return {k: _clean_for_json(v) for k, v in obj.items() if not k.startswith("_")}
    elif isinstance(obj, list):
        return [_clean_for_json(item) for item in obj]
    elif isinstance(obj, float):
        if obj != obj:  # NaN check
            return None
        return obj
    return obj


def _print_text_result(result: Dict[str, Any], output_path: str = None):
    """Print result in human-readable text format."""
    lines = []

    if "error" in result:
        lines.append(f"ERROR: {result['error']}")
    else:
        # Metadata
        meta = result.get("metadata", {})
        lines.append("=" * 60)
        lines.append("AUDIO TRANSCRIPTION PIPELINE RESULT")
        lines.append("=" * 60)
        lines.append(
            f"Duration: {meta.get('duration', 'N/A'):.1f}s | "
            f"Segments: {meta.get('num_segments', 0)} | "
            f"Processing: {meta.get('processing_time_seconds', 0):.1f}s"
        )
        lines.append("")

        # Segments
        segments = result.get("segments", [])
        if segments:
            lines.append("--- TRANSCRIPT ---")
            for seg in segments:
                speaker = seg.get("speaker", "?")
                text = seg.get("text", "").strip()
                start = seg.get("start", 0)
                end = seg.get("end", 0)
                audience = seg.get("audience_response", "")
                conf = seg.get("confidence", 0)
                tag = f" [{audience}]" if audience and audience != "unknown" else ""
                lines.append(f"  [{start:6.1f}s-{end:6.1f}s] {speaker}: {text}{tag}")
            lines.append("")

        # Audience responses
        ar = result.get("audience_responses", [])
        if ar:
            lines.append("--- AUDIENCE RESPONSES ---")
            for resp in ar[:10]:  # Show first 10
                lines.append(
                    f"  [{resp.get('start', 0):.1f}s-{resp.get('end', 0):.1f}s] "
                    f"{resp.get('response_class', '?')} "
                    f"(conf: {resp.get('confidence', 0):.2f})"
                )
            if len(ar) > 10:
                lines.append(f"  ... and {len(ar) - 10} more")
            lines.append("")

        # Summary
        summary = result.get("summary", {})
        if summary and summary.get("overview"):
            lines.append("--- SUMMARY ---")
            lines.append(f"  Overview: {summary.get('overview', 'N/A')}")
            decisions = summary.get("decisions", [])
            if decisions:
                lines.append("  Decisions:")
                for d in decisions:
                    lines.append(f"    - {d}")
            actions = summary.get("action_items", [])
            if actions:
                lines.append("  Action Items:")
                for a in actions:
                    lines.append(f"    - {a}")
            topics = summary.get("topics", [])
            if topics:
                lines.append("  Topics:")
                for t in topics:
                    lines.append(f"    - {t}")
            lines.append("")

        # ASCII frames (show summary count only)
        frames = result.get("ascii_frames", [])
        if frames:
            lines.append(f"--- ASCII SPECTROGRAM ---")
            lines.append(f"  {len(frames)} frames generated")
            lines.append(f"  First frame preview:")
            first = frames[0]
            if isinstance(first, dict):
                lines.append(f"    t={first.get('timestamp', 0):.1f}s")
                frame_text = first.get("frame", "")
                for line in frame_text.split("\n")[:5]:
                    lines.append(f"    |{line}")
            else:
                lines.append(f"    {str(first)[:60]}...")
            lines.append("")

    output = "\n".join(lines)
    if output_path:
        with open(output_path, "w") as f:
            f.write(output)
        logger.info(f"Output written to {output_path}")
    else:
        print(output)


def cmd_transcribe(args):
    """Transcribe an audio file."""
    # Hybrid mode: use Encoder-Projector-LLM pipeline
    if getattr(args, "hybrid", False):
        from hybrid_model.infer import run_pipeline

        # Determine paths
        models_dir = "models"
        hybrid_dir = "hybrid_model"
        qwen_gguf = os.path.join(models_dir, "Qwen3-8B-Q4_K_M.gguf")
        qwen_fallback = os.path.join(models_dir, "qwen2.5-0.5b-instruct-q4_k_m.gguf")
        projector_ckpt = os.path.join(hybrid_dir, "projector_checkpoint_best.pt")

        if os.path.exists(qwen_gguf):
            llm_path = qwen_gguf
        elif os.path.exists(qwen_fallback):
            llm_path = qwen_fallback
            logger.warning("Qwen3-8B not found, using Qwen2.5-0.5B fallback")
        else:
            logger.error("No LLM GGUF found for hybrid mode!")
            sys.exit(1)

        logger.info(f"Hybrid mode: LLM={llm_path}, Projector={projector_ckpt}")
        result_text = run_pipeline(
            audio_path=args.audio_path,
            llm_path=llm_path,
            projector_checkpoint=projector_ckpt,
            max_new_tokens=200,
            temperature=0.1,
        refine_with_llm_flag=True,
        )

        # Estimate duration from audio for metadata
        import librosa
        try:
            audio_dur = librosa.get_duration(path=args.audio_path)
        except Exception:
            audio_dur = 0.0

        # Build result dict
        result = {
            "metadata": {
                "hybrid_mode": True,
                "llm": os.path.basename(llm_path),
                "duration": audio_dur,
                "num_segments": 1,
                "processing_time_seconds": 0.0,
            },
            "segments": [{"start": 0, "end": audio_dur, "speaker": "SPEAKER_00", "text": result_text}],
        }
        _output_result(result, args.output, getattr(args, "output_file", None))
        return

    # Standard pipeline mode
    mods = _import_pipeline()
    pipeline = mods["AudioPipeline"](enable_summarizer=False, enable_ascii=False)
    result = pipeline.process_file(
        args.audio_path,
        language=args.language,
        vad_filter=not args.no_vad,
    )
    _output_result(result, args.output, getattr(args, "output_file", None))


def cmd_summarize(args):
    """Transcribe and summarize a meeting audio."""
    mods = _import_pipeline()
    pipeline = mods["AudioPipeline"](
        enable_summarizer=True, enable_ascii=False
    )
    result = pipeline.process_file(
        args.audio_path,
        language=args.language,
        vad_filter=not args.no_vad,
    )
    _output_result(result, args.output, getattr(args, "output_file", None))


def cmd_audience(args):
    """Classify audience responses in an audio file."""
    mods = _import_pipeline()
    classifier = mods["AudienceResponseClassifier"]()
    if args.output == "json":
        result = classifier.classify_file(args.audio_path)
        print(json.dumps(result, indent=2))
    else:
        result = classifier.classify_file(args.audio_path)
        print(f"Audience Responses ({len(result)} detections):")
        for resp in result:
            print(
                f"  [{resp['start']:.1f}s-{resp['end']:.1f}s] "
                f"{resp['response_class']} ({resp['confidence']:.2f})"
            )


def cmd_ascii_viz(args):
    """Generate ASCII spectrogram visualization."""
    import librosa

    mods = _import_pipeline()
    audio, sr = librosa.load(args.audio_path, sr=16000, mono=True)
    duration = len(audio) / sr

    viz = mods["AsciiSpectrogram"](
        columns=args.columns,
        rows=args.rows,
        fps=args.fps,
        mode=args.mode,
    )

    frames = list(viz.generate_frames(audio, sr))
    print(f"Generated {len(frames)} frames from {duration:.1f}s audio")

    if args.output == "file" or args.output_file:
        path = args.output_file or f"ascii_output_{os.path.basename(args.audio_path)}.txt"
        with open(path, "w") as f:
            for ascii_text, timestamp in frames:
                f.write(f"--- Frame @ {timestamp:.1f}s ---\n{ascii_text}\n\n")
        print(f"ASCII frames written to: {path}")
    else:
        # Print first few frames
        for i, (ascii_text, timestamp) in enumerate(frames[:5]):
            print(f"\n=== Frame {i + 1} @ {timestamp:.1f}s ===")
            print(ascii_text)
        if len(frames) > 5:
            print(f"\n... and {len(frames) - 5} more frames")


def cmd_stream(args):
    """Process stream from microphone or file."""
    # For now, just read a file in streaming mode
    import librosa
    import soundfile as sf
    import tempfile

    mods = _import_pipeline()
    pipeline = mods["AudioPipeline"](enable_summarizer=False, enable_ascii=False)

    # Load audio and chunk it
    audio, sr = librosa.load(args.audio_path, sr=16000, mono=True)
    chunk_size = sr * 2  # 2-second chunks

    def chunk_gen():
        for i in range(0, len(audio), chunk_size):
            yield audio[i : i + chunk_size]

    print(f"Stream processing: {args.audio_path}")
    for i, result in enumerate(pipeline.process_stream(chunk_gen(), sr)):
        segments = result.get("segments", [])
        print(f"Chunk {i + 1}: {len(segments)} segments")
        for seg in segments[:2]:
            print(f"  [{seg.get('start', 0):.1f}s] {seg.get('speaker', '?')}: {seg.get('text', '')[:60]}")


def cmd_all(args):
    """Run full pipeline (transcribe + audience + diarize + summarize + optional ASCII)."""
    mods = _import_pipeline()
    pipeline = mods["AudioPipeline"](
        enable_summarizer=True,
        enable_ascii=args.ascii,
    )
    result = pipeline.process_file(
        args.audio_path,
        language=args.language,
        vad_filter=not args.no_vad,
    )
    _output_result(result, args.output, getattr(args, "output_file", None))


def main():
    parser = argparse.ArgumentParser(
        description="Audio Transcription Pipeline — transcribe, classify, "
        "diarize, summarize, and visualize audio",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  python main.py transcribe meeting.wav
  python main.py transcribe meeting.wav --output json --no-vad
  python main.py summarize meeting.wav --output json
  python main.py all meeting.wav --ascii --output json
  python main.py audience presentation.wav
  python main.py ascii-viz music.mp3 --mode combined --columns 100
        """,
    )
    parser.add_argument(
        "--verbose", "-v", action="store_true", help="Enable verbose logging"
    )

    subparsers = parser.add_subparsers(dest="command", help="Command to execute")

    # transcribe
    p_transcribe = subparsers.add_parser(
        "transcribe", help="Transcribe audio to text"
    )
    p_transcribe.add_argument("audio_path", help="Path to audio file")
    p_transcribe.add_argument(
        "--output", choices=["text", "json"], default="text", help="Output format"
    )
    p_transcribe.add_argument(
        "--output-file", "-o", help="Write output to file instead of stdout"
    )
    p_transcribe.add_argument(
        "--language", "-l", help="Language code (e.g., 'en'). Default: auto-detect"
    )
    p_transcribe.add_argument(
        "--no-vad", action="store_true", help="Disable VAD filtering"
    )
    p_transcribe.add_argument(
    "--hybrid", action="store_true",
    help="Use hybrid Encoder-Projector-LLM model (Whisper+Qwen3-8B)"
    )
    p_transcribe.set_defaults(func=cmd_transcribe)

    # summarize
    p_summarize = subparsers.add_parser(
        "summarize", help="Transcribe and summarize meeting audio"
    )
    p_summarize.add_argument("audio_path", help="Path to audio file")
    p_summarize.add_argument(
        "--output", choices=["text", "json"], default="text", help="Output format"
    )
    p_summarize.add_argument(
        "--output-file", "-o", help="Write output to file instead of stdout"
    )
    p_summarize.add_argument(
        "--language", "-l", help="Language code (default: auto-detect)"
    )
    p_summarize.add_argument(
        "--no-vad", action="store_true", help="Disable VAD filtering"
    )
    p_summarize.set_defaults(func=cmd_summarize)

    # audience
    p_audience = subparsers.add_parser(
        "audience", help="Classify audience responses in audio"
    )
    p_audience.add_argument("audio_path", help="Path to audio file")
    p_audience.add_argument(
        "--output", choices=["text", "json"], default="text", help="Output format"
    )
    p_audience.set_defaults(func=cmd_audience)

    # ascii-viz
    p_ascii = subparsers.add_parser(
        "ascii-viz", help="Generate ASCII spectrogram visualization"
    )
    p_ascii.add_argument("audio_path", help="Path to audio file")
    p_ascii.add_argument(
        "--mode",
        choices=["spectrogram", "waveform", "combined"],
        default="spectrogram",
        help="Visualization mode",
    )
    p_ascii.add_argument("--columns", type=int, default=80, help="ASCII width")
    p_ascii.add_argument("--rows", type=int, default=20, help="ASCII height")
    p_ascii.add_argument("--fps", type=int, default=10, help="Frames per second")
    p_ascii.add_argument(
        "--output", choices=["text", "file"], default="text", help="Output format"
    )
    p_ascii.add_argument(
        "--output-file", "-o", help="Write output to file"
    )
    p_ascii.set_defaults(func=cmd_ascii_viz)

    # stream
    p_stream = subparsers.add_parser(
        "stream", help="Stream process audio (chunked from file)"
    )
    p_stream.add_argument("audio_path", help="Path to audio file to stream")
    p_stream.add_argument(
        "--output", choices=["text", "json"], default="text", help="Output format"
    )
    p_stream.set_defaults(func=cmd_stream)

    # all
    p_all = subparsers.add_parser(
        "all", help="Run full pipeline (transcribe + audience + diarize + summarize)"
    )
    p_all.add_argument("audio_path", help="Path to audio file")
    p_all.add_argument(
        "--output", choices=["text", "json"], default="text", help="Output format"
    )
    p_all.add_argument(
        "--output-file", "-o", help="Write output to file instead of stdout"
    )
    p_all.add_argument(
        "--language", "-l", help="Language code (default: auto-detect)"
    )
    p_all.add_argument(
        "--no-vad", action="store_true", help="Disable VAD filtering"
    )
    p_all.add_argument(
        "--ascii", action="store_true", help="Include ASCII spectrogram frames"
    )
    p_all.set_defaults(func=cmd_all)

    args = parser.parse_args()

    if args.verbose:
        logging.getLogger().setLevel(logging.DEBUG)

    if not hasattr(args, "func"):
        parser.print_help()
        sys.exit(1)

    args.func(args)


if __name__ == "__main__":
    main()