import base64,io,json,re,threading
from pathlib import Path
import numpy as np
from PIL import Image
import torch
from .networks import SketchEncoder,SpikingReadout,PongSNN,CATEGORIES,NavigationPolicy

ROOT=Path(__file__).resolve().parents[1]
DEVICE='mps' if torch.backends.mps.is_available() else 'cpu'
torch.set_num_threads(12)
LOCK=threading.RLock()
CACHE={}

def decode_image(data):
    if not isinstance(data,str) or len(data)>12_000_000: raise ValueError('Image is too large.')
    raw=base64.b64decode(data.split(',',1)[-1],validate=True)
    im=Image.open(io.BytesIO(raw))
    if im.width*im.height>16_000_000: raise ValueError('Image exceeds 16 megapixels.')
    return im.convert('RGB')

def encode_image(im):
    b=io.BytesIO(); im.save(b,format='PNG')
    return 'data:image/png;base64,'+base64.b64encode(b.getvalue()).decode()

def load(name):
    if name in CACHE: return CACHE[name]
    path=ROOT/'models'/'hf'/name
    if not (path/'config.json').exists() and name!='image':
        raise RuntimeError(f'{name} model has not been downloaded. Run scripts/download_models.py.')
    from transformers import AutoTokenizer,AutoModelForCausalLM,CLIPProcessor,CLIPModel,BlipProcessor,BlipForConditionalGeneration
    if name=='language':
        tok=AutoTokenizer.from_pretrained(path,local_files_only=True)
        model=AutoModelForCausalLM.from_pretrained(path,local_files_only=True,torch_dtype=torch.float32).to(DEVICE).eval()
        obj=(tok,model)
    elif name=='vision':
        obj=(CLIPProcessor.from_pretrained(path,local_files_only=True),CLIPModel.from_pretrained(path,local_files_only=True).to(DEVICE).eval())
    elif name=='caption':
        obj=(BlipProcessor.from_pretrained(path,local_files_only=True),BlipForConditionalGeneration.from_pretrained(path,local_files_only=True).to(DEVICE).eval())
    elif name=='speech':
        from transformers import WhisperProcessor,WhisperForConditionalGeneration
        obj=(WhisperProcessor.from_pretrained(path,local_files_only=True),WhisperForConditionalGeneration.from_pretrained(path,local_files_only=True).to(DEVICE).eval())
    elif name=='image':
        from diffusers import AutoPipelineForText2Image
        # Half-precision VAE decoding overflows on some Apple GPUs. Load fp16 files
        # into float32 for stable local generation, including the decoder.
        obj=AutoPipelineForText2Image.from_pretrained(path,local_files_only=True,torch_dtype=torch.float32,variant='fp16',use_safetensors=True).to(DEVICE)
        obj.enable_attention_slicing()
    else: raise ValueError('Unknown model')
    CACHE[name]=obj
    return obj

def language(prompt,max_tokens=160):
    with LOCK,torch.inference_mode():
        tok,model=load('language')
        text=tok.apply_chat_template([{'role':'system','content':'You are a helpful assistant for a classroom AI brain exhibit. Follow the requested format exactly. Be brief.'},{'role':'user','content':prompt}],tokenize=False,add_generation_prompt=True)
        inputs=tok(text,return_tensors='pt').to(DEVICE)
        out=model.generate(**inputs,max_new_tokens=max_tokens,do_sample=False,pad_token_id=tok.eos_token_id)
        return tok.decode(out[0,inputs.input_ids.shape[1]:],skip_special_tokens=True).strip()

def image_features(im):
    with LOCK,torch.inference_mode():
        proc,model=load('vision'); data=proc(images=im,return_tensors='pt').to(DEVICE)
        feat=model.get_image_features(**data); feat=feat/feat.norm(dim=-1,keepdim=True)
        return feat[0].cpu().numpy()

def text_features(texts):
    with LOCK,torch.inference_mode():
        proc,model=load('vision'); data=proc(text=texts,return_tensors='pt',padding=True,truncation=True).to(DEVICE)
        feat=model.get_text_features(**data); feat=feat/feat.norm(dim=-1,keepdim=True)
        return feat.cpu().numpy()

def classify_image(im,labels):
    feat=image_features(im); refs=text_features(['a photo of '+label for label in labels])
    scores=torch.tensor(refs@feat*30).softmax(0).numpy()
    return sorted([{'label':l,'probability':float(p)} for l,p in zip(labels,scores)],key=lambda a:-a['probability']),feat

def caption(im):
    with LOCK,torch.inference_mode():
        proc,model=load('caption'); data=proc(images=im,return_tensors='pt').to(DEVICE)
        out=model.generate(**data,max_new_tokens=45)
        return proc.decode(out[0],skip_special_tokens=True)

