"""Transcribe an episode with a Whisper instance (docs/transcription.md, "Transcribing").

The server downloads the episode (as silence.py does, with the same size limit), decodes it
once to 16 kHz mono, cuts it into pieces of about eight minutes where there is a pause (so no
word is cut), sends the pieces in parallel to the instance's OpenAI-compatible
`/audio/transcriptions` and joins the answers into one transcript, which is kept on the job
(TranscriptionJob.data).

A job runs in a daemon thread, one at a time per server; the others wait as `queued` and are
taken oldest first (by id). A job found `downloading` or `transcribing` when the process has
just started (a restart killed it) is queued again. Cancelling deletes the job's row: the
worker notices when it next writes progress and stops.
"""
import datetime
import logging
import os
import subprocess
import tempfile
import threading
import time
from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait

import requests
from django.conf import settings
from django.db import close_old_connections
from django.utils import timezone

from . import chapters, silence, transcripts
from .models import TranscriptionJob, TranscriptionSettings
from .transcription import DEFAULT_MODEL, for_user, token_of

log = logging.getLogger(__name__)

CHUNK_SECONDS = 8 * 60      # about this long a piece
CUT_WINDOW = 30             # a cut moves to the longest pause within this many seconds
MIN_TAIL = 60               # a last piece shorter than this is joined to the one before
RATE = 16000                # Hz, mono, 16-bit
WINDOW = silence.WINDOW
RETRIES = 3
RETRY_PAUSES = (3, 10, 30)  # seconds before each retry
UPLOAD_TIMEOUT = (10, 3600)  # a CPU instance may need many minutes for a piece
DECODE_TIMEOUT = 3600
OPUS_BITRATE = '32k'
MAX_PROMPT = 500            # characters
ALIVE_CHECK = 5             # seconds between looks whether the job was cancelled
WAIT_RETRY = 600            # seconds between tries while the built-in instance can't be reached
WAIT_MAX = 48 * 3600        # after this long since the button was pressed a waiting job gives up
WAIT_IDLE = 30              # seconds the runner sleeps while every queued job is waiting

_lock = threading.Lock()
_runner = None              # the thread working through the queue, or None
_resumed = False            # whether jobs left over from before a restart were queued again
_not_before = {}            # job id -> time.time() before which a waiting job isn't tried again
WAIT = 'wait'               # _next_job: jobs are queued, but none is due yet


class Cancelled(Exception):
    """The job's row is gone: the user cancelled it."""


class Failed(Exception):
    """The job can't succeed; the message is shown to the user."""


class Unreachable(Exception):
    """The built-in instance can't be reached (e.g. the machine it runs on is off): the job waits
    and tries again later instead of failing."""


class Permanent(Exception):
    """A request the instance refused for good (a client error): not worth a retry."""


# ---------------------------------------------------------------------------
# Queue
# ---------------------------------------------------------------------------

def kick():
    """Make sure queued jobs are being worked on. Called by the endpoints, so after a restart
    the first request brings back what was running."""
    global _runner, _resumed
    with _lock:
        if not _resumed:
            _resumed = True
            TranscriptionJob.objects.filter(
                state__in=[TranscriptionJob.DOWNLOADING, TranscriptionJob.TRANSCRIBING],
            ).update(state=TranscriptionJob.QUEUED, chunks_done=0, chunks_total=0)
        if _runner is not None:
            return
        if not TranscriptionJob.objects.filter(state=TranscriptionJob.QUEUED).exists():
            return
        _runner = threading.Thread(target=_loop, name='transcribe', daemon=True)
        _runner.start()


def _next_job():
    global _runner
    with _lock:
        queued = list(TranscriptionJob.objects.filter(state=TranscriptionJob.QUEUED)
                      .order_by('id').values_list('id', flat=True))
        if not queued:
            _runner = None   # decided under the lock, so a job queued right now starts a new runner
            return None
        now = time.time()
        due = next((i for i in queued if _not_before.get(i, 0) <= now), None)
        return due if due is not None else WAIT


def _loop():
    global _runner
    try:
        while True:
            close_old_connections()
            job_id = _next_job()
            if job_id is None:
                return
            if job_id == WAIT:
                time.sleep(WAIT_IDLE)
                continue
            try:
                run(job_id)
            except Exception:  # run() reports its own failures; this keeps the queue going
                log.exception('transcription job %s crashed', job_id)
    except BaseException:
        with _lock:
            _runner = None
        raise
    finally:
        close_old_connections()


