import os
import time
import uuid
import asyncio
import tempfile
import shutil
import threading
import logging
import torch
import torchaudio
from fastapi import FastAPI, Request, HTTPException, BackgroundTasks
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


class HealthAndStatusFilter(logging.Filter):
    """جلوی چاپ لاگ‌های /status/ و /health رو می‌گیره تا کنسول شلوغ نشه (خودِ درخواست‌ها عادی پردازش میشن)"""
    def filter(self, record):
        message = record.getMessage()
        return "/status/" not in message and "/health" not in message


logging.getLogger("uvicorn.access").addFilter(HealthAndStatusFilter())

torch.set_num_threads(2)  # جلوی رقابت بی‌فایده‌ی رشته‌ها روی CPU محدود رو می‌گیره

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
JOB_TTL_SECONDS = 3 * 60 * 60  # کارهای قدیمی بعد از ۳ ساعت پاک میشن (برای فایل‌های سنگین که ممکنه پردازششون طولانی بشه)
STAGE_TIMEOUT_SEC = 3600  # سقف زمانی برای هر مرحله‌ی سنگین (۲۰ دقیقه - فقط محافظ آخر در برابر گیرکردن واقعی)

# ─── نگهداری وضعیت کارها (Jobs) توی حافظه ───────────────────────────────────
jobs = {}
jobs_lock = threading.Lock()


def new_job(single_speaker_mode=False):
    job_id = uuid.uuid4().hex
    tmp_dir = tempfile.mkdtemp(prefix=f"subflow_{job_id}_")
    with jobs_lock:
        jobs[job_id] = {
            "status": "uploading",
            "detail": "در انتظار دریافت فایل...",
            "created_at": time.time(),
            "tmp_dir": tmp_dir,
            "chunk_paths": [],
            "srt": None,
            "language": None,
            "duration": None,
            "error": None,
            "single_speaker_mode": single_speaker_mode,
        }
    return job_id


def get_job(job_id):
    with jobs_lock:
        job = jobs.get(job_id)
    if job is None:
        raise HTTPException(status_code=404, detail="این کار پیدا نشد یا منقضی شده")
    return job


def update_job(job_id, **kwargs):
    with jobs_lock:
        if job_id in jobs:
            jobs[job_id].update(kwargs)


def cleanup_expired_jobs():
    now = time.time()
    with jobs_lock:
        expired = [jid for jid, j in jobs.items() if now - j["created_at"] > JOB_TTL_SECONDS]
    for jid in expired:
        with jobs_lock:
            job = jobs.pop(jid, None)
        if job and os.path.isdir(job["tmp_dir"]):
            shutil.rmtree(job["tmp_dir"], ignore_errors=True)
        print(f"[cleanup] removed expired job {jid}")


# ─── منطق تشخیص گوینده و هم‌ترازی ────────────────────────────────────────────

def build_speaker_lookup(diarization):
    """
    لیست بازه‌های گوینده‌ها رو یه‌بار مرتب می‌کنه تا جست‌وجوی «این لحظه مال کدوم
    گوینده‌ست» برای هر جمله، به‌جای گشتن کامل لیست، سریع (با جست‌وجوی دودویی) انجام بشه.
    """
    turns = sorted(
        [(turn.start, turn.end, speaker) for turn, _, speaker in diarization.itertracks(yield_label=True)],
        key=lambda x: x[0],
    )
    return turns


def get_speaker_at(turns, time_sec):
    import bisect
    starts = [t[0] for t in turns]
    idx = bisect.bisect_right(starts, time_sec) - 1
    if idx >= 0:
        start, end, speaker = turns[idx]
        if start <= time_sec <= end:
            return speaker
    return None


def get_speaker_for_segment(turns, seg_start, seg_end):
    """
    به‌جای اعتماد به یه لحظه‌ی تنها، چند نقطه در طول جمله رو چک می‌کنه
    و گوینده‌ای که بیشترین رأی رو آورده برمی‌گردونه.
    """
    duration = seg_end - seg_start
    if duration <= 0:
        return get_speaker_at(turns, 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(turns, 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(turns, whisper_start, search_back_sec=4.0):
    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(turns, 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):
    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):
    """
    نکته‌ی مهم: model.transcribe فوراً برنمی‌گرده - یه نتیجه‌ی "تنبل" میده که کار
    واقعیش رو دیر انجام میده. برای اینکه کل زمان واقعی پردازش زیر نظارت سقف زمانی و
    پیام‌های زنده باشه، همینجا مجبورش می‌کنیم کارش رو کامل انجام بده (list کردن).
    """
    segments_raw, info = 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,
        ),
    )
    segments_list = list(segments_raw)
    return segments_list, info


