import shutil
import sys
import os
import time
import requests
from requests.exceptions import ProxyError, ConnectionError, ConnectTimeout, ReadTimeout
import random
import json
import zipfile
import torch
import shlex
import torchaudio
import subprocess
import patoolib
import numpy as np
import pysrt
from silero_vad import get_speech_timestamps, read_audio, save_audio
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
from celery_tasks.celery import app
from config import TELEGRAM_API,SUBSOURCE_BASE_URL,SUBSOURCE_API_KEY,MOVIEO_TELEGRAM_API
from base.utils.funcs import identify_seasons_from_irregular_filename,identify_episodes_from_irregular_filename,identify_release,identify_episodes,identify_seasons,run_command,log_request,send_request
from base.utils import MovieoDatabase, YOUR_REQUEST_CLOSED, download_video, get_video_duration_seconds, \
    get_random_second, convert_to_mkv
from base.utils.subtitle_editor import *

def detect_archive_type_from_file(path):
    with open(path, "rb") as f:
        sig = f.read(8)

    if sig.startswith(b'PK\x03\x04'):
        return "zip"
    elif sig.startswith(b'Rar!\x1A\x07'):
        return "rar"
    return "unknown"


def convert_rar_to_zip(subtitle_id):
    temp_rar = f"{subtitle_id}.rar"
    os.rename(f'{DOWNLOAD_DIR}{subtitle_id}.zip', temp_rar)

    try:
        temp_dir = "temp_extraction"
        patoolib.extract_archive(temp_rar, outdir=temp_dir)

        files = [os.path.join(temp_dir, f) for f in os.listdir(temp_dir)]
        patoolib.create_archive(f'{DOWNLOAD_DIR}{subtitle_id}.zip', files)

    finally:
        if os.path.exists(temp_rar): os.remove(temp_rar)
        if os.path.exists(temp_dir): shutil.rmtree(temp_dir)

def request_to_subsource(endpoint: str, query: str, headers: dict):
    proxy = "http://apiec803a20e08d2684:driS9IZ2@res.proxy-seller.com:10000"
    proxies = {"http": proxy, "https": proxy}

    try:
        resp = requests.get(f'{SUBSOURCE_BASE_URL}{endpoint}{query}', headers=headers, proxies=proxies, timeout=15)
    except (ProxyError, ConnectionError, ConnectTimeout, ReadTimeout) as e:
        print(f"Proxy/Connection error ({e}) — retrying request with new IP...")
        time.sleep(5)
        return request_to_subsource(endpoint, query, headers)

    if resp.status_code == 200:
        resp_headers = {k.lower(): v.lower() for k, v in resp.headers.items()}
        cf_mitigated = resp_headers.get("cf-mitigated")
        connection_type = resp_headers.get("connection")
        cache_control = resp_headers.get("cache-control", "")
        is_challenge = (
            cf_mitigated == "challenge" or
            ("no-cache" in cache_control and "private" in cache_control and connection_type == "close")
        )

        if is_challenge:
            print("Cloudflare challenge detected — retrying request with new IP...")
            return request_to_subsource(endpoint, query, headers)

        return resp
    else:
        print(f"{resp.status_code} error detected — retrying request with new IP...")
        time.sleep(5)
        return request_to_subsource(endpoint, query, headers)


def get_sub_times(sub_path):
    subs = pysrt.open(sub_path, encoding='utf-8')
    starts = [s.start.ordinal / 1000.0 for s in subs if s.text.strip()]
    ends = [s.end.ordinal / 1000.0 for s in subs if s.text.strip()]
    return starts, ends

def compare_subs_timing(sub_times_list):
    n = len(sub_times_list)
    similarity_matrix = np.zeros((n, n))
    for i in range(n):
        for j in range(i + 1, n):
            min_len = min(len(sub_times_list[i][0]), len(sub_times_list[j][0]))
            if min_len == 0:
                diff = 9999
            else:
                diffs = np.abs(
                    np.array(sub_times_list[i][0][:min_len]) -
                    np.array(sub_times_list[j][0][:min_len])
                )
                diff = np.mean(diffs)
            similarity_matrix[i, j] = similarity_matrix[j, i] = diff
    return similarity_matrix

def find_outliers(similarity_matrix, threshold=2.0):
    mean_diffs = similarity_matrix.mean(axis=1)
    avg = np.mean(mean_diffs)
    std = np.std(mean_diffs)
    keep = [i for i, m in enumerate(mean_diffs) if m < avg + threshold * std]
    return keep, mean_diffs


