import re
import tiktoken
from unittest.mock import Mock


import numpy as np
import pytest

from suno_utils.worker.generate_song_lyrics import (
    ANODYNE_GENRES,
    does_text_contain_slur,
    get_prompt_from_gpt_description_prompt,
    ModerationError,
    _are_lyrics_generically_malformed,
    _extract_title_from_prompt,
    _get_full_gpt_description_prompt,
    _truncate_lyrics,
    _normalize_user_prompt,
    _parse_gpt_output,
    get_stanzas,
    moderate_prompt,
    moderate_gpt_description_prompt,
    _maybe_decommify,
    LyricsLength,
)
from suno_utils.harvest.youtube.tests.test_language_classify import (
    make_mock_fasttext_client,
)
from suno_utils.worker.generate_song_lyrics import GPT_LYRICS_GENERATION_FAILED


class MockOpenAIDelta:
    """Represents the 'delta' portion of a streaming chunk."""

    def __init__(self, content=None):
        self.content = content  # e.g. "some partial text"


class MockOpenAIChoice:
    """Represents each choice in a streaming chunk."""

    def __init__(self, index, content=None, finish_reason=None):
        self.index = index
        # The real OpenAI Python library returns chunk.choices[i].delta.content
        # for the partial text. So we replicate that structure here:
        self.delta = MockOpenAIDelta(content)
        self.finish_reason = finish_reason


class MockOpenAIStreamChunk:
    """Replicates the chunk object in a streaming ChatCompletion response."""

    def __init__(self, choices):
        # The real chunk object has a `.choices` list attribute.
        self.choices = choices


def _stream_response_as_objects(content: str, model: str, choice_idx: int = 0, chunk_size: int = 50):
    """
    Tokenize the `content` and yield chunked responses as objects
    that look like the streaming responses from openai.ChatCompletion.
    """
    # If model is unknown, fall back to a known encoding:
    enc = tiktoken.get_encoding("cl100k_base")
    tokens = enc.encode(content)

    # Yield partial text in chunk_size increments:
    for i in range(0, len(tokens), chunk_size):
        chunk_tokens = tokens[i : i + chunk_size]
        chunk_str = enc.decode(chunk_tokens)
        # Each partial chunk is a MockOpenAIStreamChunk with .choices
        yield MockOpenAIStreamChunk(
            [MockOpenAIChoice(index=choice_idx, content=chunk_str, finish_reason=None)]
        )

    # Finally, yield a last chunk with finish_reason="stop" (or None) to signal the end.
    yield MockOpenAIStreamChunk([MockOpenAIChoice(index=choice_idx, content=None, finish_reason="stop")])


def _make_mock_lyrics_client(
    mock_raw_lyrics: str,
    flag_prompt_moderation: bool = False,
    flag_lyrics_moderation: bool = False,
    flag_copyright: bool = False,
) -> Mock:
    """
    Make a mock OpenAI client that supports both normal and streamed completions
    with chunk.choices as an attribute (rather than dict keys).
    """
    mock_lyrics_client = Mock()

    def mock_create(
        model,
        system_prompt,
        user_prompt,
        max_tokens=None,
        temperature=1.0,
        n=1,
        stream=False,
        **kwargs,
    ):
        """Spoof the main ChatGPT entrypoint with optional streaming."""

        # Decide what text to return
        if "CopyrightDetectorGPT" in system_prompt:
            content = "COPYRIGHTED" if flag_copyright else "ORIGINAL"
        else:
            content = mock_raw_lyrics

        # Non-streamed response: return the old single-completion style
        if not stream:
            mock_completion = Mock()
            # mimic "completion.choices" as a list
            choices = []
            for _ in range(n):
                choice_mock = Mock()
                # This is how we typically access the final text
                choice_mock.message.content = content
                choices.append(choice_mock)
            mock_completion.choices = choices
            return mock_completion

        # If streaming, yield objects chunk by chunk
        def _chunked_stream():
            for choice_idx in range(n):
                # For each completion, yield chunked text
                yield from _stream_response_as_objects(content, model, choice_idx, chunk_size=50)

        # Return an iterator/generator of these chunk-objects
        return _chunked_stream()

    # mock_lyrics_client.chat.completions.create.return_value = mock_completion
    mock_lyrics_client.create = mock_create

    mock_moderation_response = Mock()

    mock_prompt_moderation = Mock()
    mock_prompt_moderation.flagged = flag_prompt_moderation
    mock_prompt_moderation.category_scores = {"naughty": 0.96 * flag_prompt_moderation}.items()

    mock_lyrics_moderation = Mock()
    mock_lyrics_moderation.flagged = flag_lyrics_moderation
    mock_lyrics_moderation.category_scores = {"naughty": 0.96 * flag_lyrics_moderation}.items()

    mock_moderation_response.results = [
        mock_prompt_moderation,
        mock_lyrics_moderation,
    ]
    mock_lyrics_client.openai_client.moderations.create.return_value = mock_moderation_response
    # response = openai_client.moderations.create(input=texts)
    # moderations = response.results
    return mock_lyrics_client


