from pathlib import Path
from typing import Optional, Union, List, Generator, Any, Dict
import numpy as np
import torch
import requests
import json
import asyncio
import logging
from .standard import VieNeuTTS
from .utils import _linear_overlap_add
from vieneu_utils.phonemize_text import phonemize_with_dict
from vieneu_utils.core_utils import split_text_into_chunks, join_audio_chunks

logger = logging.getLogger("Vieneu.Remote")

class RemoteVieNeuTTS(VieNeuTTS):
    """
    Client for VieNeu-TTS running on a remote LMDeploy server.
    """

    def __init__(
        self,
        api_base: str = "http://localhost:23333/v1",
        model_name: str = "pnnbao-ump/VieNeu-TTS",
        codec_repo: str = "neuphonic/distill-neucodec",
        codec_device: str = "cpu",
        hf_token: Optional[str] = None
    ):
        self.api_base = api_base.rstrip('/')
        self.model_name = model_name

        super().__init__(
            backbone_repo=None,
            codec_repo=codec_repo,
            codec_device=codec_device,
            hf_token=hf_token
        )

        self.streaming_frames_per_chunk = 10
        self.streaming_lookforward = 5
        self.streaming_lookback = 50
        self.streaming_stride_samples = self.streaming_frames_per_chunk * self.hop_length
        self._load_voices_from_repo(model_name, hf_token)

    def _load_backbone(self, backbone_repo, backbone_device, hf_token=None):
        pass

    def _format_prompt(self, ref_codes: Union[List[int], torch.Tensor, np.ndarray], ref_text: str, input_text: str) -> str:
        if isinstance(ref_codes, (torch.Tensor, np.ndarray)):
            ref_codes_list = ref_codes.flatten().tolist()
        else:
            ref_codes_list = ref_codes

        ref_text_phones = phonemize_with_dict(ref_text)
        input_text_phones = phonemize_with_dict(input_text, skip_normalize=True)
        codes_str = "".join([f"<|speech_{idx}|>" for idx in ref_codes_list])
        return (
            f"user: Convert the text to speech:<|TEXT_PROMPT_START|>{ref_text_phones} {input_text_phones}"
            f"<|TEXT_PROMPT_END|>\nassistant:<|SPEECH_GENERATION_START|>{codes_str}"
        )

    def infer(self, text: str, ref_audio: Optional[Union[str, Path]] = None, ref_codes: Optional[Union[np.ndarray, torch.Tensor]] = None, ref_text: Optional[str] = None, max_chars: int = 256, silence_p: float = 0.15, crossfade_p: float = 0.0, voice: Optional[Dict[str, Any]] = None, temperature: float = 1.0, top_k: int = 50, skip_normalize: bool = False) -> np.ndarray:

        ref_codes, ref_text = self._resolve_ref_voice(voice, ref_audio, ref_codes, ref_text)

        if not skip_normalize:
            text = self.normalizer.normalize(text)

        chunks = split_text_into_chunks(text, max_chars=max_chars)
        if not chunks:
            return np.array([], dtype=np.float32)

        all_wavs = []
        for chunk in chunks:
            prompt = self._format_prompt(ref_codes, ref_text, chunk)
            payload = {
                "model": self.model_name,
                "messages": [{"role": "user", "content": prompt}],
                "max_tokens": 2048,
                "temperature": temperature,
                "top_k": top_k,
                "stop": ["<|SPEECH_GENERATION_END|>"],
                "stream": False
            }
            try:
                response = requests.post(f"{self.api_base}/chat/completions", json=payload, timeout=60)
                response.raise_for_status()
                output_str = response.json()["choices"][0]["message"]["content"]
                wav = self._decode(output_str)
                all_wavs.append(wav)
            except Exception as e:
                logger.error(f"Error during remote inference: {e}")
                continue

        final_wav = join_audio_chunks(all_wavs, self.sample_rate, silence_p, crossfade_p)
        return self._apply_watermark(final_wav)

    def infer_stream(self, text: str, ref_audio: Optional[Union[str, Path]] = None, ref_codes: Optional[Union[np.ndarray, torch.Tensor]] = None, ref_text: Optional[str] = None, max_chars: int = 256, voice: Optional[Dict[str, Any]] = None, temperature: float = 1.0, top_k: int = 50, skip_normalize: bool = False) -> Generator[np.ndarray, None, None]:

        ref_codes, ref_text = self._resolve_ref_voice(voice, ref_audio, ref_codes, ref_text)

        if not skip_normalize:
            text = self.normalizer.normalize(text)

        chunks = split_text_into_chunks(text, max_chars=max_chars)
        for chunk in chunks:
            yield from self._infer_stream_chunk(chunk, ref_codes, ref_text, temperature, top_k)

    def _infer_stream_chunk(self, chunk, ref_codes, ref_text, temperature, top_k):
        prompt = self._format_prompt(ref_codes, ref_text, chunk)
        payload = {
            "model": self.model_name,
            "messages": [{"role": "user", "content": prompt}],
            "max_tokens": 2048,
            "temperature": temperature,
            "top_k": top_k,
            "stop": ["<|SPEECH_GENERATION_END|>"],
            "stream": True
        }

        if isinstance(ref_codes, (torch.Tensor, np.ndarray)):
            ref_codes_list = ref_codes.flatten().tolist()
        else:
            ref_codes_list = ref_codes

        audio_cache: List[np.ndarray] = []
        token_cache: List[str] = [f"<|speech_{idx}|>" for idx in ref_codes_list]
        n_decoded_samples: int = 0
        n_decoded_tokens: int = len(ref_codes_list)

        try:
             with requests.post(f"{self.api_base}/chat/completions", json=payload, stream=True, timeout=60) as r:
                r.raise_for_status()
                for line in r.iter_lines():
                    if not line: continue
                    line_str = line.decode('utf-8')
                    if not line_str.startswith('data: '): continue
                    data_str = line_str[6:]
                    if data_str == '[DONE]': break
                    try:
                        content = json.loads(data_str)["choices"][0]["delta"].get("content", "")
                        if content:
                             token_cache.append(content)
                             if len(token_cache[n_decoded_tokens:]) >= self.streaming_frames_per_chunk + self.streaming_lookforward:
                                tokens_start = max(n_decoded_tokens - self.streaming_lookback - self.streaming_overlap_frames, 0)
                                tokens_end = n_decoded_tokens + self.streaming_frames_per_chunk + self.streaming_lookforward + self.streaming_overlap_frames
                                sample_start = (n_decoded_tokens - tokens_start) * self.hop_length
                                sample_end = sample_start + (self.streaming_frames_per_chunk + 2 * self.streaming_overlap_frames) * self.hop_length
                                curr_codes = token_cache[tokens_start:tokens_end]
                                recon = self._decode("".join(curr_codes))
                                recon = self._apply_watermark(recon)
                                recon = recon[sample_start:sample_end]
                                audio_cache.append(recon)
                                processed_recon = _linear_overlap_add(audio_cache, stride=self.streaming_stride_samples)
                                new_samples_end = len(audio_cache) * self.streaming_stride_samples
                                processed_recon = processed_recon[n_decoded_samples:new_samples_end]
                                n_decoded_samples = new_samples_end
                                n_decoded_tokens += self.streaming_frames_per_chunk
                                yield processed_recon
                    except json.JSONDecodeError: continue
        except Exception as e:
            logger.error(f"Error streaming chunk: {e}")
            return

        remaining_tokens = len(token_cache) - n_decoded_tokens
        if remaining_tokens > 0:
            tokens_start = max(len(token_cache) - (self.streaming_lookback + self.streaming_overlap_frames + remaining_tokens), 0)
            sample_start = (len(token_cache) - tokens_start - remaining_tokens - self.streaming_overlap_frames) * self.hop_length
            curr_codes = token_cache[tokens_start:]
            recon = self._decode("".join(curr_codes))
            recon = self._apply_watermark(recon)
            recon = recon[sample_start:]
            audio_cache.append(recon)
            processed_recon = _linear_overlap_add(audio_cache, stride=self.streaming_stride_samples)
            processed_recon = processed_recon[n_decoded_samples:]
            yield processed_recon

    async def infer_async(self, text: str, ref_audio: Optional[Union[str, Path]] = None, ref_codes: Optional[Union[np.ndarray, torch.Tensor]] = None, ref_text: Optional[str] = None, max_chars: int = 256, silence_p: float = 0.15, crossfade_p: float = 0.0, voice: Optional[Dict[str, Any]] = None, temperature: float = 1.0, top_k: int = 50, session=None, skip_normalize: bool = False) -> np.ndarray:
        try:
            import aiohttp
        except ImportError:
            raise ImportError("Async requires 'aiohttp'.")

        ref_codes, ref_text = self._resolve_ref_voice(voice, ref_audio, ref_codes, ref_text)

        if not skip_normalize:
            text = self.normalizer.normalize(text)

        chunks = split_text_into_chunks(text, max_chars=max_chars)
        if not chunks:
            return np.array([], dtype=np.float32)

        should_close_session = False
        if session is None:
            session = aiohttp.ClientSession()
            should_close_session = True

        try:
            tasks = [self._infer_chunk_async(session, chunk, ref_codes, ref_text, temperature, top_k) for chunk in chunks]
            wavs = await asyncio.gather(*tasks)
            final_wav = join_audio_chunks(wavs, self.sample_rate, silence_p, crossfade_p)
            return self._apply_watermark(final_wav)
        finally:
            if should_close_session:
                await session.close()

    async def _infer_chunk_async(self, session, chunk, ref_codes, ref_text, temperature, top_k):
        prompt = self._format_prompt(ref_codes, ref_text, chunk)
        payload = {
            "model": self.model_name,
            "messages": [{"role": "user", "content": prompt}],
            "max_tokens": 2048,
            "temperature": temperature,
            "top_k": top_k,
            "stop": ["<|SPEECH_GENERATION_END|>"],
            "stream": False
        }
        try:
            async with session.post(f"{self.api_base}/chat/completions", json=payload, timeout=60) as resp:
                resp.raise_for_status()
                data = await resp.json()
                output_str = data["choices"][0]["message"]["content"]
                return self._decode(output_str)
        except Exception as e:
            logger.error(f"Error in async chunk: {e}")
            return np.array([], dtype=np.float32)

    async def infer_batch_async(self, texts: List[str], ref_audio: Optional[Union[str, Path]] = None, ref_codes: Optional[Union[np.ndarray, torch.Tensor]] = None, ref_text: Optional[str] = None, max_chars: int = 256, silence_p: float = 0.15, crossfade_p: float = 0.0, voice: Optional[Dict[str, Any]] = None, temperature: float = 1.0, top_k: int = 50, concurrency_limit: int = 50, skip_normalize: bool = False) -> List[np.ndarray]:
        try:
            import aiohttp
        except ImportError:
            raise ImportError("Async requires 'aiohttp'.")

        if not skip_normalize:
            texts = [self.normalizer.normalize(t) for t in texts]

        ref_codes, ref_text = self._resolve_ref_voice(voice, ref_audio, ref_codes, ref_text)

        sem = asyncio.Semaphore(concurrency_limit)
        async with aiohttp.ClientSession() as session:
            async def bounded_infer(text):
                async with sem:
                    return await self.infer_async(
                        text, ref_codes=ref_codes, ref_text=ref_text,
                        max_chars=max_chars, silence_p=silence_p, crossfade_p=crossfade_p,
                        temperature=temperature, top_k=top_k,
                        session=session, skip_normalize=True
                    )
            tasks = [bounded_infer(text) for text in texts]
            results = await asyncio.gather(*tasks)
        return results
