import json
import logging
import re
import typing
from dataclasses import dataclass

from elevenlabs.core.request_options import RequestOptions
from elevenlabs.music.client import AsyncMusicClient as AutogeneratedAsyncMusicClient
from elevenlabs.music.client import MusicClient as AutogeneratedMusicClient
from elevenlabs.types.allowed_output_formats import AllowedOutputFormats
from elevenlabs.types.music_prompt import MusicPrompt

# this is used as the default value for optional parameters
OMIT = typing.cast(typing.Any, ...)

logger = logging.getLogger(__name__)


def _find_audio_start(data: bytes, start: int) -> int:
    """Find the blank line separating headers from body, returning the index of the body start."""
    crlf_pos = data.find(b"\r\n\r\n", start)
    if crlf_pos != -1:
        return crlf_pos + 4

    lf_pos = data.find(b"\n\n", start)
    if lf_pos != -1:
        return lf_pos + 2

    raise ValueError("Could not find header/body separator in audio part")


@dataclass
class SongMetadata:
    title: str
    description: str
    genres: typing.List[str]
    languages: typing.List[str]
    is_explicit: bool


@dataclass
class MultipartResponse:
    json: typing.Dict[str, typing.Any]  # Contains compositionPlan and songMetadata
    audio: bytes
    filename: str
    song_id: typing.Optional[str] = None


class MusicClient(AutogeneratedMusicClient):
    """
    A client to handle ElevenLabs music-related functionality
    Extends the autogenerated client to include custom music methods
    """

    def compose_detailed(  # type: ignore[override]
        self,
        *,
        output_format: typing.Optional[AllowedOutputFormats] = None,
        prompt: typing.Optional[str] = OMIT,
        composition_plan: typing.Optional[MusicPrompt] = OMIT,
        music_length_ms: typing.Optional[int] = OMIT,
        model_id: typing.Optional[typing.Literal["music_v1"]] = OMIT,
        force_instrumental: typing.Optional[bool] = OMIT,
        store_for_inpainting: typing.Optional[bool] = OMIT,
        with_timestamps: typing.Optional[bool] = OMIT,
        sign_with_c2pa: typing.Optional[bool] = OMIT,
        request_options: typing.Optional[RequestOptions] = None,
        **kwargs: typing.Any,
    ) -> MultipartResponse:
        """
        Compose a song from a prompt or a composition plan with detailed response parsing.
        This method calls the original compose_detailed and then parses the stream response.

        Returns a MultipartResponse containing parsed JSON metadata, audio bytes, filename,
        and song_id (if store_for_inpainting was True).
        """
        # Use the raw client directly to access response headers (for song_id)
        with self._raw_client.compose_detailed(
            output_format=output_format,
            prompt=prompt,
            composition_plan=composition_plan,
            music_length_ms=music_length_ms,
            model_id=model_id,
            force_instrumental=force_instrumental,
            store_for_inpainting=store_for_inpainting,
            with_timestamps=with_timestamps,
            sign_with_c_2_pa=sign_with_c2pa,
            request_options=request_options,
        ) as r:
            result = self._parse_multipart(r.data)
            result.song_id = r.headers.get("song-id")
            return result

    def _parse_multipart(self, stream: typing.Iterator[bytes]) -> MultipartResponse:
        """
        Reads a byte stream containing multipart data and parses it into JSON and audio parts.

        Args:
            stream: Iterator of bytes from ElevenLabs music API response

        Returns:
            MultipartResponse containing parsed JSON metadata, audio bytes, and filename
        """
        # Collect all chunks into a single bytes object
        chunks = []
        for chunk in stream:
            chunks.append(chunk)

        # Combine all chunks into a single buffer
        response_bytes = b''.join(chunks)

        # Parse the multipart content
        response_text = response_bytes.decode('utf-8', errors='ignore')
        lines = response_text.split('\n')

        if not lines:
            raise ValueError("Empty response from music API")

        boundary = lines[0].strip()

        # Find the JSON part (should be early in the response)
        json_data = None
        filename = 'generated_music.mp3'

        # Parse JSON from the text representation
        for i in range(min(10, len(lines))):
            if 'Content-Type: application/json' in lines[i] and i + 2 < len(lines):
                json_line = lines[i + 2]
                if json_line.strip() and json_line.startswith('{'):
                    try:
                        json_data = json.loads(json_line)
                    except json.JSONDecodeError as e:
                        logger.error(f'Failed to parse composition plan JSON: {e}')
                    break

        # Extract filename from headers
        for i in range(min(20, len(lines))):
            if 'filename=' in lines[i]:
                match = re.search(r'filename="([^"]+)"', lines[i])
                if match:
                    filename = match.group(1)
                    break

        # Find where the audio data starts (after the second boundary and headers)
        boundary_bytes = boundary.encode('utf-8')
        first_boundary = -1
        second_boundary = -1

        for i in range(len(response_bytes) - len(boundary_bytes) + 1):
            if response_bytes[i:i + len(boundary_bytes)] == boundary_bytes:
                if first_boundary == -1:
                    first_boundary = i
                elif second_boundary == -1:
                    second_boundary = i
                    break

        if second_boundary == -1:
            raise ValueError('Could not find audio part boundary')

        # Find the start of audio data (after headers and empty line)
        audio_start = _find_audio_start(response_bytes, second_boundary + len(boundary_bytes))

        # Audio goes until the end (or until we find another boundary)
        audio_buffer = response_bytes[audio_start:]

        if not json_data:
            raise ValueError('Could not parse JSON data')

        return MultipartResponse(
            json=json_data,
            audio=audio_buffer,
            filename=filename
        )


