main.py

import json
import os
import subprocess
import uuid
import shutil
import threading
import tempfile
from datetime import datetime, timezone
from pathlib import Path

import librosa
import soundfile as sf
from dotenv import load_dotenv
from fastapi import FastAPI, File, Form, UploadFile, HTTPException
from fastapi.responses import FileResponse, JSONResponse
from fastapi.staticfiles import StaticFiles
from huggingface_hub import HfApi, create_repo, hf_hub_download, upload_file


# ── Env ────────────────────────────────────────────────────────────────────
load_dotenv()

HF_TOKEN       = os.getenv("HF_TOKEN")
DATASET_REPO   = os.getenv("HF_DATASET_REPO")
MODEL_REPO     = os.getenv("HF_MODEL_REPO")                       # e.g. "your-username/tiny-whisper-cpp-q8"
MODEL_FILENAME = os.getenv("HF_MODEL_FILENAME", "ggml-tiny-q8_0.bin")

# ── Paths ──────────────────────────────────────────────────────────────────
BASE_DIR    = Path(__file__).parent
STATIC_DIR  = BASE_DIR / "static"
DATASET_DIR = BASE_DIR / "dataset"
AUDIO_DIR   = DATASET_DIR / "audio"
MODEL_DIR   = BASE_DIR / "models"
MANIFEST    = DATASET_DIR / "transcripts.jsonl"

STATIC_DIR.mkdir(exist_ok=True)
DATASET_DIR.mkdir(exist_ok=True)
AUDIO_DIR.mkdir(exist_ok=True)
MODEL_DIR.mkdir(exist_ok=True)

WHISPER_BIN = os.getenv("WHISPER_CLI_BIN", "/usr/local/bin/whisper-cli")

# ── HuggingFace setup ──────────────────────────────────────────────────────
_hf_api: HfApi | None = None

def get_hf_api() -> HfApi | None:
    global _hf_api
    if _hf_api is None and HF_TOKEN:
        _hf_api = HfApi(token=HF_TOKEN)
        try:
            create_repo(
                repo_id=DATASET_REPO,
                repo_type="dataset",
                exist_ok=True,
                token=HF_TOKEN,
            )
        except Exception as e:
            print(f"[HF] Dataset repo check: {e}")
    return _hf_api

# ── Model (ggml weights pulled from the Hub, lazy on first request) ───────
_MODEL_PATH: str | None = None

def get_model_path() -> str:
    global _MODEL_PATH
    if _MODEL_PATH is None:
        local_path = MODEL_DIR / MODEL_FILENAME
        if local_path.exists():
            _MODEL_PATH = str(local_path)
        else:
            print(f"Downloading {MODEL_FILENAME} from {MODEL_REPO} …")
            _MODEL_PATH = hf_hub_download(
                repo_id=MODEL_REPO,
                filename=MODEL_FILENAME,
                token=HF_TOKEN,
                local_dir=str(MODEL_DIR),
            )
            print(f"Model ready at {_MODEL_PATH}.")
    return _MODEL_PATH

# ── App ────────────────────────────────────────────────────────────────────
app = FastAPI(title="Kasanoma ASR", version="3.0.0")

_csv_lock = threading.Lock()

# ── Helpers ────────────────────────────────────────────────────────────────

def audio_duration(path: Path) -> float:
    try:
        y, sr = librosa.load(str(path), sr=None, mono=True)
        return round(len(y) / sr, 2)
    except Exception:
        return 0.0


def save_entry(entry: dict) -> None:
    with MANIFEST.open("a", encoding="utf-8") as f:
        f.write(json.dumps(entry, ensure_ascii=False) + "\n")


