import os
import re

from setuptools import find_packages, setup  # type: ignore


def read_file(filename: str) -> str:
    """Read package file as text to get name and version."""
    cwd = os.path.abspath(os.path.dirname(__file__))
    with open(os.path.join(cwd, "suno_utils", filename)) as f:
        return f.read()


def find_version() -> str:
    """Only define version in one place."""
    text = read_file("__init__.py")
    m_version = re.search(r"^__version__ = [\"\']([^\"\']*)[\"\']$", text, re.M)
    if m_version:
        return m_version.group(1)
    raise RuntimeError("Could not find version string.")


def find_name() -> str:
    """Only define name in one place."""
    text = read_file("__init__.py")
    m_name = re.search(r"^__package_name__ = [\"\']([^\"\']*)[\"\']$", text, re.M)
    if m_name:
        return m_name.group(1)
    raise RuntimeError("Could not find name string.")


def find_long_description() -> str:
    """Return the content of the README.md file."""
    return read_file("../README.md")


MUSIC_REQS = [
    "transformers==4.18.0",
    "ddsp==3.3.4",
    "t5==0.9.3",
    "note-seq==0.0.3",
    "pretty-midi==0.2.9",
    "einops==0.4.1",
]
EXTRAS_REQUIRES = {
    "ffmpeg": ["ffmpeg-python"],
    "dev": ["ffmpeg-python"],
    "music": MUSIC_REQS,
}

setup(
    name=find_name(),
    version=find_version(),
    description="Suno audio tools.",
    long_description=find_long_description(),
    long_description_content_type="text/markdown",
    # url="https://github.com/suno-ai/suno_utils",
    author="suno.ai",
    # author_email="code@suno.com",
    # license="Apache 2.0",
    packages=find_packages(),
    install_requires=[
        "beautifulsoup4>=4.10.0",
        "ctc-segmentation>=1.7.4",
        "dagster>0.14",
        "dagster-aws>0.14",
        "descript-audiotools>=0.7.1",
        "editdistance>=0.6.0",
        "encodec>=0.1.1",
        "fasttext",
        "feedparser>=6.0.8",
        "ffmpeg-python>=0.2.0",
        "frozendict",
        "gql[all]>=3.3.0",
        "funcy>=1.17",
        "joblib>=1.2.0",
        "numpy>=1.20.3,<2.0",
        "opensearch-py",
        "pandas>=1.3.4",
        "portion>=2.2.0",
        "pygtrie>=2.1,<3.0",
        "python-speech-features>=0.6",
        "requests>=2.26.0",
        "scipy>=1.7.3",
        "sox>=1.4.1",
        "torch>=2.5.0,<3",
        "tqdm>=4.62.3",
        "tinytag>=1.8.1",
        "transformers",
        "unidecode>=1.3.2",
        "phonemizer>=3.2.1",
        "sentencepiece>=0.1.97",
        "tiktoken==0.1.2",
        "openai-whisper",
        "rpyc>=6.0.0",
    ],
    extras_require=EXTRAS_REQUIRES,
    package_data={
        "": [
            "tests/sample_data/cats.opus",
            "tests/sample_data/cats.txt",
            "tests/sample_data/cats_gentle.json",
            "tasks/nano/tag_map.json",
        ]
    },
    dependency_links=[],
)