def run(job_id):
    """Work one job through to `done` or `error`, in this thread."""
    started = TranscriptionJob.objects.filter(pk=job_id, state=TranscriptionJob.QUEUED).update(
        state=TranscriptionJob.DOWNLOADING, started_at=timezone.now(), finished_at=None,
        error='', chunks_done=0, chunks_total=0)
    if not started:
        return   # cancelled while it waited
    try:
        data = _transcribe(job_id)
        _write(job_id, state=TranscriptionJob.DONE, data=data, error='', finished_at=timezone.now())
    except Cancelled:
        log.info('transcription job %s cancelled', job_id)
    except Unreachable as e:
        _wait(job_id, str(e))
    except Failed as e:
        _fail(job_id, str(e))
    except Exception as e:
        log.exception('transcription job %s failed', job_id)
        _fail(job_id, f'{type(e).__name__}: {e}')


def _wait(job_id, reason):
    """Back into the queue to try again in WAIT_RETRY seconds, or give up when it waited too long."""
    job = TranscriptionJob.objects.filter(pk=job_id).only('created_at').first()
    if job is None:
        return
    if (timezone.now() - job.created_at).total_seconds() > WAIT_MAX:
        _not_before.pop(job_id, None)
        _fail(job_id, 'Whisper-Instanz zwei Tage lang nicht erreichbar')
        return
    _not_before[job_id] = time.time() + WAIT_RETRY
    again = timezone.localtime(timezone.now() + datetime.timedelta(seconds=WAIT_RETRY)).strftime('%H:%M')
    log.info('transcription job %s waits for the instance (%s)', job_id, reason)
    TranscriptionJob.objects.filter(pk=job_id).update(
        state=TranscriptionJob.QUEUED, chunks_done=0, chunks_total=0, started_at=None,
        error=f'Whisper-Instanz nicht erreichbar, neuer Versuch um {again}')


def _probe(url):
    """Unreachable unless something answers at the built-in instance's address."""
    try:
        requests.get(url.rsplit('/audio/transcriptions', 1)[0], timeout=(5, 10)).close()
    except (requests.ConnectionError, requests.Timeout) as e:
        raise Unreachable(e.__class__.__name__)


def _fail(job_id, message):
    TranscriptionJob.objects.filter(pk=job_id).update(
        state=TranscriptionJob.ERROR, error=(message or 'failed')[:500], finished_at=timezone.now())


def _write(job_id, **fields):
    """Update the job; Cancelled when its row is gone."""
    if not TranscriptionJob.objects.filter(pk=job_id).update(**fields):
        raise Cancelled()


def _alive(job_id):
    if not TranscriptionJob.objects.filter(pk=job_id).exists():
        raise Cancelled()


# ---------------------------------------------------------------------------
# The work
# ---------------------------------------------------------------------------

def _target(conf):
    """(url, headers, builtin) of the instance the settings point at."""
    builtin = conf.engine == TranscriptionSettings.BUILTIN
    base = settings.WHISPER_URL if builtin else conf.url
    if not base:
        raise Failed('Keine eingebaute Instanz eingerichtet' if builtin else 'Keine Adresse eingetragen')
    headers = {'User-Agent': 'diora transcription'}
    if not builtin and conf.token_encrypted:
        headers['Authorization'] = 'Bearer ' + token_of(conf)
    return base.rstrip('/') + '/audio/transcriptions', headers, builtin


def _prompt(episode):
    parts = [episode.feed.title, episode.title, episode.feed.author]
    return '. '.join(p.strip() for p in parts if p and p.strip())[:MAX_PROMPT]


