import os
import asyncio
import tempfile
import torch
import torchaudio
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.middleware.cors import CORSMiddleware
from faster_whisper import WhisperModel
from df.enhance import enhance, init_df, load_audio, save_audio
from pyannote.audio import Pipeline as DiarizationPipeline
from pipeline.postprocess import process

app = FastAPI()

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

print("Loading Whisper model...", flush=True)
model = WhisperModel("large-v3", device="cpu", compute_type="int8")
print("Whisper model ready!")

print("Loading DeepFilterNet model...", flush=True)
df_model, df_state, _ = init_df()
print("DeepFilterNet model ready!")

USE_DEEPFILTERNET = False  # غیرفعال - شواهد نشون داد کمکی نمی‌کنه و بعضی صداها رو خراب می‌کنه

print("Loading word-alignment model...", flush=True)
align_bundle = torchaudio.pipelines.MMS_FA
align_model = align_bundle.get_model().to("cpu")
align_labels = align_bundle.get_labels()
align_dictionary = {c: i for i, c in enumerate(align_labels)}
print("Word-alignment model ready!")

print("Loading speaker-diarization model...", flush=True)
HF_TOKEN = os.environ.get("HF_TOKEN")
diarize_pipeline = DiarizationPipeline.from_pretrained(
    "pyannote/speaker-diarization-3.1", use_auth_token=HF_TOKEN
)
print("Speaker-diarization model ready!")

MIN_NOISE_BLOCK_SEC = 1.0  # حداقل طول یه بلوک بی‌گوینده تا «نویز» حساب بشه
HEARTBEAT_SEC = 3  # هر چند ثانیه یه پیام زنده بفرستیم تا اتصال ساکت نمونه


def get_speaker_at(diarization, time_sec):
    for turn, _, speaker in diarization.itertracks(yield_label=True):
        if turn.start <= time_sec <= turn.end:
            return speaker
    return None


def get_speaker_for_segment(diarization, seg_start, seg_end):
    """
    به‌جای اعتماد به یه لحظه‌ی تنها (که ممکنه نزدیک مرز بین دو گوینده باشه)،
    چند نقطه در طول جمله رو چک می‌کنه و گوینده‌ای که بیشترین رأی رو آورده
    برمی‌گردونه - این دقت رو بهتر می‌کنه ولی صد در صد تضمین‌شده نیست.
    """
    duration = seg_end - seg_start
    if duration <= 0:
        return get_speaker_at(diarization, seg_start)

    sample_points = [seg_start + duration * frac for frac in (0.2, 0.4, 0.5, 0.6, 0.8)]
    votes = {}
    for t in sample_points:
        speaker = get_speaker_at(diarization, t)
        if speaker is not None:
            votes[speaker] = votes.get(speaker, 0) + 1

    if not votes:
        return None
    return max(votes, key=votes.get)


def find_leading_noise_cutoff(diarization, whisper_start, search_back_sec=4.0):
    """
    برمی‌گرده به عقب از شروعِ حدسیِ Whisper، دنبال یه بلوک نویز/بی‌گوینده‌ی
    طولانی (حداقل MIN_NOISE_BLOCK_SEC) که دقیقاً قبل از این جمله باشه.
    نوسان‌های خیلی کوتاه (کمتر از STABLE_SPEAKER_SEC) که یه گوینده برای یه
    لحظه‌ی گذرا ظاهر میشه رو نادیده می‌گیره - فقط با دیدنِ یه گوینده‌ی
    "پایدار" واقعاً متوقف میشه.
    اگه پیدا کرد، زمان پایانِ اون بلوک نویز رو برمی‌گردونه (یعنی «از اینجا به بعد
    اعتماد کن»)؛ اگه پیدا نکرد، None برمی‌گردونه (یعنی چیزی رو نبُر).
    """
    STABLE_SPEAKER_SEC = 0.2
    window_start = max(0, whisper_start - search_back_sec)
    step = 0.05
    t = whisper_start
    noise_end = None
    noise_start = None
    speaker_run_start = None

    while t > window_start:
        speaker = get_speaker_at(diarization, t)

        if speaker is None:
            speaker_run_start = None
            if noise_end is None:
                noise_end = t
            noise_start = t
        else:
            if speaker_run_start is None:
                speaker_run_start = t
            run_len = speaker_run_start - t
            if run_len >= STABLE_SPEAKER_SEC:
                break
            if noise_end is None:
                noise_end = t
            noise_start = t

        t -= step

    if noise_start is not None and noise_end is not None:
        block_len = noise_end - noise_start
        if block_len >= MIN_NOISE_BLOCK_SEC:
            return noise_end
    return None