def generate_image(prompt,seed=21):
    with LOCK,torch.inference_mode():
        pipe=load('image'); gen=torch.Generator(device='cpu').manual_seed(seed)
        im=pipe(prompt=prompt,num_inference_steps=2,guidance_scale=0.,height=512,width=512,generator=gen).images[0]
        return encode_image(im)

def sketch_input(im):
    gray=255-np.asarray(im.convert('L'))
    ys,xs=np.where(gray>30)
    if len(xs)<8: return None
    bbox=(int(xs.min()),int(ys.min()),int(xs.max())+1,int(ys.max())+1)
    crop=Image.fromarray(gray).crop(bbox)
    crop.thumbnail((25,25),Image.Resampling.LANCZOS)
    out=Image.new('L',(28,28)); out.paste(crop,((28-crop.width)//2,(28-crop.height)//2))
    return torch.tensor(np.asarray(out,dtype=np.float32)[None,None]/255)

def sketch(im,threshold=1.):
    x=sketch_input(im)
    if x is None: return {'candidates':[],'spikes':[],'membrane':[],'message':'Add a few strokes to begin.'}
    with LOCK,torch.inference_mode():
        if 'sketch' not in CACHE:
            ck=torch.load(ROOT/'models'/'sketch.pt',map_location='cpu',weights_only=True)
            enc=SketchEncoder(); enc.load_state_dict(ck['encoder']); enc.eval()
            net=SpikingReadout(); net.load_state_dict(ck['snn']); net.eval()
            CACHE['sketch']=(enc,net,ck['mean'],ck['std'])
        enc,net,mean,std=CACHE['sketch']
        features=(enc.features(x)-mean)/std
        logits,spikes,mem=net(features,threshold,True)
        scores=logits[0].softmax(0).numpy()
        ids=np.argsort(-scores)[:5]
        return {'candidates':[{'label':CATEGORIES[i],'probability':float(scores[i])} for i in ids],
            'spikes':spikes[:,0,:32].int().tolist(),'membrane':mem[:,0,0].tolist(),
            'spike_count':int(spikes.sum()),'threshold':threshold,'model':'CNN encoder → trained LIF layer → linear class readout'}

def pong_config():
    receipt=json.loads((ROOT/'models'/'pong_v2_receipt.json').read_text())
    return {'model':'Target-position LIF SNN','optimal_threshold':receipt['optimal_threshold'],
            'steps':10,'neurons':64,'validation_target_mae':receipt['threshold_validation_mae'][str(receipt['optimal_threshold'])],
            'imitation_accuracy':receipt['imitation_accuracy'],'test_states':receipt['test_states'],
            'benchmark_scope':receipt['benchmark_scope']}

def pong(state,threshold):
    if len(state)!=5 or not np.isfinite(state).all():raise ValueError('Use five finite game observations.')
    with LOCK,torch.inference_mode():
        if 'pong' not in CACHE:
            net=PongSNN();net.load_state_dict(torch.load(ROOT/'models'/'pong_v2.pt',weights_only=True));net.eval();CACHE['pong']=net
        goal,spikes,mem=CACHE['pong'](torch.tensor([state],dtype=torch.float32),threshold,True)
        active=bool(spikes.sum()>0)
        target=float(goal[0].clamp(.095,.905)) if active else float(state[4])
        delta=target-state[4]
        action=-1 if delta<-.025 else 1 if delta>.025 else 0
        return {'action':action,'target_y':target,'decision_active':active,'model':'Target-position LIF SNN',
            'spikes':spikes[:,0,:32].int().tolist(),'membrane':mem[:,0,0].tolist(),
            'spike_count':int(spikes.sum()),'threshold':threshold,'neurons':64,'steps':10}

def navigate(dx,dy,blocked):
    with LOCK,torch.inference_mode():
        if 'navigation' not in CACHE:
            net=NavigationPolicy(); net.load_state_dict(torch.load(ROOT/'models'/'navigation.pt',weights_only=True)); net.eval(); CACHE['navigation']=net
        logits=CACHE['navigation'](torch.tensor([[dx,dy,*blocked]],dtype=torch.float32))[0]
        # Physical collision gate rejects impossible movement, independently of policy.
        for i,b in enumerate(blocked):
            if b: logits[i]=-1e6
        return int(logits.argmax()),logits.softmax(0).tolist()

def transcribe(path):
    import subprocess,tempfile,soundfile as sf
    with tempfile.NamedTemporaryFile(suffix='.wav') as out:
        subprocess.run(['ffmpeg','-v','error','-y','-i',str(path),'-ar','16000','-ac','1',out.name],check=True,capture_output=True)
        audio,sr=sf.read(out.name)
    if len(audio)>16000*30: raise ValueError('Please record 30 seconds or less.')
    with LOCK,torch.inference_mode():
        proc,model=load('speech'); inp=proc(audio,sampling_rate=16000,return_tensors='pt').input_features.to(DEVICE)
        out=model.generate(inp,max_new_tokens=128)
        return proc.batch_decode(out,skip_special_tokens=True)[0].strip()