#  Copyright (C) 2025. Hao Zheng
#  All rights reserved.

from __future__ import annotations

import re
import subprocess
import time
import unicodedata
from pathlib import Path
from typing import TYPE_CHECKING, Any

if TYPE_CHECKING:
    from spacy.language import Language as SpacyLanguage

import audioread
import ffmpeg
import filetype
import jaconvV2
import langcodes
import spacy
import spacy.cli
import tiktoken
import torch
from lingua import LanguageDetectorBuilder

from openlrc.defaults import supported_languages_lingua
from openlrc.logger import logger


def extract_audio(path: Path) -> Path:
    """
    Extract audio from video.
    :return: Audio path
    """
    file_type = get_file_type(path)
    if file_type == "audio":
        return path

    probe = ffmpeg.probe(path)
    audio_streams = next((stream for stream in probe["streams"] if stream["codec_type"] == "audio"), None)
    if audio_streams is None:
        raise RuntimeError(f"No audio stream found in {path}")
    sample_rate = audio_streams["sample_rate"]
    logger.info(f"File {path}: Audio sample rate: {sample_rate}")

    audio, err = (
        ffmpeg.input(path)
        .output("pipe:", format="wav", acodec="pcm_s16le", ar=sample_rate, loglevel="quiet")
        .run(capture_stdout=True)
    )

    if err:
        raise RuntimeError(f"ffmpeg error: {err}")

    audio_path = path.with_suffix(".wav")
    with open(audio_path, "wb") as f:
        f.write(audio)

    return audio_path


def get_preprocessed_path(audio_path: str | Path) -> Path:
    """
    Get the expected preprocessed audio file path for a given input.

    This function returns the path where the preprocessed audio file would be
    stored after preprocessing. It is useful when you need to check if
    preprocessing has already been done, or when running preprocessing and
    transcription in separate stages.

    Args:
        audio_path (Union[str, Path]): Original audio file path.

    Returns:
        Path: Expected path of the preprocessed audio file.

    Example:
        >>> get_preprocessed_path('/data/audio.mp3')
        PosixPath('/data/preprocessed/audio_preprocessed.wav')
    """
    audio_path = Path(audio_path)
    return audio_path.parent / "preprocessed" / f"{audio_path.stem}_preprocessed.wav"


def get_file_type(path: Path) -> str:
    if path.suffix == ".ts":
        return "video"

    try:
        guess = filetype.guess(path)
        if guess is None:
            raise RuntimeError(f"File {path} is not a valid file.")
        file_type = guess.mime.split("/")[0]
    except (TypeError, AttributeError) as e:
        raise RuntimeError(f"File {path} is not a valid file.") from e

    if file_type not in ["audio", "video"]:
        raise RuntimeError(f"File {path} is not a valid file. Should be audio or video file.")

    return file_type


def get_audio_duration(path: str | Path) -> float:
    with audioread.audio_open(str(path)) as audio:
        return audio.duration


def release_memory(model: torch.nn.Module) -> None:
    # gc.collect()
    torch.cuda.empty_cache()
    del model


def normalize(text):
    """
    Normalize strings using str.lower(), and unicodedata.normalize
    """
    # unicodedata can't handle quotes as expected'’'
    quotes_table = str.maketrans("〈〉゛“”‘’", '<>"""\'\'')
    text = text.translate(quotes_table)

    text = unicodedata.normalize("NFKC", text.lower())

    # unicodedata can't handle kana
    text = jaconvV2.z2h(text)

    # Special case
    special_table = str.maketrans("〇①②③④⑤⑥⑦⑧⑨", "0123456789")
    text = text.translate(special_table)

    return text


def get_text_token_number(text: str, model: str = "gpt-3.5-turbo") -> int:
    tokens = tiktoken.encoding_for_model(model).encode(text)

    return len(tokens)


def get_messages_token_number(messages: list[dict[str, Any]], model: str = "gpt-3.5-turbo") -> int:
    total = sum([get_text_token_number(element["content"], model=model) for element in messages])

    return total


def extend_filename(filename: Path, extend: str) -> Path:
    """Extend a filename with some string."""
    return filename.with_stem(filename.stem + extend)


class Timer:
    def __init__(self, task=""):
        self._start: float | None = None
        self._stop: float | None = None
        self.task = task

    def start(self):
        if self.task:
            logger.info(f"Start {self.task}")
        self._start = time.perf_counter()

    def stop(self):
        self._stop = time.perf_counter()
        logger.info(f"{self.task} Elapsed: {self._elapsed:.2f}s")

    @property
    def _elapsed(self) -> float:
        assert self._start is not None and self._stop is not None, "Timer not started/stopped"
        return self._stop - self._start

    @property
    def duration(self) -> float:
        assert self._start is not None, "Timer not started"
        return time.perf_counter() - self._start

    def __enter__(self):
        self.start()
        return self

    def __exit__(self, *args):
        self.stop()