def _transcribe(job_id):
    job = TranscriptionJob.objects.select_related('episode__feed').get(pk=job_id)
    episode = job.episode
    conf = for_user(episode.feed.user)
    url, headers, builtin = _target(conf)
    if builtin:
        _probe(url)   # before downloading: is the instance there at all?
    options = {'model': conf.model or DEFAULT_MODEL, 'response_format': 'verbose_json'}
    if conf.language:
        options['language'] = conf.language
    prompt = _prompt(episode)
    if prompt:
        options['prompt'] = prompt

    with tempfile.TemporaryDirectory(prefix='diora-transcribe-') as tmp:
        source = os.path.join(tmp, 'episode.audio')
        _download(episode.audio_url, source, job_id)
        pcm = os.path.join(tmp, 'episode.pcm')
        _decode(source, pcm)
        os.remove(source)
        total = os.path.getsize(pcm) // 2 / RATE
        if total <= 0:
            raise Failed('Kein Ton in der Datei')
        pieces = plan(total, pauses_of(pcm))
        _write(job_id, state=TranscriptionJob.TRANSCRIBING, chunks_total=len(pieces), chunks_done=0)

        results = {}
        stop = threading.Event()   # set when the job ends early: no piece starts after that
        suffix, mime = ('.flac', 'audio/flac') if builtin else ('.ogg', 'audio/ogg')

        def work(index):
            # Only the HTTP side runs here; the database is touched by this thread alone.
            if stop.is_set():
                return index, None
            start, end = pieces[index]
            path = _encode(pcm, start, end, tmp, index, opus=not builtin)
            try:
                with open(path, 'rb') as f:
                    body = f.read()
            finally:
                os.remove(path)
            return index, _send(url, headers, builtin, f'piece-{index}{suffix}', mime, body, options, stop)

        pool = ThreadPoolExecutor(max_workers=max(1, min(conf.parallel, 8)))
        try:
            pending = {pool.submit(work, i) for i in range(len(pieces))}
            while pending:
                finished, pending = wait(pending, timeout=ALIVE_CHECK, return_when=FIRST_COMPLETED)
                _alive(job_id)   # also while a long piece is out
                for f in finished:
                    index, answer = f.result()   # the first failure ends the job
                    results[index] = answer
                if finished:
                    _write(job_id, chunks_done=len(results))
        finally:
            stop.set()
            pool.shutdown(wait=True, cancel_futures=True)
    return build(pieces, [results[i] for i in range(len(pieces))], options['model'], conf.language)


def _download(url, path, job_id):
    size = 0
    with open(path, 'wb') as out, requests.get(
            url, stream=True, timeout=(10, 60), headers={'User-Agent': 'diora transcription'}) as r:
        if r.status_code >= 400:
            raise Failed(f'Download: HTTP {r.status_code}')
        if int(r.headers.get('Content-Length') or 0) > silence.MAX_BYTES:
            raise Failed('too large')
        for n, chunk in enumerate(r.iter_content(1 << 20)):
            size += len(chunk)
            if size > silence.MAX_BYTES:
                raise Failed('too large')
            out.write(chunk)
            if n % 8 == 0:
                _alive(job_id)
    if size == 0:
        raise Failed('Download: leere Datei')


def _run_ffmpeg(args, input=None, timeout=DECODE_TIMEOUT):
    import imageio_ffmpeg
    proc = subprocess.run(
        [imageio_ffmpeg.get_ffmpeg_exe(), '-nostdin', '-hide_banner', '-loglevel', 'error', *args],
        input=input, stdout=subprocess.PIPE, stderr=subprocess.PIPE, timeout=timeout)
    if proc.returncode != 0:
        err = proc.stderr.decode('utf-8', 'replace').strip().splitlines()
        raise Failed('ffmpeg: ' + (err[-1] if err else f'exit {proc.returncode}')[:200])


def _decode(source, pcm):
    """The audio at `source` as raw 16 kHz mono 16-bit samples in the file `pcm`."""
    _run_ffmpeg(['-vn', '-sn', '-dn', '-i', source, '-ac', '1', '-ar', str(RATE), '-f', 's16le', '-y', pcm])


def pauses_of(pcm):
    """The pauses of the decoded audio as [(middle, length)] in seconds, found with
    silence.py's rules (what it cuts out of a pause is centred on the pause)."""
    size = int(RATE * WINDOW) * 2
    levels = []
    with open(pcm, 'rb') as f:
        while True:
            data = f.read(size * 100000)
            if not data:
                break
            for i in range(0, len(data) - size + 1, size):
                levels.append(silence._level_db(data[i:i + size]))
    return [((a + b) / 2, b - a) for a, b in silence.cuts(levels)]


