"""Self-contained inference for Gooya v2 experimental releases."""
from pathlib import Path
import json,re,warnings
import numpy as np
import torch
from torch import nn
from torch.nn import functional as F
from torch.nn.utils.rnn import pack_padded_sequence,pad_packed_sequence
from torch.nn.utils import weight_norm
from safetensors.torch import load_file
from transformers import AutoTokenizer,AutoModelForSeq2SeqLM
from .networks import PhonemeEncoder,MelDecoder,Phoneme2Mel
from .hifigan import Generator
from .normalize import normalize_persian_for_g2p
from .words import surface_words

class ContextEncoder(nn.Module):
    def __init__(self,local):
        super().__init__(); self.local=local; self.context=nn.GRU(80,128,batch_first=True,bidirectional=True); self.project=nn.Linear(256,32); self.dropout=nn.Dropout(.1)
    def forward(self,phones):
        x=self.local.embed(phones); lengths=(phones!=0).sum(-1).cpu()
        packed=pack_padded_sequence(x,lengths,batch_first=True,enforce_sorted=False)
        y,_=self.context(packed); y,_=pad_packed_sequence(y,batch_first=True,total_length=phones.shape[1])
        return self.local(phones)+self.project(self.dropout(y))

class SmoothDuration(nn.Module):
    def __init__(self,original):
        super().__init__(); self.original=original; self.original.duration=False
    def forward(self,x): return F.softplus(self.original(x))+1
    def get_embedding(self,*args,**kwargs): return self.original.get_embedding(*args,**kwargs)

class AttrDict(dict):
    __getattr__=dict.__getitem__

class GooyaTTS:
    """Loads all model weights locally. CPU is the default; CUDA is optional."""
    def __init__(self,model_dir,device='cpu',num_threads=4):
        self.root=Path(model_dir); self.device=torch.device(device)
        if num_threads: torch.set_num_threads(num_threads)
        self.config=json.loads((self.root/'gooya_config.json').read_text()); self.sample_rate=self.config['sample_rate']; self.ids=self.config['vocabulary']
        self.tokenizer=AutoTokenizer.from_pretrained(self.root/'frontend',local_files_only=True)
        self.g2p=AutoModelForSeq2SeqLM.from_pretrained(self.root/'frontend',local_files_only=True).to(self.device).eval()
        stats=self.config['stats']; encoder=PhonemeEncoder(pitch_stats=stats['pitch'][:2],energy_stats=stats['energy'][:2]); decoder=MelDecoder()
        encoder.encoder.embed=nn.Embedding(len(self.ids),80,padding_idx=0)
        encoder.boundary_token_ids=[i for s,i in self.ids.items() if s.startswith('<')]
        encoder.encoder=ContextEncoder(encoder.encoder); encoder.duration_decoder=SmoothDuration(encoder.duration_decoder)
        decoder.mel_linear_down=nn.Linear(320,100); self.acoustic=Phoneme2Mel(encoder,decoder)
        self.acoustic.load_state_dict(load_file(str(self.root/'acoustic.safetensors')),strict=True); self.acoustic.to(self.device).eval()
        self.vocoder=Generator(AttrDict(self.config['vocoder_config'])); self.vocoder.conv_pre=weight_norm(nn.Conv1d(100,128,7,1,padding=3))
        self.vocoder.load_state_dict(load_file(str(self.root/'vocoder.safetensors')),strict=True); self.vocoder.to(self.device).eval()
        self.overlay={}
        if self.config['frontend']['overlay']:
            data=json.loads((self.root/'frontend/overlay.json').read_text()); self.overlay={(r['surface'],r['raw']):r['target'] for r in data['rules']}
        count=sum(p.numel() for m in [self.g2p,self.acoustic,self.vocoder] for p in m.parameters())
        if count!=self.config['parameters']['total']: raise ValueError(f'Parameter count mismatch: {count}')
    @classmethod
    def from_pretrained(cls,repo_id,revision=None,**kwargs):
        from huggingface_hub import snapshot_download
        return cls(snapshot_download(repo_id,revision=revision),**kwargs)
    @torch.inference_mode()
    def synthesize(self,text,return_details=False):
        if not isinstance(text,str) or not text.strip(): raise ValueError('Text must be a nonempty string.')
        chunks=re.findall(r'[^،,؟?\.\n]+[،,؟?\.]?',text); waves=[]; records=[]
        for chunk in chunks:
            chunk=chunk.strip()
            if not chunk: continue
            normalized=normalize_persian_for_g2p(chunk)
            x=self.tokenizer(normalized,return_tensors='pt',add_special_tokens=False).to(self.device)
            if x['input_ids'].shape[1]>1024: raise ValueError('Clause exceeds 1024 input tokens; use shorter sentences or punctuation.')
            output=self.g2p.generate(**x,max_new_tokens=self.config['frontend']['max_new_tokens'],num_beams=self.config['frontend']['num_beams'])
            if self.tokenizer.eos_token_id not in output[0,1:].tolist(): raise RuntimeError('G2P did not terminate. Split this clause into shorter phrases; no incomplete audio was returned.')
            phones=self.tokenizer.decode(output[0],skip_special_tokens=True).strip().replace('?','Q')
            words=surface_words(normalized); pw=phones.split(); matched=len(words)==len(pw)
            if self.overlay and matched: phones=' '.join(self.overlay.get((a,b),b) for a,b in zip(words,pw))
            if not matched: warnings.warn('G2P word count differs from input; inspect pronunciation in returned details.',RuntimeWarning)
            if not phones: raise RuntimeError('G2P returned an empty pronunciation.')
            tokens=['<sil>']+['<wb>' if c==' ' else c for c in phones]+['<comma>' if chunk[-1] in '،,' else '<stop>']
            unknown=set(tokens)-set(self.ids)
            if unknown: raise ValueError(f'Unsupported G2P symbols: {sorted(unknown)}')
            encoded=torch.tensor([[self.ids[t] for t in tokens]],device=self.device)
            pred=self.acoustic.encoder({'phoneme':encoded},train=False); mel=self.acoustic.decoder(pred['features'])
            if mel.shape[1]>24000//256*120: raise RuntimeError('Predicted clause duration exceeds 120 seconds.')
            wave=self.vocoder(F.pad(mel.transpose(1,2),(4,4),mode='replicate'))[:,:,1024:-1024].squeeze().cpu().numpy()
            if not np.isfinite(wave).all(): raise RuntimeError('Nonfinite audio output.')
            waves.append(wave); gap=self.config['comma_gap_seconds'] if chunk[-1] in '،,' else self.config['sentence_gap_seconds']; waves.append(np.zeros(round(gap*self.sample_rate),dtype=np.float32))
            records.append({'text':chunk,'normalized':normalized,'phonemes':phones,'word_count_matches':matched,'duration_seconds':len(wave)/self.sample_rate})
        if not waves: raise ValueError('No speakable text found.')
        result=np.concatenate(waves[:-1]); details={'text':text,'model':self.config['name'],'parameters':self.config['parameters'],'sample_rate':self.sample_rate,'segments':records,'duration_seconds':len(result)/self.sample_rate}
        return (result,details) if return_details else result