def parse_timestamp(time_stamp: str, fmt: str) -> float:
    """
    Parse a timestamp from a subtitle file and convert it to seconds.

    Args:
        time_stamp (str): The timestamp to parse.
        fmt (str): The format of `time_stamp`. Supported values are:
            - 'lrc' for LRC format, e.g., '1:23.45'
            - 'srt' for SRT format, e.g., '01:23:45,678'

    Returns:
        float: The timestamp in seconds.

    Raises:
        ValueError: If `time_stamp` does not match the expected format for the specified `fmt`.
    """

    if fmt == "lrc":
        if not re.match(r"^\d+:\d+\.\d+$", time_stamp):
            raise ValueError(f"Invalid timestamp format for LRC: {time_stamp}")
        minutes, seconds = time_stamp.split(":")
        seconds, hundredths_of_sec = seconds.split(".")
        return int(minutes) * 60 + int(seconds) + int(hundredths_of_sec) / 100.0
    elif fmt == "srt":
        if not re.match(r"^\d+:\d+:\d+,\d+$", time_stamp):
            raise ValueError(f"Invalid timestamp format for SRT: {time_stamp}")
        hours, minutes, seconds = time_stamp.split(":")
        seconds, milliseconds = seconds.split(",")
        return int(hours) * 3600 + int(minutes) * 60 + int(seconds) + int(milliseconds) / 1000.0
    else:
        raise ValueError(f"Unsupported timestamp format: {fmt}")


def format_timestamp(seconds: float, fmt: str = "lrc") -> str:
    """
    Converts a timestamp in seconds into a string in the specified format.

    Args:
        seconds (float): Timestamp in seconds.
        fmt (str): Format of the output string. Supported values are:
            - 'lrc' for LRC format, e.g., '1:23.45'
            - 'srt' for SRT format, e.g., '01:23:45,678'

    Returns:
        str: A string representation of the timestamp in the specified format.
    """
    # Ensure that the timestamp is non-negative.
    # assert seconds >= 0, "non-negative timestamp expected"
    if seconds < 0:
        logger.warning(f"Negative timestamp: {seconds}")
        if fmt == "lrc":
            return "0:00.00"
        elif fmt == "srt":
            return "00:00:00,000"
        else:
            raise ValueError(f"Unsupported timestamp format: {fmt}")

    # Convert seconds into milliseconds.
    milliseconds = round(seconds * 1000.0)

    # Extract hours, minutes, seconds, and milliseconds from milliseconds.
    hours = milliseconds // 3600000
    milliseconds %= 3600000
    minutes = milliseconds // 60000
    milliseconds %= 60000
    seconds = milliseconds // 1000
    milliseconds %= 1000

    # Return the timestamp in the specified format.
    if fmt == "lrc":
        return f"{minutes:02d}:{seconds:02d}.{milliseconds // 10:02d}"
    elif fmt == "srt":
        return f"{hours:02d}:{minutes:02d}:{seconds:02d},{milliseconds:03d}"
    else:
        raise ValueError(f"Unsupported timestamp format: {fmt}")


def detect_lang(text) -> str:
    detector = LanguageDetectorBuilder.from_languages(*supported_languages_lingua).build()
    detected = detector.detect_language_of(text)
    if detected is None:
        raise RuntimeError(f"Could not detect language for text: {text[:50]!r}")
    name = detected.name.lower()
    lang_code = langcodes.Language.find(name).language
    if lang_code is None:
        raise RuntimeError(f"Could not resolve language code for: {name!r}")
    return lang_code


def get_spacy_lib(lang):
    special_case = {"core_web": ["zh", "en"], "ent_wiki": ["xx"]}

    mid_str = "core_news"
    for k, v in special_case.items():
        if lang in v:
            mid_str = k

    return f"{lang}_{mid_str}_sm"


def spacy_load(lang) -> SpacyLanguage:
    lib_name = get_spacy_lib(lang)
    try:
        nlp = spacy.load(lib_name)
    except (ImportError, OSError):
        logger.warning(f"Spacy model {lib_name} missed, downloading")
        spacy.cli.download(lib_name)  # pyright: ignore[reportPrivateImportUsage]
        nlp = spacy.load(lib_name)

    return nlp


def get_similarity(text1, text2):
    lang1 = detect_lang(text1)
    lang2 = detect_lang(text2)

    if lang1 != lang2:
        raise ValueError(f'language of "{text1}" ({lang1}) is not the same as "{text2}" ({lang2})')

    nlp = spacy_load(lang1)

    doc1 = nlp(text1)
    doc2 = nlp(text2)

    return doc1.similarity(doc2)


def merge_subtitle(video_path, subtitle_path, output_path):
    # check ffmpeg
    try:
        subprocess.check_output(["ffmpeg", "-version"])
    except FileNotFoundError:
        raise RuntimeError("ffmpeg is not installed. Please install ffmpeg first.") from None

    subtitle_abs = str(Path(subtitle_path).absolute())
    style = "FontSize=24,Bold=1"
    subprocess.call(
        [
            "ffmpeg",
            "-y",
            "-i",
            str(video_path),
            "-vf",
            f"subtitles={subtitle_abs}:force_style='{style}'",
            str(output_path),
        ]
    )

    logger.info(f"Subtitled video saved to {output_path}")


def remove_stop(text, stop_sequences):
    """
    Remove stop sequences from the text.
    """
    if not text or not stop_sequences:
        return text

    for stop_seq in stop_sequences:
        text = text.rstrip(stop_seq)

    return text.strip()