def realign_word_timestamps(waveform_16k, words, whisper_start, whisper_end, cutoff_start=None):
    """
    برای یه جمله، به‌جای زمان‌بندی حدسی خودِ Whisper، با یه مدل تخصصی جدا
    (wav2vec2 از طریق torchaudio) دقیقاً چک می‌کنه هر کلمه از کجای صدا گفته شده.
    اگه cutoff_start داده بشه (یعنی یه بلوک نویز قبل از جمله شناسایی شده)،
    جستجو از همون نقطه شروع میشه، نه از قبل‌ترش.
    اگه به هر دلیلی هم‌ترازی جواب نداد، همون زمان‌بندی قبلی Whisper دست‌نخورده می‌مونه.
    """
    try:
        margin = 0.5
        raw_start = whisper_start - margin
        effective_start = max(raw_start, cutoff_start) if cutoff_start is not None else raw_start

        start_sample = max(0, int(effective_start * 16000))
        end_sample = int((whisper_end + margin) * 16000)
        clip = waveform_16k[:, start_sample:end_sample]

        clean_words = []
        for w in words:
            token = "".join(ch for ch in w["word"].lower() if ch in align_dictionary)
            if token:
                clean_words.append((w, token))

        if not clean_words:
            return words

        transcript = "".join(tok for _, tok in clean_words)
        tokens = [align_dictionary[c] for c in transcript]

        with torch.inference_mode():
            emission, _ = align_model(clip)

        targets = torch.tensor([tokens], dtype=torch.int32)
        input_lengths = torch.tensor([emission.shape[1]])
        target_lengths = torch.tensor([len(tokens)])

        aligned_tokens, scores = torchaudio.functional.forced_align(
            emission, targets, input_lengths, target_lengths, blank=0
        )
        spans = torchaudio.functional.merge_tokens(aligned_tokens[0], scores[0])

        ratio = clip.shape[1] / emission.shape[1] / 16000
        new_words = []
        span_idx = 0
        for w, tok in clean_words:
            n_chars = len(tok)
            matched = [s for s in spans[span_idx:] if s.token != 0][:n_chars]
            if len(matched) < n_chars:
                new_words.append(w)
                continue
            w_start = matched[0].start * ratio + effective_start
            w_end = matched[-1].end * ratio + effective_start
            new_words.append({"word": w["word"], "start": round(max(0, w_start), 3), "end": round(max(0, w_end), 3)})
            span_idx += n_chars

        return new_words if len(new_words) == len(words) else words
    except Exception as e:
        print(f"  [align] skipped due to: {type(e).__name__}: {e}")
        return words


def run_transcription(audio_path):
    """فقط تشخیص گفتار - بلاک‌کننده، توی یه thread جدا اجرا میشه"""
    return model.transcribe(
        audio_path,
        language="en",
        word_timestamps=True,
        vad_filter=True,
        vad_parameters=dict(
            min_silence_duration_ms=300,
            speech_pad_ms=200,
            threshold=0.3,
        ),
    )


def run_diarization(audio_path):
    """فقط تشخیص گوینده - بلاک‌کننده، توی یه thread جدا اجرا میشه"""
    return diarize_pipeline(audio_path)


def run_alignment_pass(segments_raw, align_waveform, diarization):
    """
    برای همه‌ی جمله‌ها، هم‌ترازی + تشخیص گوینده + برش نویز رو اجرا می‌کنه.
    بلاک‌کننده‌ست (حلقه‌ی سنگین روی هر جمله)، توی یه thread جدا اجرا میشه.
    """
    segments = []
    for seg in segments_raw:
        words = list(seg.words)
        if not words:
            continue
        word_dicts = [{"word": w.word, "start": w.start, "end": w.end} for w in words]

        cutoff = find_leading_noise_cutoff(diarization, word_dicts[0]["start"])
        if cutoff is not None:
            print(f"  [noise-trim] '{seg.text.strip()[:40]}...' cutoff moved from {word_dicts[0]['start']:.2f}s to {cutoff:.2f}s")

        realigned = realign_word_timestamps(
            align_waveform, word_dicts, word_dicts[0]["start"], word_dicts[-1]["end"], cutoff_start=cutoff
        )

        speaker = get_speaker_for_segment(diarization, realigned[0]["start"], realigned[-1]["end"])

        segments.append({
            "text": seg.text.strip(),
            "start": realigned[0]["start"],
            "end": realigned[-1]["end"],
            "words": realigned,
            "speaker": speaker,
        })
    return segments


async def run_with_heartbeat(websocket, blocking_fn, *args, status, detail):
    """
    یه تابع سنگین و بلاک‌کننده رو توی یه thread جدا اجرا می‌کنه، و تا وقتی تموم بشه،
    هر چند ثانیه یه پیام "هنوز کار می‌کنم" روی همون status می‌فرسته - این هم جلوی
    قطع‌شدن اتصال به‌خاطر سکوت طولانی رو می‌گیره، هم UI رو زنده نگه می‌داره.
    """
    task = asyncio.create_task(asyncio.to_thread(blocking_fn, *args))
    elapsed = 0
    await websocket.send_json({"status": status, "detail": detail})
    while not task.done():
        try:
            await asyncio.wait_for(asyncio.shield(task), timeout=HEARTBEAT_SEC)
        except asyncio.TimeoutError:
            elapsed += HEARTBEAT_SEC
            await websocket.send_json({
                "status": status,
                "detail": f"{detail} ({elapsed} ثانیه)",
            })
    return await task


