"""Offline, disk-backed Codec2 / turbo / neural channel pipeline."""
import csv
import hashlib
import io
import json
import math
import os
from pathlib import Path
import subprocess
import sys
import time
import uuid

import numpy as np
import soundfile as sf
from scipy.signal import resample_poly
import torch

from ccsds_turbo import TurboCode
from sequence_channel import SequenceChannel

ROOT = Path(__file__).resolve().parent
RUNS = ROOT / 'runs'
MAX_SECONDS = 300
ITERATIONS = 12


def configure_cpu(cores=20):
    allowed = sorted(os.sched_getaffinity(0)) if hasattr(os, 'sched_getaffinity') else list(range(os.cpu_count() or 1))
    representatives = {}
    if Path('/proc/cpuinfo').exists():
        for section in Path('/proc/cpuinfo').read_text().split('\n\n'):
            values = dict(line.split(':', 1) for line in section.splitlines() if ':' in line)
            values = {k.strip(): v.strip() for k, v in values.items()}
            if 'processor' in values and int(values['processor']) in allowed:
                key = (values.get('physical id', '0'), values.get('core id', values['processor']))
                representatives.setdefault(key, int(values['processor']))
    selected = sorted(representatives.values() or allowed)[:max(1, cores)]
    if hasattr(os, 'sched_setaffinity'):
        os.sched_setaffinity(0, selected)
    torch.set_num_threads(len(selected))
    torch.set_num_interop_threads(1)
    return selected


def save_json(path, value):
    path.write_text(json.dumps(value, indent=2) + '\n')


def new_run(kind):
    run_id = uuid.uuid4().hex
    folder = RUNS / run_id
    folder.mkdir(parents=True)
    return folder, dict(id=run_id, kind=kind, created_unix=time.time())


def get_run(run_id):
    if not isinstance(run_id, str) or len(run_id) != 32 or any(c not in '0123456789abcdef' for c in run_id):
        raise ValueError('Invalid run ID')
    folder = RUNS / run_id
    if not (folder / 'manifest.json').is_file():
        raise ValueError('Run not found')
    return folder, json.loads((folder / 'manifest.json').read_text())


def asset(folder, name):
    return f'/runs/{folder.name}/{name}'


def codec_worker(operation, source, destination, samples=None):
    command = [sys.executable, str(ROOT / 'codec2_worker.py'), operation, str(source), str(destination)]
    if samples is not None:
        command += ['--samples', str(samples)]
    env = dict(os.environ, OPENBLAS_NUM_THREADS='1', OMP_NUM_THREADS='1')
    process = subprocess.run(command, capture_output=True, text=True, timeout=180, env=env)
    if process.returncode:
        raise RuntimeError('Codec2 operation failed: ' + process.stderr[-2000:])