def load_manifest() -> list[dict]:
    if not MANIFEST.exists():
        return []
    entries = []
    with MANIFEST.open(encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if line:
                entries.append(json.loads(line))
    return entries


def _upload_in_background(audio_path: Path, relative_audio_path: str) -> None:
    api = get_hf_api()
    if api is None:
        return
    try:
        upload_file(
            path_or_fileobj=str(audio_path),
            path_in_repo=relative_audio_path,
            repo_id=DATASET_REPO,
            repo_type="dataset",
            token=HF_TOKEN,
        )
        with _csv_lock:
            upload_file(
                path_or_fileobj=str(MANIFEST),
                path_in_repo="transcripts.jsonl",
                repo_id=DATASET_REPO,
                repo_type="dataset",
                token=HF_TOKEN,
            )
    except Exception as e:
        print(f"[Background upload error] {e}")


def transcribe_path(audio_path: Path) -> tuple[str, str]:
    """Run inference via the whisper.cpp CLI against the int8 ggml model."""
    model_path = get_model_path()

    # librosa decodes any format (webm, mp4, ogg, wav) via ffmpeg;
    # whisper-cli wants a 16kHz mono wav on disk.
    audio_data, _ = librosa.load(str(audio_path), sr=16000, mono=True)

    with tempfile.TemporaryDirectory() as tmp_dir:
        wav_path    = Path(tmp_dir) / "input.wav"
        out_prefix  = Path(tmp_dir) / "out"
        sf.write(wav_path, audio_data, 16000)

        result = subprocess.run(
            [
                WHISPER_BIN,
                "-m", model_path,
                "-f", str(wav_path),
                "-l", "auto",
                "-oj",
                "-of", str(out_prefix),
                "-nt",              # no timestamps in stdout
            ],
            capture_output=True,
            text=True,
            timeout=60,
        )
        if result.returncode != 0:
            raise RuntimeError(result.stderr.strip() or "whisper-cli failed")

        json_path = out_prefix.with_suffix(".json")
        if not json_path.exists():
            raise RuntimeError("whisper-cli did not produce an output file")

        with json_path.open(encoding="utf-8") as f:
            out = json.load(f)

    segments = out.get("transcription", [])
    transcription = " ".join(seg.get("text", "") for seg in segments).strip()

    return transcription, "tw"


# ── Routes ─────────────────────────────────────────────────────────────────

@app.get("/", include_in_schema=False)
async def root():
    index = STATIC_DIR / "index.html"
    if not index.exists():
        raise HTTPException(404, "Frontend not found. Place index.html in static/")
    return FileResponse(index)


@app.post("/transcribe")
async def transcribe(audio: UploadFile = File(...)):
    suffix   = Path(audio.filename or "audio.webm").suffix or ".webm"
    tmp_path = BASE_DIR / f"_tmp_{uuid.uuid4().hex}{suffix}"

    try:
        with tmp_path.open("wb") as f:
            shutil.copyfileobj(audio.file, f)

        transcription, detected_lang = transcribe_path(tmp_path)

        return JSONResponse({
            "transcription":     transcription,
            "detected_language": detected_lang,
        })

    except Exception as e:
        print(f"[Transcribe error] {e}")
        raise HTTPException(500, detail=str(e))

    finally:
        tmp_path.unlink(missing_ok=True)


@app.post("/save")
async def save(
    audio: UploadFile = File(...),
    transcription: str = Form(...),
):
    if not transcription.strip():
        raise HTTPException(422, "Transcription must not be empty.")

    entry_id       = uuid.uuid4().hex
    suffix         = Path(audio.filename or "audio.webm").suffix or ".webm"
    audio_filename = f"{entry_id}{suffix}"
    audio_path     = AUDIO_DIR / audio_filename

    with audio_path.open("wb") as f:
        shutil.copyfileobj(audio.file, f)

    duration = audio_duration(audio_path)

    entry = {
        "id":            entry_id,
        "audio_file":    f"dataset/audio/{audio_filename}",
        "transcription": transcription.strip(),
        "language":      "twi_en",
        "duration_s":    duration,
        "created_at":    datetime.now(timezone.utc).isoformat(),
    }

    with _csv_lock:
        save_entry(entry)

    relative_audio_path = f"audio/{audio_filename}"
    threading.Thread(
        target=_upload_in_background,
        args=(audio_path, relative_audio_path),
        daemon=True,
    ).start()

    total = len(load_manifest())
    return JSONResponse({
        "id":          entry_id,
        "total_saved": total,
        "duration_s":  duration,
    })


@app.get("/dataset/stats")
async def dataset_stats():
    entries = load_manifest()
    if not entries:
        return JSONResponse({"total": 0, "total_duration_s": 0, "total_words": 0})

    total_duration = sum(e.get("duration_s", 0) for e in entries)
    total_words    = sum(len(e["transcription"].split()) for e in entries)

    return JSONResponse({
        "total":            len(entries),
        "total_duration_s": round(total_duration, 1),
        "total_words":      total_words,
    })


@app.get("/dataset/entries")
async def dataset_entries(limit: int = 50, offset: int = 0):
    entries = load_manifest()
    entries.reverse()
    return JSONResponse({
        "entries": entries[offset : offset + limit],
        "total":   len(entries),
    })


# ── Static files ───────────────────────────────────────────────────────────
app.mount("/static", StaticFiles(directory=str(STATIC_DIR)), name="static")
app.mount("/assets", StaticFiles(directory="assets"), name="assets")