def find_central_sub(mean_diffs):
    try:
        n = len(mean_diffs)
        distances = np.zeros(n)

        for i in range(n):
            distances[i] = np.mean([abs(mean_diffs[i] - mean_diffs[j]) for j in range(n) if j != i])

        central_idx = int(np.argmin(distances))
    except:
        central_idx = 0

    return central_idx


def calculate_smart_anchor_score(subtitle_path, vad_segments, silence_threshold=1.5, min_speech_duration=1.0, sync_window=0.5, sr=16000):
    try:
        subs = pysrt.open(subtitle_path)
    except Exception as e:
        print(f"Error opening subtitle: {e}")
        return 0,0,0
    try:
        sub_starts = []
        sub_ends = []
        valid_sub_objects = []

        isolated_sub_indices = set()

        last_sub_end_sec = 0.0

        temp_index = 0
        if subs:
            for sub in subs:
                if sub.text and len(sub.text.strip()) > 1:
                    start_seconds = sub.start.ordinal / 1000.0
                    end_seconds = sub.end.ordinal / 1000.0

                    gap_before_sub = start_seconds - last_sub_end_sec
                    if gap_before_sub >= silence_threshold:
                        isolated_sub_indices.add(temp_index)

                    sub_starts.append(start_seconds)
                    sub_ends.append(end_seconds)
                    valid_sub_objects.append(sub)

                    last_sub_end_sec = end_seconds
                    temp_index += 1
            sub_starts = np.array(sub_starts)

            audio_anchors = []
            last_vad_end_sec = 0.0

            sorted_vad = sorted(vad_segments, key=lambda x: x['start'])

            for seg in sorted_vad:
                start_sec = seg['start'] / sr
                end_sec = seg['end'] / sr

                silence_duration = start_sec - last_vad_end_sec
                speech_duration = end_sec - start_sec

                if silence_duration >= silence_threshold and speech_duration >= min_speech_duration:
                    audio_anchors.append(start_sec)

                last_vad_end_sec = end_sec

            if not audio_anchors:
                return 0,0,0

            matched_anchors = 0
            valid_comparisons = 0

            for anchor_time in audio_anchors:
                differences = np.abs(sub_starts - anchor_time)
                closest_idx = np.argmin(differences)
                min_diff = differences[closest_idx]

                if closest_idx not in isolated_sub_indices:
                    continue

                valid_comparisons += 1

                if min_diff <= sync_window:
                    matched_anchors += 1

            final_score = matched_anchors / valid_comparisons

            return round(final_score,2),valid_comparisons,matched_anchors
        else:
            return 0,0,0
    except:
        return 0,0,0