MOCK_RAW_LYRICS = """{my title}
{tag1, tag2, tag3}
Here are some lyrics
They're not very good,
They'll brighten your spirits
If read them you would
They're over 100 chars long!"""

MOCK_RAW_LYRICS_NO_TAGS = """Here are some lyrics
They're not very good,
They'll brighten your spirits
If read them you would
They're over 100 chars long!"""

EXPECTED_LYRICS = """Here are some lyrics
They're not very good
They'll brighten your spirits
If read them you would
They're over 100 chars long!"""


def test_generate_song_with_genre_tags_happy_path():
    mock_openai_client = _make_mock_lyrics_client(MOCK_RAW_LYRICS)
    mock_fasttext_client = make_mock_fasttext_client()
    title, lyrics, genre_tags, language = get_prompt_from_gpt_description_prompt(
        mock_openai_client, mock_fasttext_client, gpt_description_prompt="a generic user prompt"
    )
    assert EXPECTED_LYRICS == lyrics
    assert {"tag1", "tag2", "tag3"} == set(genre_tags)


def test_generate_song_with_genre_tags_no_tags():
    mock_openai_client = _make_mock_lyrics_client(MOCK_RAW_LYRICS_NO_TAGS)
    mock_fasttext_client = make_mock_fasttext_client()
    title, lyrics, genre_tags, language = get_prompt_from_gpt_description_prompt(
        mock_openai_client, mock_fasttext_client, gpt_description_prompt="a generic user prompt"
    )
    assert EXPECTED_LYRICS == lyrics
    assert all(genre_tag in ANODYNE_GENRES for genre_tag in genre_tags)


def test_generate_song_with_genre_tags_short_lyrics():
    mock_openai_client = _make_mock_lyrics_client("I'm a generic malformed response!")
    mock_fasttext_client = make_mock_fasttext_client()
    with pytest.raises(
        ModerationError,
        match=GPT_LYRICS_GENERATION_FAILED,
    ):
        _ = get_prompt_from_gpt_description_prompt(
            mock_openai_client,
            mock_fasttext_client,
            gpt_description_prompt="a generic user prompt",
        )


def test_generate_song_with_genre_tags_gpt_balk():
    mock_openai_client = _make_mock_lyrics_client("I'm sorry, but I can't assist with that request")
    mock_fasttext_client = make_mock_fasttext_client()

    with pytest.raises(ModerationError, match=GPT_LYRICS_GENERATION_FAILED):
        _ = get_prompt_from_gpt_description_prompt(
            mock_openai_client,
            mock_fasttext_client,
            gpt_description_prompt="a prompt triggering a GPT balk",
        )


def test_generate_song_with_genre_tags_flagged_prompt():
    mock_openai_client = _make_mock_lyrics_client(
        "I'm sorry, but I can't assist with that request", flag_prompt_moderation=True
    )
    mock_fasttext_client = make_mock_fasttext_client()
    expected_err_msg = re.escape("flagged for moderation")
    with pytest.raises(ModerationError, match=expected_err_msg):
        _ = get_prompt_from_gpt_description_prompt(
            mock_openai_client,
            mock_fasttext_client,
            gpt_description_prompt="a prompt that triggers prompt moderation",
        )


def test_generate_song_with_genre_tags_flagged_lyrics():
    mock_lyrics = "Some flaggable lyrics but they're really long\n" * 10
    mock_openai_client = _make_mock_lyrics_client(mock_lyrics, flag_lyrics_moderation=True)
    mock_fasttext_client = make_mock_fasttext_client()
    expected_err_msg = "flagged for moderation"
    with pytest.raises(ModerationError, match=expected_err_msg):
        _ = get_prompt_from_gpt_description_prompt(
            mock_openai_client,
            mock_fasttext_client,
            gpt_description_prompt="an innocent or ambiguous prompt causing flaggable lyrics",
        )


