import io
import logging
import random
import re

import boto3

_logger = logging.getLogger(__name__)

POLLY_VOICES = {
    "Danielle": {"gender": "Female", "accent": "US"},
    "Gregory": {"gender": "Male", "accent": "US"},
    "Kevin": {"gender": "Male", "accent": "US"},
    "Ruth": {"gender": "Female", "accent": "US"},
    "Stephen": {"gender": "Male", "accent": "US"},
    "Arthur": {"gender": "Male", "accent": "GB"},
    "Olivia": {"gender": "Female", "accent": "GB"},
    "Ayanda": {"gender": "Female", "accent": "ZA"},
    "Aria": {"gender": "Female", "accent": "NZ"},
    "Kajal": {"gender": "Female", "accent": "IN"},
    "Niamh": {"gender": "Female", "accent": "IE"},
}

FINAL_CUE_TEXT = (
    "This audio recording, for the listening exercise, has finished."
)
FINAL_CUE_VOICE = "Stephen"

MAX_CHUNK_SIZE = 3000


class EncoachPollyService:

    def __init__(self, env):
        self.env = env
        icp = env["ir.config_parameter"].sudo()
        self.client = boto3.client(
            "polly",
            aws_access_key_id=icp.get_param("encoach.aws_access_key_id", ""),
            aws_secret_access_key=icp.get_param("encoach.aws_secret_access_key", ""),
            region_name=icp.get_param("encoach.aws_region", "us-east-1"),
        )

    def synthesize_speech(
        self, text, voice, engine="neural", output_format="mp3"
    ):
        """Synthesize a single text block into MP3 bytes."""
        chunks = self._chunk_text(text)
        audio_parts = []
        for chunk in chunks:
            resp = self.client.synthesize_speech(
                Text=chunk,
                VoiceId=voice,
                Engine=engine,
                OutputFormat=output_format,
            )
            audio_parts.append(resp["AudioStream"].read())
        return b"".join(audio_parts)

    def text_to_speech(self, dialog, include_final_cue=True):
        """Generate full MP3 audio from a conversation or monologue.

        dialog: either a string (monologue) or a list of dicts with
            'name', 'gender', 'text' keys.
        """
        audio_segments = []

        if isinstance(dialog, str):
            voice = random.choice(list(POLLY_VOICES.keys()))
            audio_segments.append(self.synthesize_speech(dialog, voice))
        else:
            voice_map = self._assign_voices(dialog)
            for line in dialog:
                voice = voice_map.get(line["name"], "Stephen")
                audio_segments.append(
                    self.synthesize_speech(line["text"], voice)
                )

        if include_final_cue:
            audio_segments.append(
                self.synthesize_speech(FINAL_CUE_TEXT, FINAL_CUE_VOICE)
            )

        return b"".join(audio_segments)

    @staticmethod
    def _assign_voices(dialog):
        """Assign distinct Polly voices to each speaker based on gender."""
        speakers = {}
        male_voices = [v for v, info in POLLY_VOICES.items() if info["gender"] == "Male"]
        female_voices = [v for v, info in POLLY_VOICES.items() if info["gender"] == "Female"]
        male_idx, female_idx = 0, 0

        for line in dialog:
            name = line.get("name", "")
            if name in speakers:
                continue
            gender = (line.get("gender") or "").lower()
            if gender == "female" and female_idx < len(female_voices):
                speakers[name] = female_voices[female_idx]
                female_idx += 1
            elif male_idx < len(male_voices):
                speakers[name] = male_voices[male_idx]
                male_idx += 1
            else:
                speakers[name] = random.choice(list(POLLY_VOICES.keys()))
        return speakers

    @staticmethod
    def _chunk_text(text, max_size=MAX_CHUNK_SIZE):
        """Split text at sentence boundaries respecting max chunk size."""
        if len(text) <= max_size:
            return [text]

        sentences = re.split(r"(?<=[.!?])\s+", text)
        chunks = []
        current = ""

        for sentence in sentences:
            if len(current) + len(sentence) + 1 > max_size:
                if current:
                    chunks.append(current.strip())
                current = sentence
            else:
                current = f"{current} {sentence}" if current else sentence

        if current:
            chunks.append(current.strip())
        return chunks