class SingleSpeakerResult:
    """
    یه نتیجه‌ی «قلابی» که همون شکل خروجی diarize_pipeline رو تقلید می‌کنه (متد itertracks داره)
    ولی هیچ پردازش واقعی‌ای روی کل فایل انجام نداده - فقط می‌گه «همه‌جا همون یه گوینده‌ست».
    فقط وقتی استفاده میشه که خودِ کاربر صراحتاً بگه فایل تک‌گوینده‌ست.
    """
    def __init__(self, duration):
        self.duration = duration

    def itertracks(self, yield_label=False):
        class FakeTurn:
            def __init__(self, start, end):
                self.start = start
                self.end = end
        yield FakeTurn(0, self.duration), None, "SPEAKER_00"


def run_diarization(audio_path, single_speaker_mode=False, job_id=None):
    if single_speaker_mode:
        print("  [diarization] single-speaker mode selected by user — skipping full diarization entirely")
        info = torchaudio.info(audio_path)
        duration = info.num_frames / info.sample_rate
        return SingleSpeakerResult(duration)

    return diarize_pipeline(audio_path, max_speakers=8)


def run_alignment_pass(segments_raw, align_waveform, diarization):
    turns = build_speaker_lookup(diarization)
    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(turns, 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(turns, 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


class ProcessingTimeout(Exception):
    pass


def run_blocking_with_timeout(job_id, blocking_fn, *args, status, detail, max_seconds=STAGE_TIMEOUT_SEC):
    """
    یه تابع سنگین و بلاک‌کننده رو توی یه thread جدا اجرا می‌کنه، و تا وقتی تموم بشه،
    هر چند ثانیه یه پیام "هنوز کار می‌کنم" روی همون status می‌فرسته - این هم جلوی
    قطع‌شدن اتصال به‌خاطر سکوت طولانی رو می‌گیره، هم UI رو زنده نگه می‌داره.
    اگه از max_seconds بیشتر طول بکشه، به‌جای قفل ابدی، خطای مشخص میده.
    """
    result_holder = {}
    error_holder = {}

    def runner():
        try:
            result_holder["value"] = blocking_fn(*args)
        except Exception as e:
            error_holder["error"] = e

    update_job(job_id, status=status, detail=detail)
    t = threading.Thread(target=runner, daemon=True)
    t.start()

    elapsed = 0
    while t.is_alive():
        t.join(timeout=3)
        if t.is_alive():
            elapsed += 3
            if elapsed >= max_seconds:
                raise ProcessingTimeout(f"مرحله‌ی «{detail}» بیشتر از {max_seconds} ثانیه طول کشید و متوقف شد.")
            update_job(job_id, status=status, detail=f"{detail} ({elapsed} ثانیه)")

    if "error" in error_holder:
        raise error_holder["error"]
    return result_holder.get("value")


def process_job(job_id):
    """کل پایپلاین پردازش - این تابع توی یه thread پس‌زمینه اجرا میشه، نه توی درخواست HTTP."""
    job = get_job(job_id)
    tmp_dir = job["tmp_dir"]
    single_speaker_mode = job.get("single_speaker_mode", False)
    tmp_path = os.path.join(tmp_dir, "input.wav")
    enhanced_path = None

    try:
        with open(tmp_path, "wb") as f:
            for chunk_path in job["chunk_paths"]:
                with open(chunk_path, "rb") as c:
                    f.write(c.read())
        print(f"Job {job_id}: assembled {len(job['chunk_paths'])} chunks into {tmp_path}")

        audio_to_transcribe = tmp_path

        if USE_DEEPFILTERNET:
            update_job(job_id, status="cleaning", detail="در حال حذف نویز پس‌زمینه")
            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 = os.path.join(tmp_dir, "enhanced.wav")
            save_audio(enhanced_path, enhanced_audio, df_state.sr())
            audio_to_transcribe = enhanced_path

        segments_raw_list, info = run_blocking_with_timeout(
            job_id, run_transcription, audio_to_transcribe,
            status="transcribing", detail="مدل large-v3 در حال شنیدن صداست",
        )
        update_job(job_id, detail=f"{len(segments_raw_list)} جمله تشخیص داده شد")

        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)

        if single_speaker_mode:
            diarization = run_diarization(audio_to_transcribe, single_speaker_mode=True)
            update_job(job_id, status="aligning", detail="تشخیص گوینده رد شد (حالت تک‌گوینده)")
        else:
            diarization = run_blocking_with_timeout(
                job_id, run_diarization, audio_to_transcribe, False, job_id,
                status="aligning", detail="در حال تشخیص گوینده‌ها",
            )
        update_job(job_id, detail="گوینده‌ها شناسایی شدند")

        segments = run_blocking_with_timeout(
            job_id, run_alignment_pass, segments_raw_list, align_waveform, diarization,
            status="aligning", detail=f"در حال هماهنگ‌سازی زمان‌بندی {len(segments_raw_list)} جمله",
        )
        update_job(job_id, detail=f"{len(segments)} خط شناسایی شد")

        print(f"Job {job_id}: got {len(segments)} segments")
        for i, seg in enumerate(segments):
            print(f"  seg[{i}]: '{seg['text']}' | speaker={seg['speaker']} | dur={seg['end']-seg['start']:.2f}s")

        update_job(job_id, status="processing", detail=f"در حال چیدمان {len(segments)} خط زیرنویس")
        srt_content = process(segments)
        update_job(job_id, detail="زیرنویس ساخته شد")

        update_job(
            job_id,
            status="done",
            detail="زیرنویس آماده شد",
            srt=srt_content,
            language=info.language,
            duration=info.duration,
        )
        print(f"Job {job_id}: done!")

    except ProcessingTimeout as e:
        print(f"Job {job_id}: TIMEOUT: {e}")
        update_job(job_id, status="error", error=str(e))
    except Exception as e:
        print(f"Job {job_id}: ERROR: {type(e).__name__}: {e}")
        import traceback
        traceback.print_exc()
        update_job(job_id, status="error", error=str(e))
    finally:
        pass  # پاک‌سازی جداگانه و فقط موقع شروع کارهای جدید انجام میشه، نه بلافاصله بعد از اتمام همین کار


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


@app.post("/upload/start")
async def upload_start(request: Request):
    cleanup_expired_jobs()
    body = {}
    try:
        body = await request.json()
    except Exception:
        pass
    single_speaker_mode = bool(body.get("single_speaker_mode", False))
    job_id = new_job(single_speaker_mode)
    return {"job_id": job_id}


@app.post("/upload/chunk/{job_id}")
async def upload_chunk(job_id: str, request: Request):
    job = get_job(job_id)
    chunk_bytes = await request.body()
    chunk_index = len(job["chunk_paths"])
    chunk_path = os.path.join(job["tmp_dir"], f"chunk_{chunk_index:06d}.bin")
    with open(chunk_path, "wb") as f:
        f.write(chunk_bytes)
    with jobs_lock:
        jobs[job_id]["chunk_paths"].append(chunk_path)
    return {"received": chunk_index}


@app.post("/upload/finish/{job_id}")
async def upload_finish(job_id: str, background_tasks: BackgroundTasks):
    job = get_job(job_id)
    total_bytes = sum(os.path.getsize(p) for p in job["chunk_paths"])
    print(f"Job {job_id}: upload finished, {len(job['chunk_paths'])} chunks, {total_bytes} bytes total")
    update_job(job_id, status="receiving", detail="فایل آماده‌ی پردازش شد")
    background_tasks.add_task(process_job, job_id)
    return {"status": "processing_started"}


@app.get("/status/{job_id}")
async def status(job_id: str):
    job = get_job(job_id)
    response = {
        "status": job["status"],
        "detail": job["detail"],
    }
    if job["status"] == "done":
        response["srt"] = job["srt"]
        response["language"] = job["language"]
        response["duration"] = job["duration"]
    elif job["status"] == "error":
        response["message"] = job["error"]
    return response