@app.task(queue='sync_subtitles')
def find_and_synchronize_subtitle(admin_user_id: int, request_id: int, request_type_str: str, imdb_id: str, title_type: int, file_name: str, file_id: str):
    # try:
    if request_type_str in ['Persian SoftSub','Persian HardSub']:
        subtitle_lang = 'farsi_persian'
    elif request_type_str in ['English SoftSub','English HardSub']:
        subtitle_lang = 'english'

    if 'SoftSub' in request_type_str:
        task_type = 1
    else:
        task_type = 2

    if title_type == 2:
        file_season = identify_seasons(file_name)
        _, _, file_episode = identify_episodes(file_name)
        query = f'?searchType=imdb&imdb={imdb_id}&season={file_season}'
    else:
        file_season = file_episode = None
        query = f'?searchType=imdb&imdb={imdb_id}'
    file_release_type = identify_release(file_name)

    print(f'\n**********************\n--- imdb id: {imdb_id}\n--- title type: {title_type}\n--- file release type: {file_release_type}\n--- file season: {file_season}\n--- file episode: {file_episode}\n**********************\n')

    headers = {"X-API-Key": SUBSOURCE_API_KEY}
    movie = request_to_subsource('movies/search', query, headers).json()
    try:
        subsource_id = movie['data'][0]['movieId']
    except:
        with MovieoDatabase() as db:
            user_request_db = db.get_users_requests('id',request_id,'status')
        request_status = user_request_db[0][0]
        if request_status != 3:
            chat_id, movieo_id, title_name = log_request(request_id, 3)
            data = {
                'chat_id': chat_id,
                'text': YOUR_REQUEST_CLOSED(request_type_str, title_name)
            }
            send_request(MOVIEO_TELEGRAM_API, 'sendMessage', data)
        exit()
    subtitles = request_to_subsource('subtitles', f'?movieId={subsource_id}&language={subtitle_lang}&limit=100', headers).json()['data']
    def extract_info(release_info):
        if title_type == 2:
            season = identify_seasons_from_irregular_filename(release_info)
            episode_found, _, episode = identify_episodes(release_info)
            if not episode_found:
                episode = identify_episodes_from_irregular_filename(release_info)
        else:
            season = episode = None
        release_type = identify_release(release_info)
        return season, episode, release_type
    subtitle_ids = []
    if not subtitles:
        with MovieoDatabase() as db:
            user_request_db = db.get_users_requests('id',request_id,'status')
        request_status = user_request_db[0][0]
        if request_status != 3:
            chat_id, movieo_id, title_name = log_request(request_id, 3)
            data = {
                'chat_id': chat_id,
                'text': YOUR_REQUEST_CLOSED(request_type_str, title_name)
            }
            send_request(MOVIEO_TELEGRAM_API, 'sendMessage', data)
        exit()
    for sub in subtitles:
        for release in sub['releaseInfo']:
            s_season, s_episode, _ = extract_info(release)
            files_count = sub['files'] or 0
            if (file_season == s_season and file_episode == s_episode) or (file_season == s_season and s_episode == None) or 'complete' in release.lower() or 'all episodes' in release.replace('.', ' ').lower() or files_count > 2:
                print(f'\n**********************\n--- subtitle {release} added.\n**********************')
                subtitle_ids.append(sub['subtitleId'])
                break
        if len(subtitle_ids) > 10:
            break

    if subtitle_ids:
        print(f'\n**********************\n--- {len(subtitle_ids)} subtitles download and extraction has begun. \n**********************')

        i = 1
        subtitles = []
        for subtitle_id in subtitle_ids:
            headers = {
                "X-API-Key": SUBSOURCE_API_KEY,
                "contentType": "application/zip",
                "contentDisposition": f"attachment; filename=\"{subtitle_id}.zip\"",
                "body": "ZIP file stream"
            }
            res = request_to_subsource('subtitles', f'/{subtitle_id}/download', headers)
            if res.status_code == 200:
                content_type = res.headers.get("Content-Type", "")
                if "text/html" in content_type.lower():
                    exit()
                else:
                    try:
                        with open(f'{DOWNLOAD_DIR}{subtitle_id}.zip', "wb") as f:
                            for chunk in res.iter_content(chunk_size=8192):
                                if chunk:
                                    f.write(chunk)
                        archive_type = detect_archive_type_from_file(f'{DOWNLOAD_DIR}{subtitle_id}.zip')
                        if archive_type == "rar":
                            convert_rar_to_zip(subtitle_id)
                        with zipfile.ZipFile(f'{DOWNLOAD_DIR}{subtitle_id}.zip', 'r') as zip_ref:
                            subtitle_names = zip_ref.namelist()
                            if len(subtitle_names) > 1:
                                for sub_name in subtitle_names:
                                    if title_type == 2:
                                        _, s_episode, _ = extract_info(sub_name)
                                        if file_episode == s_episode:
                                            subtitle = sub_name
                                            print(f'\n**********************\n--- chosen subtitle: {subtitle}\n**********************')
                                            break
                                    else:
                                        _, sub_ext = os.path.splitext(sub_name)
                                        if sub_ext in ['.srt','.ass','.ssa','.vtt']:
                                            subtitle = sub_name
                                            print(f'\n**********************\n--- chosen subtitle: {subtitle}\n**********************')
                                            break
                            else:
                                subtitle = subtitle_names[0]
                                print(f'\n**********************\n--- chosen subtitle: {subtitle}\n**********************')
                            zip_ref.extract(subtitle,DOWNLOAD_DIR)
                        os.remove(f'{DOWNLOAD_DIR}{subtitle_id}.zip')
                        _, subtitle_format = os.path.splitext(f'{subtitle}')
                        if subtitle_format != '.srt':
                            print(f'\n**********************\n--- subtitle format is {subtitle_format}. we have to change it to srt. \n**********************')
                            convert_to_srt(f'{DOWNLOAD_DIR}{subtitle}',subtitle_format)
                            subtitle = subtitle.replace(subtitle_format,'.srt')
                        os.rename(f'{DOWNLOAD_DIR}{subtitle}',f'{DOWNLOAD_DIR}{i}_{file_id}.srt')
                        print(f'\n**********************\n--- cleaning and preparation of the subtitles began.\n**********************')
                        subtitle = f'{i}_{file_id}.srt'
                        convert_to_utf8(subtitle)
                        repair_subtitle(subtitle)
                        if subtitle_lang == 'farsi_persian':
                            find_ads_in_subtitle(subtitle)
                        if os.path.exists(f'{SUBTITLE_DIR}{subtitle}'):
                            os.remove(f'{DOWNLOAD_DIR}{subtitle}')
                        else:
                            os.rename(f'{DOWNLOAD_DIR}{subtitle}',f'{SUBTITLE_DIR}{subtitle}')
                        subtitles.append(subtitle)
                    except:
                        pass
            i += 1
        print(f'\n**********************\n--- download {file_name} has begun. \n**********************')
        server_file_name, file_format = download_video(TELEGRAM_API,file_id,file_name)
        if server_file_name:
            original_audio_wav = f'{DOWNLOAD_DIR}{file_id}_original.wav'
            speech_audio_wav = f"{DOWNLOAD_DIR}speech_{file_id}.wav"
            print(f'\n**********************\n--- extracting audio from video file... \n**********************')
            cmd_start_time = f'ffprobe -i "{DOWNLOAD_DIR}{server_file_name}" -show_entries stream=start_time -select_streams a -v quiet -of csv="p=0"'
            audio_start_output = subprocess.check_output(shlex.split(cmd_start_time)).decode().strip()
            audio_start = float(audio_start_output.split('\n')[0])

            cmd_duration = f'ffprobe -i "{DOWNLOAD_DIR}{server_file_name}" -show_entries format=duration -v quiet -of csv="p=0"'
            video_duration_output = subprocess.check_output(shlex.split(cmd_duration)).decode().strip()
            video_duration = float(video_duration_output)


            if audio_start > 0:
                delay_ms = int(audio_start * 1000)
                af_filter = f"adelay={delay_ms}|{delay_ms},aresample=16000"
            elif audio_start < 0:
                trim_start = abs(audio_start)
                af_filter = f"atrim=start={trim_start},aresample=16000"
            else:
                af_filter = "aresample=16000"

            command_convert = (
                f'ffmpeg -i "{DOWNLOAD_DIR}{server_file_name}" '
                f'-vn -af "{af_filter}" -ac 1 -t {video_duration} "{original_audio_wav}" -y'
            )
            if not run_command(command_convert, "Initial conversion to WAV failed."):
                exit()
            print(f'\n**********************\n--- speech recognition (VAD) \n**********************')
            try:
                model, utils = torch.hub.load(
                    repo_or_dir='snakers4/silero-vad',
                    model='silero_vad',
                    force_reload=False
                )
                (get_speech_timestamps, save_audio, read_audio, VADIterator, collect_chunks) = utils

                wav = read_audio(original_audio_wav, sampling_rate=16000)
                sr = 16000
                speech_timestamps = get_speech_timestamps(
                    wav,
                    model,
                    sampling_rate=sr,
                    threshold=0.45,
                    min_speech_duration_ms=350,
                    min_silence_duration_ms=350,
                    speech_pad_ms=100
                )
            except Exception as e:
                print(f'\n**********************\n--- error loading or running Silero VAD model: {e} \n**********************')
                exit()
            output_wav = torch.zeros_like(wav)
            for segment in speech_timestamps:
                output_wav[segment['start']:segment['end']] = wav[segment['start']:segment['end']]
            save_audio(speech_audio_wav, output_wav, sampling_rate=sr)
            print(f'\n**********************\n--- syncing subtitles with Alass \n**********************')
            original_subtitles = []
            sub_times_list = []
            for subtitle in subtitles:
                shutil.copy2(f'{SUBTITLE_DIR}{subtitle}', f'{SUBTITLE_DIR}original_{subtitle}')
                original_subtitles.append(f'original_{subtitle}')
                command_alass = f"/usr/local/bin/alass \"{speech_audio_wav}\" \"{SUBTITLE_DIR}{subtitle}\" \"{SUBTITLE_DIR}fixed_{subtitle}\""
                run_command(command_alass, "Alass process was unsuccessful.")
                try:
                    subs = pysrt.open(f'{SUBTITLE_DIR}fixed_{subtitle}', encoding='utf-8')

                    def round_time_to_100ms(time_obj):
                        total_ms = time_obj.ordinal
                        rounded_ms = round(total_ms / 100) * 100
                        return pysrt.SubRipTime.from_ordinal(rounded_ms)

                    for s in subs:
                        s.start = round_time_to_100ms(s.start)
                        s.end = round_time_to_100ms(s.end)

                    subs.save(f'{SUBTITLE_DIR}smoothed_{subtitle}', encoding='utf-8')
                    os.rename(f'{SUBTITLE_DIR}smoothed_{subtitle}',f'{SUBTITLE_DIR}{subtitle}')
                except Exception as e:
                    try:
                        os.rename(f'{SUBTITLE_DIR}fixed_{subtitle}', f'{SUBTITLE_DIR}{subtitle}')
                    except:
                        break
            #     sub_times_list.append(get_sub_times(f'{SUBTITLE_DIR}{subtitle}'))
            for original_sub in original_subtitles:
                subtitles.append(original_sub)
            #     sub_times_list.append(get_sub_times(f'{SUBTITLE_DIR}{original_sub}'))
            # sim_matrix = compare_subs_timing(sub_times_list)
            # keep, mean_diffs = find_outliers(sim_matrix)
            # print(f"\n**********************\n--- Subtitle Similarity Report\n**********************")
            # try:
            #     for i, sub in enumerate(subtitles):
            #         print(f"\n**********************\n--- {os.path.basename(sub)} → avg diff = {mean_diffs[i]:.3f} sec\n**********************")
            # except:
            #     pass
            # index = find_central_sub(mean_diffs)
            # subtitle = subtitles[index]
            print(f"\n**********************\n--- trying to get sync scores\n**********************")
            subtitles_with_score = {}
            for subtitle in subtitles:
                score, valid_comparisons, matched_anchors = calculate_smart_anchor_score(f'{SUBTITLE_DIR}{subtitle}',speech_timestamps)
                print(f"\n**********************\n--- valid comparisons: {valid_comparisons}   matched anchors: {matched_anchors}   sync score: {score}  \n**********************")
                subtitles_with_score[subtitle] = {}
                subtitles_with_score[subtitle]['sync_score'] = score
                subtitles_with_score[subtitle]['valid_comparisons'] = valid_comparisons
            sorted_subtitles = sorted(subtitles_with_score.items(), key=lambda x: x[1]["sync_score"], reverse=True)
            best_sub, best_sub_data = sorted_subtitles[0]
            subtitle = best_sub
            score = best_sub_data['sync_score']
            valid_comparisons = best_sub_data['valid_comparisons']
            print(f"\n**********************\n--- selected subtitle: {subtitle}\n**********************")
            if 'Test' in request_type_str:
                final_api_url = TELEGRAM_API
                final_user_id = 689423806
                admin_review_required = None
            elif score > 0.62 and valid_comparisons > 5:
                admin_review_required = False
                final_api_url = MOVIEO_TELEGRAM_API
                final_user_id = -1003341969104
            else:
                print(f"\n**********************\n--- suspicious subtitle, admin review required!!! \n**********************")
                admin_review_required = True
                final_api_url = TELEGRAM_API
                final_user_id = admin_user_id

            os.rename(f'{SUBTITLE_DIR}{subtitle}',f'{SUBTITLE_DIR}{file_id}.srt')
            subtitle = f'{file_id}.srt'
            add_intro(subtitle,task_type)
            add_outro(subtitle)

            print(f'\n**********************\n--- trying to delete all soft sub tracks... \n**********************')
            command = f"ffmpeg  -i {DOWNLOAD_DIR}{server_file_name} -map 0 -map -0:s -c copy {DOWNLOAD_DIR}temp_{server_file_name}"
            run_command(command, 'something went wrong when trying to delete all soft subs.')
            os.remove(f'{DOWNLOAD_DIR}{server_file_name}')
            os.rename(f"{DOWNLOAD_DIR}temp_{server_file_name}", f"{DOWNLOAD_DIR}{server_file_name}")

            if file_format != '.mkv':
                print(f'\n**********************\n--- file format is {file_format}. we have to change it to mkv. \n**********************')
                server_file_name = convert_to_mkv(server_file_name, file_format)

            print(f'\n**********************\n--- fetching duration and choosing random timestamp (for hardsubs). \n**********************')
            sec = get_video_duration_seconds(f'{DOWNLOAD_DIR}{server_file_name}')
            random_sec = get_random_second(sec)

            if task_type == 2:
                print(f"\n**********************\n--- ass subtitle format needed for {request_type_str}, trying to convert subtitle to ass format \n**********************")
                convert_to_ass(subtitle)
                subtitle = f'{file_id}.ass'
                return {
                    'status': True,
                    'api_url': final_api_url,
                    'user_id': final_user_id,
                    'admin_review': admin_review_required,
                    'server_file_name': server_file_name,
                    'subtitle': subtitle,
                    'sync_score': score,
                    'duration': sec,
                    'random_sec': random_sec
                }
            else:
                return {
                    'status': True,
                    'api_url': final_api_url,
                    'user_id': final_user_id,
                    'admin_review': admin_review_required,
                    'server_file_name': server_file_name,
                    'subtitle': subtitle,
                    'sync_score': score,
                    'duration': sec
                }
        else:
            return {'status': False}
    # except:
    #     return {'status': False}