def plan(total, pauses):
    """[(start, end)] in seconds covering 0..total: pieces of about CHUNK_SECONDS, each cut at
    the longest pause within CUT_WINDOW of where it would fall (right there when there is none)."""
    cuts = []
    last = 0.0
    k = 1
    while k * CHUNK_SECONDS < total - MIN_TAIL:
        target = k * CHUNK_SECONDS
        k += 1
        near = [p for p in pauses if abs(p[0] - target) <= CUT_WINDOW and p[0] > last + 1]
        at = max(near, key=lambda p: p[1])[0] if near else target
        if at <= last or at >= total:
            continue
        cuts.append(at)
        last = at
    edges = [0.0, *cuts, total]
    return [(edges[i], edges[i + 1]) for i in range(len(edges) - 1)]


def _encode(pcm, start, end, tmp, index, opus):
    """The piece [start, end) of the decoded audio as FLAC or Opus (ogg), in a file."""
    with open(pcm, 'rb') as f:
        f.seek(int(round(start * RATE)) * 2)
        data = f.read(int(round((end - start) * RATE)) * 2)
    out = os.path.join(tmp, f'piece-{index}.' + ('ogg' if opus else 'flac'))
    codec = ['-c:a', 'libopus', '-b:a', OPUS_BITRATE, '-f', 'ogg'] if opus else ['-c:a', 'flac', '-f', 'flac']
    _run_ffmpeg(['-f', 's16le', '-ar', str(RATE), '-ac', '1', '-i', 'pipe:0', *codec, '-y', out], input=data)
    return out


def _post(url, headers, builtin, name, mime, body, options):
    """One request; the parsed answer. Permanent for a refusal that won't change, RuntimeError
    for trouble that may pass."""
    files = {'file': (name, body, mime)}
    try:
        if builtin:
            # The operator's own instance (DIORA_WHISPER_URL) may sit on a private network.
            try:
                r = requests.post(url, headers=headers, files=files, data=options, timeout=UPLOAD_TIMEOUT)
            except requests.ConnectionError as e:
                raise Unreachable(e.__class__.__name__)
            if r.status_code >= 400:
                if r.status_code < 500 and r.status_code not in (408, 429):
                    raise Permanent(f'HTTP {r.status_code}')
                raise RuntimeError(f'HTTP {r.status_code}')
        else:
            try:
                r = chapters._open(url, headers, method='POST', timeout=UPLOAD_TIMEOUT, files=files, data=options)
            except chapters.ChaptersInvalid as e:
                raise Permanent(str(e))
            except chapters.ChaptersUnavailable as e:
                raise RuntimeError(str(e))
        try:
            return r.json()
        finally:
            r.close()
    except requests.RequestException as e:
        raise RuntimeError(e.__class__.__name__)
    except ValueError:
        raise RuntimeError('Antwort ist kein JSON')


def _send(url, headers, builtin, name, mime, body, options, stop):
    """A piece to the instance, retried up to RETRIES times with a growing pause."""
    for attempt in range(RETRIES + 1):
        if stop.is_set():
            raise Cancelled()
        try:
            return _post(url, headers, builtin, name, mime, body, options)
        except Permanent as e:
            raise Failed(f'Whisper: {e}')
        except Unreachable:
            if attempt >= RETRIES:
                raise
            time.sleep(RETRY_PAUSES[min(attempt, len(RETRY_PAUSES) - 1)])
        except RuntimeError as e:
            if attempt >= RETRIES:
                raise Failed(f'Whisper: {e}')
            log.info('whisper piece %s failed (%s), retrying', name, e)
            time.sleep(RETRY_PAUSES[min(attempt, len(RETRY_PAUSES) - 1)])


def build(pieces, answers, model, language):
    """{segments, source} from the pieces' places and the instance's verbose_json answers."""
    fragments = []
    detected = ''
    for (start, end), answer in zip(pieces, answers):
        if not isinstance(answer, dict):
            continue
        detected = detected or (answer.get('language') or '')
        segments = answer.get('segments')
        if not isinstance(segments, list):   # a plain {"text": …} answer: one segment for the piece
            text = (answer.get('text') or '').strip()
            segments = [{'start': 0, 'end': end - start, 'text': text}] if text else []
        for s in segments:
            if not isinstance(s, dict):
                continue
            text = ' '.join(str(s.get('text') or '').split())
            try:
                a, b = float(s.get('start', 0)), float(s.get('end', 0))
            except (TypeError, ValueError):
                continue
            if text:
                fragments.append((start + a, start + b, '', text))
    merged = transcripts.merge(fragments)
    if not merged:
        raise Failed('Kein Text erkannt')
    return {'segments': merged, 'source': {'type': 'whisper', 'model': model, 'language': language or detected}}