def read_voice(wav):
    if len(wav) < 44 or wav[:4] not in (b'RIFF', b'RF64') or wav[8:12] != b'WAVE':
        raise ValueError('Please upload a WAV file')
    with sf.SoundFile(io.BytesIO(wav)) as stream:
        if not 8000 <= stream.samplerate <= 192000 or not 1 <= stream.channels <= 8:
            raise ValueError('Use a WAV with 1–8 channels and a sample rate from 8 to 192 kHz')
        if not 0 < stream.frames <= MAX_SECONDS * stream.samplerate:
            raise ValueError(f'Voice must be between one sample and {MAX_SECONDS} seconds long')
        rate, channels, subtype = stream.samplerate, stream.channels, stream.subtype
        samples = stream.read(dtype='float64', always_2d=True).mean(axis=1)
    if not np.isfinite(samples).all():
        raise ValueError('WAV contains non-finite samples')
    if rate != 16000:
        divisor = math.gcd(rate, 16000)
        samples = resample_poly(samples, 16000 // divisor, rate // divisor)
    return np.clip(samples, -1, 32767 / 32768), dict(sample_rate=rate, channels=channels, subtype=subtype)


def prepare_voice(wav, progress):
    folder,manifest=new_run('voice');progress(3,'Reading and resampling voice')
    samples,original_format=read_voice(wav);count=len(samples);seconds=math.ceil(count/16000)
    sf.write(folder/'source.wav',samples,16000,subtype='PCM_16');pcm,_=sf.read(folder/'source.wav',dtype='int16')
    sf.write(folder/'padded_source.wav',np.pad(pcm,(0,seconds*16000-count)),16000,subtype='PCM_16')
    progress(20,'Encoding Codec2 at 3200 bit/s');codec_worker('encode',folder/'padded_source.wav',folder/'source_codec2.bit')
    compressed=(folder/'source_codec2.bit').read_bytes()
    if len(compressed)!=400*seconds:raise ValueError('Unexpected Codec2 3200 payload length')
    payload=np.unpackbits(np.frombuffer(compressed,dtype=np.uint8),bitorder='big').reshape(seconds,3200)
    np.save(folder/'source_payload.npy',payload);np.packbits(payload.reshape(-1),bitorder='big').tofile(folder/'source_payload.bin')
    progress(65,'Decoding the Codec2-only reference');codec_worker('decode',folder/'source_codec2.bit',folder/'codec2_reference.wav',(count+1)//2)
    manifest.update(seconds=seconds,original_samples_16k=count,duration=count/16000,padding_seconds=seconds-count/16000,
        input_format=original_format,codec2_bits=seconds*3200,payload_bits=seconds*3200,
        files={name:asset(folder,name) for name in ['source.wav','codec2_reference.wav','source_codec2.bit','source_payload.bin','source_payload.npy']})
    save_json(folder/'manifest.json',manifest);return manifest

class NeuralChannel:
    def __init__(self):
        self.codec=TurboCode();self.channel=SequenceChannel(ROOT/'models')
    def encode(self,source_id,progress):
        source_folder,source=get_run(source_id)
        if source['kind']!='voice':raise ValueError('Choose a Codec2 voice first')
        folder,manifest=new_run('carrier');seconds=source['seconds'];payload=np.load(source_folder/'source_payload.npy',allow_pickle=False)
        coded=np.stack([self.codec.encode(u) for u in payload]);progress(2,'Generating the 10-frame neural carrier')
        self.channel.generate(coded,folder/'carrier.wav',lambda p:progress(5+94*p,'Generating carrier'))
        np.save(folder/'coded_vectors.npy',coded.reshape(seconds,10,10,64))
        manifest.update(source_id=source_id,seconds=seconds,duration=seconds,frames=seconds*100,coded_bits=seconds*6400,
            original_samples_16k=source['original_samples_16k'],model_checkpoint=self.channel.metadata['checkpoint']['step'],
            format='ccsds-derived-3200-6400-v1',files={name:asset(folder,name) for name in ['carrier.wav','coded_vectors.npy']})
        save_json(folder/'manifest.json',manifest);return manifest
    def import_carrier(self,wav):
        try:info=sf.info(io.BytesIO(wav))
        except Exception as e:raise ValueError('Cannot read carrier WAV') from e
        if info.format not in ('WAV','WAVEX') or info.samplerate!=16000 or info.channels!=1 or info.subtype!='PCM_16' or not 0<info.frames<=MAX_SECONDS*16000 or info.frames%16000:
            raise ValueError('Carrier must be whole seconds of mono PCM16 WAV at 16 kHz (1–300 seconds)')
        folder,manifest=new_run('carrier');(folder/'carrier.wav').write_bytes(wav)
        manifest.update(source_id=None,seconds=info.frames//16000,duration=info.duration,original_samples_16k=info.frames,
            format='ccsds-derived-3200-6400-v1',files={'carrier.wav':asset(folder,'carrier.wav')})
        save_json(folder/'manifest.json',manifest);return manifest
    def decode(self,carrier_id,progress):
        cf,carrier=get_run(carrier_id)
        if carrier['kind']!='carrier':raise ValueError('Choose a carrier first')
        folder,manifest=new_run('decoded');seconds=carrier['seconds'];progress(1,'Reading saved PCM16 and analyzing 100 ms blocks')
        logits=self.channel.analyze(cf/'carrier.wav',lambda p:progress(2+18*p,'Analyzing the saved carrier'))
        probabilities=1/(1+np.exp(-np.clip(logits,-80,80)));recovered=np.empty((seconds,3200),dtype=np.uint8);diagnostics=[]
        # This loop has no access to source payloads, filler, or generated codewords.
        for second in range(seconds):
            detail=self.codec.decode_llr(-logits[second],iterations=ITERATIONS,return_details=True)
            recovered[second]=detail['bits'];diagnostics.append({'second':second+1,'iterations':detail['iterations'],'stable':detail['converged_stable']})
            progress(20+73*(second+1)/seconds,f'Turbo decoding: {second+1}/{seconds} seconds')
        compressed=np.packbits(recovered.reshape(-1),bitorder='big').tobytes()
        (folder/'recovered_payload.bin').write_bytes(compressed);(folder/'recovered_codec2.bit').write_bytes(compressed)
        np.savez_compressed(folder/'received_arrays.npz',probability_one=probabilities,logits=logits,hard_bits=(logits>=0).astype(np.uint8),recovered_payload=recovered)
        progress(95,'Decoding recovered Codec2 speech');codec_worker('decode',folder/'recovered_codec2.bit',folder/'decoded.wav',(carrier['original_samples_16k']+1)//2)
        statistics=None;blocks=[]
        # Source references enter ONLY after channel analysis and turbo decoding.
        if carrier.get('source_id'):
            source_folder,_=get_run(carrier['source_id']);original=np.load(source_folder/'source_payload.npy',allow_pickle=False)
            coded=np.load(cf/'coded_vectors.npy',allow_pickle=False).reshape(seconds,6400);hard=(logits>=0).astype(np.uint8)
            def stat(actual,predicted):
                flipped=int(np.count_nonzero(actual!=predicted));total=int(actual.size)
                return dict(total=total,correct=total-flipped,flipped=flipped,error_percent=100*flipped/total,
                    mean_correct_per_second=(total-flipped)/seconds,mean_flipped_per_second=flipped/seconds)
            reference,_=sf.read(source_folder/'codec2_reference.wav',dtype='int16');decoded,_=sf.read(folder/'decoded.wav',dtype='int16')
            statistics=dict(coded=stat(coded,hard),payload=stat(original,recovered),codec2=stat(original,recovered),
                codec2_bytes_identical=(source_folder/'source_codec2.bit').read_bytes()==compressed,
                decoded_pcm_identical_to_reference=bool(np.array_equal(reference,decoded)),
                failed_payload_blocks=int(np.count_nonzero(np.any(original!=recovered,axis=1))))
            for i in range(seconds):
                raw=int(np.count_nonzero(coded[i]!=hard[i]));res=int(np.count_nonzero(original[i]!=recovered[i]))
                blocks.append(dict(second=i+1,raw_correct=6400-raw,raw_flipped=raw,payload_correct=3200-res,payload_flipped=res,codec2_flipped=res,iterations=diagnostics[i]['iterations']))
            with (folder/'block_statistics.csv').open('w',newline='') as f:
                w=csv.DictWriter(f,fieldnames=list(blocks[0]));w.writeheader();w.writerows(blocks)
        manifest.update(carrier_id=carrier_id,source_id=carrier.get('source_id'),seconds=seconds,duration=carrier['original_samples_16k']/16000,
            statistics=statistics,blocks=blocks,decoder_diagnostics=diagnostics,
            configuration=dict(codec2_bps=3200,payload_bps=3200,channel_bps=6400,code_rate=.5,iterations=ITERATIONS,algorithm='logmap',
                turbo_profile='CCSDS-derived custom k=3200; tail symbols omitted',public_randomizer='CCSDS degree17; reset each second',
                physical_permutation='(173*j+29) mod 6400',model_checkpoint=self.channel.metadata['checkpoint']['step'],
                generator_sha256=self.channel.metadata['exports']['generator.ts'],analyzer_sha256=self.channel.metadata['exports']['analyzer_logits.ts'],
                channel='Saved aligned PCM16 WAV; no added noise or acoustic transmission'),
            files={name:asset(folder,name) for name in ['decoded.wav','recovered_codec2.bit','recovered_payload.bin','received_arrays.npz','results.json']})
        if blocks:manifest['files']['block_statistics.csv']=asset(folder,'block_statistics.csv')
        save_json(folder/'results.json',manifest);save_json(folder/'manifest.json',manifest);return manifest