class AsyncMusicClient(AutogeneratedAsyncMusicClient):
    """
    An async client to handle ElevenLabs music-related functionality
    Extends the autogenerated async client to include custom music methods
    """

    async def compose_detailed(  # type: ignore[override]
        self,
        *,
        output_format: typing.Optional[AllowedOutputFormats] = None,
        prompt: typing.Optional[str] = OMIT,
        composition_plan: typing.Optional[MusicPrompt] = OMIT,
        music_length_ms: typing.Optional[int] = OMIT,
        model_id: typing.Optional[typing.Literal["music_v1"]] = OMIT,
        force_instrumental: typing.Optional[bool] = OMIT,
        store_for_inpainting: typing.Optional[bool] = OMIT,
        with_timestamps: typing.Optional[bool] = OMIT,
        sign_with_c2pa: typing.Optional[bool] = OMIT,
        request_options: typing.Optional[RequestOptions] = None,
        **kwargs: typing.Any,
    ) -> MultipartResponse:
        """
        Compose a song from a prompt or a composition plan with detailed response parsing.
        This method calls the original compose_detailed and then parses the stream response.

        Returns a MultipartResponse containing parsed JSON metadata, audio bytes, filename,
        and song_id (if store_for_inpainting was True).
        """
        # Use the raw client directly to access response headers (for song_id)
        async with self._raw_client.compose_detailed(
            output_format=output_format,
            prompt=prompt,
            composition_plan=composition_plan,
            music_length_ms=music_length_ms,
            model_id=model_id,
            force_instrumental=force_instrumental,
            store_for_inpainting=store_for_inpainting,
            with_timestamps=with_timestamps,
            sign_with_c_2_pa=sign_with_c2pa,
            request_options=request_options,
        ) as r:
            result = await self._parse_multipart_async(r.data)
            result.song_id = r.headers.get("song-id")
            return result

    async def _parse_multipart_async(self, stream: typing.AsyncIterator[bytes]) -> MultipartResponse:
        """
        Reads an async byte stream containing multipart data and parses it into JSON and audio parts.

        Args:
            stream: AsyncIterator of bytes from ElevenLabs music API response

        Returns:
            MultipartResponse containing parsed JSON metadata, audio bytes, and filename
        """
        # Collect all chunks into a single bytes object
        chunks = []
        async for chunk in stream:
            chunks.append(chunk)

        # Combine all chunks into a single buffer
        response_bytes = b''.join(chunks)

        # Parse the multipart content
        response_text = response_bytes.decode('utf-8', errors='ignore')
        lines = response_text.split('\n')

        if not lines:
            raise ValueError("Empty response from music API")

        boundary = lines[0].strip()

        # Find the JSON part (should be early in the response)
        json_data = None
        filename = 'generated_music.mp3'

        # Parse JSON from the text representation
        for i in range(min(10, len(lines))):
            if 'Content-Type: application/json' in lines[i] and i + 2 < len(lines):
                json_line = lines[i + 2]
                if json_line.strip() and json_line.startswith('{'):
                    try:
                        json_data = json.loads(json_line)
                    except json.JSONDecodeError as e:
                        logger.error(f'Failed to parse JSON: {e}')
                    break

        # Extract filename from headers
        for i in range(min(20, len(lines))):
            if 'filename=' in lines[i]:
                match = re.search(r'filename="([^"]+)"', lines[i])
                if match:
                    filename = match.group(1)
                    break

        # Find where the audio data starts (after the second boundary and headers)
        boundary_bytes = boundary.encode('utf-8')
        first_boundary = -1
        second_boundary = -1

        for i in range(len(response_bytes) - len(boundary_bytes) + 1):
            if response_bytes[i:i + len(boundary_bytes)] == boundary_bytes:
                if first_boundary == -1:
                    first_boundary = i
                elif second_boundary == -1:
                    second_boundary = i
                    break

        if second_boundary == -1:
            raise ValueError('Could not find audio part boundary')

        # Find the start of audio data (after headers and empty line)
        audio_start = _find_audio_start(response_bytes, second_boundary + len(boundary_bytes))

        # Audio goes until the end (or until we find another boundary)
        audio_buffer = response_bytes[audio_start:]

        if not json_data:
            raise ValueError('Could not parse JSON data')

        return MultipartResponse(
            json=json_data,
            audio=audio_buffer,
            filename=filename
        )