@app.get("/health")
async def health():
    return {"status": "ok", "message": "SubFlow backend is running"}


@app.websocket("/transcribe")
async def transcribe(websocket: WebSocket):
    await websocket.accept()
    tmp_path = None
    enhanced_path = None

    try:
        await websocket.send_json({"status": "receiving", "detail": "در انتظار دریافت فایل..."})
        print("Step 1: waiting for bytes")
        audio_bytes = await websocket.receive_bytes()
        print(f"Step 2: received {len(audio_bytes)} bytes")
        size_mb = len(audio_bytes) / (1024 * 1024)
        await websocket.send_json({"status": "receiving", "detail": f"{size_mb:.1f} مگابایت دریافت شد"})

        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp:
            tmp.write(audio_bytes)
            tmp_path = tmp.name
        print(f"Step 3: saved to {tmp_path}")
        await websocket.send_json({"status": "receiving", "detail": "فایل آماده‌ی پردازش شد"})

        audio_to_transcribe = tmp_path

        if USE_DEEPFILTERNET:
            await websocket.send_json({"status": "cleaning", "detail": "در حال حذف نویز پس‌زمینه"})
            print("Step 3.5: cleaning audio with DeepFilterNet")

            audio_data, _ = load_audio(tmp_path, sr=df_state.sr())
            enhanced_audio = enhance(df_model, df_state, audio_data, atten_lim_db=20)
            enhanced_path = tmp_path.replace(".wav", "_enhanced.wav")
            save_audio(enhanced_path, enhanced_audio, df_state.sr())
            audio_to_transcribe = enhanced_path
            print(f"Step 3.6: cleaned audio saved to {enhanced_path}")
        else:
            print("Step 3.5: skipping DeepFilterNet (disabled for A/B test)")

        print("Step 4: starting transcription")
        segments_raw, info = await run_with_heartbeat(
            websocket, run_transcription, audio_to_transcribe,
            status="transcribing", detail="مدل large-v3 در حال شنیدن صداست",
        )
        print("Step 5: transcription done")
        await websocket.send_json({"status": "transcribing", "detail": "تشخیص گفتار کامل شد"})

        align_waveform, sr = torchaudio.load(audio_to_transcribe)
        if sr != 16000:
            align_waveform = torchaudio.functional.resample(align_waveform, sr, 16000)
        if align_waveform.shape[0] > 1:
            align_waveform = align_waveform.mean(dim=0, keepdim=True)

        print("Step 5.5: running speaker diarization")
        diarization = await run_with_heartbeat(
            websocket, run_diarization, audio_to_transcribe,
            status="aligning", detail="در حال تشخیص گوینده‌ها",
        )
        print("Step 5.6: diarization done, now aligning words per segment")
        await websocket.send_json({"status": "aligning", "detail": "گوینده‌ها شناسایی شدند"})

        segments_raw_list = list(segments_raw)
        segments = await run_with_heartbeat(
            websocket, run_alignment_pass, segments_raw_list, align_waveform, diarization,
            status="aligning", detail=f"در حال هماهنگ‌سازی زمان‌بندی {len(segments_raw_list)} جمله",
        )

        print(f"Step 6: got {len(segments)} segments")
        await websocket.send_json({"status": "aligning", "detail": f"{len(segments)} خط شناسایی شد"})
        for i, seg in enumerate(segments):
            print(f"  seg[{i}]: '{seg['text']}' | speaker={seg['speaker']} | dur={seg['end']-seg['start']:.2f}s")

        await websocket.send_json({
            "status": "processing",
            "detail": f"در حال چیدمان {len(segments)} خط زیرنویس",
        })

        srt_content = process(segments)
        print(f"Step 7: SRT built")
        await websocket.send_json({"status": "processing", "detail": "زیرنویس ساخته شد"})

        await websocket.send_json({
            "status": "done",
            "detail": "زیرنویس آماده شد",
            "srt": srt_content,
            "language": info.language,
            "duration": info.duration,
        })
        print("Step 8: done sent!")
        await asyncio.sleep(1)

    except WebSocketDisconnect:
        print("WebSocket disconnected by client")
    except Exception as e:
        print(f"ERROR: {type(e).__name__}: {e}")
        import traceback
        traceback.print_exc()
        try:
            await websocket.send_json({"status": "error", "message": str(e)})
        except Exception:
            pass
    finally:
        if tmp_path and os.path.exists(tmp_path):
            os.remove(tmp_path)
        if enhanced_path and os.path.exists(enhanced_path):
            os.remove(enhanced_path)
        try:
            await websocket.close()
        except Exception:
            pass