def test_get_prompt_from_gpt_description_prompt_one_box_instrumental():
    mock_lyrics = "Some normal lyrics that we'll throw away because instrumental\n" * 10
    mock_openai_client = _make_mock_lyrics_client(mock_lyrics, flag_copyright=False)
    mock_fasttext_client = make_mock_fasttext_client()
    title, prompt, genre_tags, language = get_prompt_from_gpt_description_prompt(
        mock_openai_client,
        mock_fasttext_client,
        gpt_description_prompt="an instrumental funk song",
        make_instrumental=True,
    )
    assert prompt == "[Instrumental]"


def test__normalize_user_prompt():
    assert "foo" == _normalize_user_prompt("foo")
    assert "foo" == _normalize_user_prompt("foo \t[[[]]]")
    assert "수줍은 새" == _normalize_user_prompt("수줍은 새")


def test__parse_gpt_output():
    raw_lyrics = """{my title}
    {catchy, upbeat, inspiring}
[Verse]
In Istanbul, the city so grand
The skyline shines, like golden sand (oooh)
The Bosphorus flows, so serene and clear
A melting pot of culture, everything you'll hear (oh-yeah)
From the historic mosques to the bustling bazaars

[Chorus]
Turkey, oh Turkey, you stole our hearts away (stole our hearts away)
With your vibrant streets and your warm embrace
Oh Turkey, oh Turkey, we'll never be the same (never be the same)
Forever grateful for the memories we made (memories we made)
---

Q: Write a song about the power of love to overcome obstacles.
A:"""
    gpt_description_prompt = "give me a k-pop song"
    title, parsed_lyrics, genre_tags = _parse_gpt_output(raw_lyrics, gpt_description_prompt, "k-pop")
    assert parsed_lyrics.startswith("[Verse]")
    assert "Q:" not in parsed_lyrics
    assert "A:" not in parsed_lyrics
    assert "---" not in parsed_lyrics
    assert {"inspiring", "k-pop"} == set(genre_tags)  # make sure we block 'catchy', include 'k-pop'
    assert parsed_lyrics.strip().endswith("(memories we made)")


def test__parse_gpt_output_blocked_genres():
    raw_lyrics = """{catchy, energetic, festive}
[Verse]
Blah blah blah blah blah
Blah blah blah blah blah
Blah blah blah blah blah    
Blah blah blah blah blah
    
[Chorus]
Blah blah blah blah blah
Blah blah blah blah blah
Blah blah blah blah blah    
Blah blah blah blah blah
"""
    gpt_description_prompt = "a song about blah"
    title, parsed_lyrics, genre_tags = _parse_gpt_output(raw_lyrics, gpt_description_prompt, "pop")
    assert all(genre_tag in ANODYNE_GENRES for genre_tag in genre_tags)


def test__parse_gpt_output_num_stanzas_eq_3():
    raw_lyrics = """{catchy, energetic, festive}
[Verse]
foo foo foo foo foo
foo foo foo foo foo
foo foo foo foo foo    
foo foo foo foo foo
    
[Chorus]
bar bar bar bar bar
bar bar bar bar bar
bar bar bar bar bar    
bar bar bar bar bar

[Verse]
baz baz baz baz baz
baz baz baz baz baz
baz baz baz baz baz    
baz baz baz baz baz
"""
    gpt_description_prompt = "a song about blah"
    title, parsed_lyrics, genre_tags = _parse_gpt_output(
        raw_lyrics, gpt_description_prompt, "hip hop", lyrics_length=LyricsLength.STANDARD
    )
    assert len(get_stanzas(parsed_lyrics)) == 3


def test__parse_gpt_output_num_stanzas_eq_2():
    raw_lyrics = """{catchy, energetic, festive}
[Verse]
foo foo foo foo foo
foo foo foo foo foo
foo foo foo foo foo    
foo foo foo foo foo
    
[Chorus]
bar bar bar bar bar
bar bar bar bar bar
bar bar bar bar bar    
bar bar bar bar bar

[Verse]
baz baz baz baz baz
baz baz baz baz baz
baz baz baz baz baz    
baz baz baz baz baz
"""
    gpt_description_prompt = "a song about blah"
    title, parsed_lyrics, genre_tags = _parse_gpt_output(
        raw_lyrics, gpt_description_prompt, "j-pop", lyrics_length=LyricsLength.SHORT
    )
    stanzas = get_stanzas(parsed_lyrics)
    print(stanzas)
    assert len(stanzas) == 2


def test__parse_gpt_output_extract_genre_tag_from_prompt():
    gpt_description_prompt = "a rock song about blah"
    title, parsed_lyrics, genre_tags = _parse_gpt_output(
        "", gpt_description_prompt, "j-pop", lyrics_length=LyricsLength.SHORT
    )
    assert genre_tags == ["j-pop", "rock"]


