import os
import time
import uuid
import asyncio
import tempfile
import shutil
import threading
import torch
import torchaudio
from fastapi import FastAPI, Request, HTTPException, BackgroundTasks

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())

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

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 = 30 * 60  # کارهای قدیمی بعد از ۳۰ دقیقه پاک میشن
STAGE_TIMEOUT_SEC = 240  # سقف زمانی برای هر مرحله‌ی سنگین

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


def new_job():
    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,
        }
    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 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):
    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):
    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):
    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):
    return diarize_pipeline(audio_path, max_speakers=8)


def run_alignment_pass(segments_raw, align_waveform, 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(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


class ProcessingTimeout(Exception):
    pass


def run_blocking_with_timeout(job_id, blocking_fn, *args, status, detail, max_seconds=STAGE_TIMEOUT_SEC):
    """
    نسخه‌ی هم‌زمان (نه async) - چون این خودش داخل یه thread پس‌زمینه اجرا میشه،
    نیازی به هماهنگی با event loop نداره. وضعیت Job رو مستقیم آپدیت می‌کنه
    تا polling بتونه ببینتش.
    """
    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"]
    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, info = run_blocking_with_timeout(
            job_id, run_transcription, audio_to_transcribe,
            status="transcribing", detail="مدل large-v3 در حال شنیدن صداست",
        )
        update_job(job_id, 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)

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

        segments_raw_list = list(segments_raw)
        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:
        cleanup_expired_jobs()


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


@app.post("/upload/start")
async def upload_start():
    cleanup_expired_jobs()
    job_id = new_job()
    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