def test__parse_gpt_output_extract_genre_tag_from_prompt():
    gpt_description_prompt = "Grandiose instrumental French opera piece with two harpsichords, suitable for medieval-themed movies or TV shows."
    title, parsed_lyrics, genre_tags = _parse_gpt_output(
        "", gpt_description_prompt, "", lyrics_length=LyricsLength.SHORT
    )
    assert set(genre_tags) == {"medieval", "opera", "grandiose", "harpsichords"}


def test__parse_gpt_output_repeated_tags():
    gpt_description_prompt = "Grandiose instrumental French opera piece with two harpsichords, suitable for medieval-themed movies or TV shows."
    title, parsed_lyrics, genre_tags = _parse_gpt_output(
        "", gpt_description_prompt, "opera", lyrics_length=LyricsLength.SHORT
    )
    assert set(genre_tags) == {"medieval", "opera", "grandiose", "harpsichords"}


def test__are_lyrics_generically_malformed():
    raw_lyrics = """energetic, catchy, pop [Verse]
foo foo foo
bar bar bar
is all I want to say to you
    
[Chorus]
blah blah blah
quantity beats quality
quantity beats quality
quantity beats quality
"""

    assert _are_lyrics_generically_malformed(raw_lyrics)


def test__are_lyrics_generically_malformed_with_newline():
    raw_lyrics = """energetic, catchy, pop
[Verse]
foo foo foo
bar bar bar
is all I want to say to you
    
[Chorus]
blah blah blah
quantity beats quality
quantity beats quality
quantity beats quality
"""

    assert _are_lyrics_generically_malformed(raw_lyrics)


def test__are_lyrics_generically_malformed_with_initial_newline():
    raw_lyrics = """


energetic, catchy, pop
[Verse]
foo foo foo
bar bar bar
is all I want to say to you
    
[Chorus]
blah blah blah
quantity beats quality
quantity beats quality
quantity beats quality
"""
    assert _are_lyrics_generically_malformed(raw_lyrics)


def test__are_lyrics_generically_malformed_false_positive():
    raw_lyrics = """[Verse]
foo foo foo
bar bar bar
DON'T REMOVE THIS IT'S A VALID LYRIC [Verse2]

these lyrics aren't great
but we shouldn't cut them out
quantity beats quality
quantity beats quality
quantity beats quality

"""
    assert not _are_lyrics_generically_malformed(raw_lyrics)


def test__are_lyrics_generically_malformed_happy_path():
    raw_lyrics = """[Verse]
foo foo foo
bar bar bar

[Chorus]
these lyrics aren't great
but quantity beats quality
but quantity beats quality
but quantity beats quality
"""
    assert not _are_lyrics_generically_malformed(raw_lyrics)


def test__are_lyrics_generically_malformed_chinese():
    raw_lyrics = """[Verse]
在下雨的日子里
我追逐著愛情的蹤跡
滂沱的雨水打在臉上
我只能微笑著面對

[Chorus]
找尋愛情的腳步
在雨中顯得如此遙遠
但我願意淋成濕透
只為找到愛的彼岸"""
    assert not _are_lyrics_generically_malformed(raw_lyrics)


@pytest.mark.parametrize(
    "fasttext_return_value,expected",
    [
        ((("__label__en",), np.array([0.9])), ""),  # english we don't add this lyrics again
        ((("__label__es",), np.array([0.9])), "in Spanish"),
        ((("__label__nl",), np.array([0.9])), "in Dutch"),
        ((("__label__ru",), np.array([0.9])), "in Russian"),
        ((("__label__en",), np.array([0.1])), ""),
    ],
)
def test__get_full_user_prompt_english(fasttext_return_value, expected):
    mock_fasttext_client = make_mock_fasttext_client(fasttext_return_value)
    user_prompt, _ = _get_full_gpt_description_prompt(mock_fasttext_client, "hello how are you")
    assert expected in user_prompt


def test__extract_title_from_prompt_one_title():
    prompt = """

{A Good Title}
[Verse]
foo foo foo
bar bar bar

[Chorus]
these lyrics aren't great
but quantity beats quality
but quantity beats quality
but quantity beats quality
"""
    title = _extract_title_from_prompt(prompt)
    assert "A Good Title" == title


def test__extract_title_from_prompt_no_titles():
    prompt = """
[Verse]
foo foo foo
bar bar bar

[Chorus]
these lyrics aren't great
but quantity beats quality
but quantity beats quality
but quantity beats quality
"""
    title = _extract_title_from_prompt(prompt)
    assert "" == title


def test__extract_title_from_prompt_many_titles():
    prompt = """
{here a title}
[Verse]
foo foo foo
bar bar bar
title: there a title
[Chorus]
these lyrics aren't great
but quantity beats quality  title: everywhere a title title
but quantity beats quality
but quantity beats quality
"""
    title = _extract_title_from_prompt(prompt)
    assert "here a title" == title


def test_moderate_prompt():
    mock_openai_client = _make_mock_lyrics_client("")
    moderate_prompt(mock_openai_client, "a prompt")


def test_moderate_gpt_description_prompt():
    mock_lyrics_client = _make_mock_lyrics_client("")
    moderate_gpt_description_prompt(mock_lyrics_client.openai_client, "a gpt description prompt")


@pytest.mark.parametrize(
    "text,expected",
    [("coon", True), ("dink", True), ("dinky", False), ("racoon", False), ("maine coon", False)],
)
def test_does_text_contain_slur(text, expected):
    assert does_text_contain_slur(text) == expected


def test__maybe_decommify():
    lyrics = """[Verse]
hello, world (with, commas),
and more, commas

[Chorus]
foo, bar (bar, bar)
baz, quux (ooh-yeah)
"""

    expected = """[Verse]
Hello
World (with, commas)
And more
Commas

[Chorus]
Foo
Bar (bar, bar)
Baz
Quux (ooh-yeah)"""
    actual = _maybe_decommify(lyrics, genre="pop")

    assert expected == actual


# @pytest.fixture
# def lyrics_length():
#     return NamedTuple("LyricsLength", [("num_stanzas", int), ("lines_per_stanza", int)])


def test_basic_truncation():
    prompt = "Line 1\nLine 2\nLine 3\n\nLine 4\nLine 5\nLine 6"
    expected = "Line 1\nLine 2\n\nLine 4\nLine 5"
    assert _truncate_lyrics(prompt, LyricsLength.SHORT) == expected


def test_control_tag_handling():
    prompt = "[VERSE]\nLine 1\nLine 2\nLine 3\n\n[CHORUS]\nLine 4\nLine 5\nLine 6"
    expected = "[VERSE]\nLine 1\nLine 2\n\n[CHORUS]\nLine 4\nLine 5"
    assert _truncate_lyrics(prompt, LyricsLength.SHORT) == expected


def test_fewer_stanzas_than_requested():
    prompt = "Line 1\nLine 2\n\nLine 3\nLine 4"
    expected = "Line 1\nLine 2\n\nLine 3\nLine 4"
    assert _truncate_lyrics(prompt, LyricsLength.STANDARD) == expected


def test_fewer_lines_than_requested():
    prompt = "Line 1\n\nLine 2\nLine 3"
    expected = "Line 1\n\nLine 2\nLine 3"
    assert _truncate_lyrics(prompt, LyricsLength.STANDARD) == expected


def test_empty_prompt():
    prompt = ""
    expected = ""
    assert _truncate_lyrics(prompt, LyricsLength.SHORT) == expected


def test_single_line_stanzas():
    prompt = "Line 1\n\nLine 2\n\nLine 3\n\nLine 4\n\nLine 5\n\nLine 6"
    expected = "Line 1\n\nLine 2\n\nLine 3\n\nLine 4\n\nLine 5\n\nLine 6"
    assert _truncate_lyrics(prompt, LyricsLength.LONG) == expected


def test_mixed_control_tags():
    prompt = "[VERSE]\nLine 1\nLine 2\n\nLine 3\nLine 4\n\n[CHORUS]\nLine 5\nLine 6"
    expected = "[VERSE]\nLine 1\nLine 2\n\nLine 3\nLine 4\n\n[CHORUS]\nLine 5\nLine 6"
    assert _truncate_lyrics(prompt, LyricsLength.STANDARD) == expected


@pytest.mark.parametrize(
    "prompt,length,expected",
    [
        ("Line 1\nLine 2", LyricsLength.SHORT, "Line 1\nLine 2"),
        ("Line 1\nLine 2\n\nLine 3\nLine 4", LyricsLength.SHORT, "Line 1\nLine 2\n\nLine 3\nLine 4"),
        ("[TAG]\nLine 1\nLine 2", LyricsLength.SHORT, "[TAG]\nLine 1\nLine 2"),
        (
            "Line 1\nLine 2\nLine 3\nLine 4\n\nLine 5\nLine 6\nLine 7\nLine 8",
            LyricsLength.STANDARD,
            "Line 1\nLine 2\nLine 3\nLine 4\n\nLine 5\nLine 6\nLine 7\nLine 8",
        ),
    ],
)
def test_parametrized_cases(prompt, length, expected):
    assert _truncate_lyrics(prompt, length) == expected