/* SPDX-License-Identifier: MIT. RFXA/RFXZ-26 memory codec.
 * Link RFXL, RAU, LeXA and shared PackBits/BriefLZ framing sources.
 */
#include "rfxlanim.h"
#include "../../rfxl/codecs/rfxl.h"
#include "../../rau/codecs/rau.h"
#include "../../lexa/codecs/lexa.h"
#include "../../shared/packbits.h"
#include "../../shared/brieflz-container.h"
#include <stdlib.h>
#include <string.h>
#define LIMIT ((size_t)67108864)
typedef struct { const uint8_t *s; size_t n,p; int error; } reader;
static unsigned u16(const uint8_t *s) { return s[0] | (unsigned)s[1] << 8; }
static uint32_t u32(const uint8_t *s) { return s[0] | (uint32_t)s[1] << 8 | (uint32_t)s[2] << 16 | (uint32_t)s[3] << 24; }
static const uint8_t *take(reader *r, size_t n) { const uint8_t *p; if (r->error || r->p > r->n || n > r->n-r->p) { r->error=1; return NULL; } p=r->s+r->p; r->p+=n; return p; }
static unsigned get(reader *r,unsigned n) { const uint8_t *p=take(r,n); return !p?0:n==1?p[0]:n==2?u16(p):u32(p); }
static void put(uint8_t *p,size_t v,unsigned n) { unsigned i; for(i=0;i<n;i++)p[i]=(uint8_t)(v>>(8*i)); }
static const uint8_t *unwrap(const uint8_t *s,size_t n,uint8_t **owned,size_t *size) {
    *owned=NULL; *size=n;
    if (!s || n>LIMIT) return NULL;
    if(n>=4&&!memcmp(s,"RFXZ",4)) {
        if(n<16||s[4]!=26||s[5]!=14||u16(s+6)||u32(s+12))return NULL;
        *size=u32(s+8); *owned=codec_brief_lz_decode(s+16,n-16,*size); return *owned;
    } return s;
}
static int inspect_raw(const uint8_t *s,size_t n,rfxa_info *o) {
    size_t i,table,entry,expected,off,len,audio; unsigned palette; rau_info ai;
    if(!s||n<32||memcmp(s,"RFXA",4)||s[4]!=26||s[5]>4||s[6]&~3||s[7])return -1;
    memset(o,0,sizeof(*o)); o->quality=s[5];o->flags=s[6];o->width=u32(s+8);o->height=u32(s+12);o->count=u32(s+16);o->fps=(uint16_t)u16(s+20);o->audio_rate=(uint16_t)u16(s+22);o->audio_channels=s[24];o->audio_bits=s[25];palette=u16(s+26);table=u32(s+28);entry=o->flags&1?12:8;
    if(!o->width||!o->height||o->width>8192||o->height>8192||(size_t)o->width*o->height>16777216||!o->count||!o->fps||palette>256||!!palette!=!!(o->flags&2)||table!=32+palette*4||table>n||o->count>(n-table)/entry)return -1;
    if(o->flags&1) { if(!o->audio_rate||(o->audio_channels!=1&&o->audio_channels!=2)||(o->audio_bits!=8&&o->audio_bits!=16))return -1; }
    else if(o->audio_rate||o->audio_channels||o->audio_bits)return -1;
    expected=table+(size_t)o->count*entry;
    for(i=0;i<o->count;i++) {
        off=u32(s+table+i*entry);len=u32(s+table+i*entry+4);audio=entry==12?u32(s+table+i*entry+8):0;
        if(off!=expected||off>n||len<10||len>n-off||audio>n-off-len||s[off]>7)return -1;
        if(audio&&(rau_inspect(s+off+len,audio,&ai)||ai.sample_rate!=o->audio_rate||ai.channels!=o->audio_channels||ai.bits!=o->audio_bits))return -1;
        expected+=len+audio;
    }return expected==n?0:-1;
}
int rfxa_inspect(const uint8_t *s,size_t n,rfxa_info *o) { uint8_t *owned;size_t size;const uint8_t *raw;int rc;if(!o)return -1;raw=unwrap(s,n,&owned,&size);rc=inspect_raw(raw,size,o);free(owned);return rc; }
static uint8_t *indices(const uint8_t *s,size_t n,size_t count,unsigned bits) {
    uint8_t *raw=NULL,*out;size_t packed,i;unsigned method;
    if(!s||!n||!count||(bits!=1&&bits!=2&&bits!=4&&bits!=8))return NULL;
    method=*s++;n--;packed=(count*bits+7)/8;
    if(method==0||method==2) { size_t expected=method==2?count:packed;if(n!=expected)return NULL;raw=malloc(expected);if(raw)memcpy(raw,s,n); }
    else if(method==1||method==3)raw=codec_packbits_decode(s,n,method==3?count:packed);
    else if(method==4||method==5)raw=rfxl_unpack(s,n,4,method==5?count:packed);
    if(!raw)return NULL;
    if(method==2||method==3||method==5)return raw;
    out=malloc(count);if(out)for(i=0;i<count;i++)out[i]=(uint8_t)(raw[i*bits/8]>>(i*bits%8)&((1u<<bits)-1));free(raw);return out;
}
static int rect(const uint8_t *s,size_t n,uint8_t *canvas,unsigned W,unsigned H,unsigned x,unsigned y,unsigned w,unsigned h,int indexed,unsigned bits,const uint8_t *palette,unsigned colours) {
    uint8_t *values; rfxl_image image;size_t row,col,pos,dest;
    if(!w||!h||x>W||w>W-x||y>H||h>H-y)return -1;
    if(indexed)values=indices(s,n,(size_t)w*h,bits);
    else { if(rfxl_decode(s,n,&image))return -1;values=image.rgba;if(image.width!=w||image.height!=h){free(values);return -1;} }
    if(!values)return -1;
    for(row=0;row<h;row++)for(col=0;col<w;col++) {
        pos=row*w+col;dest=((row+y)*W+x+col)*4;
        if(indexed){if(values[pos]>=colours){free(values);return -1;}memcpy(canvas+dest,palette+values[pos]*4,4);}
        else memcpy(canvas+dest,values+pos*4,4);
    }
    free(values);return 0;
}
static int apply(const uint8_t *s,size_t n,uint8_t *canvas,const rfxa_info *info,unsigned frame,const uint8_t *palette,unsigned colours) {
    reader r={s,n,0,0},v;unsigned type=get(&r,1),bits=get(&r,1),x=get(&r,2),y=get(&r,2),w=get(&r,2),h=get(&r,2),W=info->width,H=info->height;
    size_t count,i,j,row,col,columns,rows,full_count=0,sparse_count=0,changes=0,full_pos=0,sparse_pos=0,block_pos=0,len;
    const uint8_t *vectors,*payload,*p;int32_t *dx=NULL,*dy=NULL,previous;uint8_t *flags=NULL,*next=NULL,*a=NULL,*b=NULL,*counts=NULL,*deltas=NULL,*sparse=NULL,*full=NULL; rfxl_image image={0,0,NULL},residual={0,0,NULL};unsigned format,prediction,bx,by,bw,bh,state;int rc=-1;
    if(r.error)return -1;
    if(type==0||type==3) { if(x||y||w!=W||h!=H)return -1;return rect(s+10,n-10,canvas,W,H,x,y,w,h,type==3,bits,palette,colours); }
    if(!frame)return -1;
    if(type==1||type==4) { if(!w&&!h)return n==10?0:-1;return rect(s+10,n-10,canvas,W,H,x,y,w,h,type==4,bits,palette,colours); }
    if(type==2||type==5) {
        count=get(&r,2);p=take(&r,count*12);if(!count||!p)return -1;
        for(i=0;i<count;i++){payload=take(&r,u32(p+i*12+8));if(!payload||rect(payload,u32(p+i*12+8),canvas,W,H,u16(p+i*12),u16(p+i*12+2),u16(p+i*12+4),u16(p+i*12+6),type==5,bits,palette,colours))return -1;}
        return r.p==n?0:-1;
    }
    if((type!=6&&type!=7)||x||y||w!=W||h!=H)return -1;
    format=get(&r,1);len=get(&r,4);vectors=take(&r,len);if(!vectors)return -1;
    columns=(W+15)/16;rows=(H+15)/16;count=columns*rows;dx=calloc(count,sizeof(*dx));dy=calloc(count,sizeof(*dy));flags=calloc(count,1);next=malloc((size_t)W*H*4);if(!dx||!dy||!flags||!next)goto done;
    if(format==1) {
        v.s=vectors;v.n=len;v.p=0;v.error=0;prediction=get(&v,1);len=get(&v,4);payload=take(&v,len);if(!payload||prediction>1)goto done;a=rfxl_unpack(payload,len,3,count);
        len=get(&v,4);payload=take(&v,len);if(!payload)goto done;b=rfxl_unpack(payload,len,3,count);if(!a||!b)goto done;
        if(type==7){free(flags);len=get(&v,4);payload=take(&v,len);if(!payload){flags=NULL;goto done;}flags=rfxl_unpack(payload,len,3,count);if(!flags)goto done;}
        if(v.error||v.p!=v.n)goto done;
        previous=0;for(i=0;i<count;i++){int value=a[i]<128?a[i]:(int)a[i]-256;previous=prediction?previous+value:value;dx[i]=previous;}
        previous=0;for(i=0;i<count;i++){int value=b[i]<128?b[i]:(int)b[i]-256;previous=prediction?previous+value:value;dy[i]=previous;}
    }else if(format==0) {
        if(rfxl_decode(vectors,len,&image)||image.width!=columns||image.height!=rows)goto done;
        for(i=0;i<count;i++){dx[i]=(int)image.rgba[i*4]-128;dy[i]=(int)image.rgba[i*4+1]-128;flags[i]=image.rgba[i*4+2];}
    }else goto done;
    for(i=0;i<count;i++)if(type==7){if(flags[i]>2)goto done;if(flags[i]==1)sparse_count++;if(flags[i]==2)full_count+=(W-i%columns*16<16?W-i%columns*16:16)*(H-i/columns*16<16?H-i/columns*16:16);}
    if(type==6){if(rfxl_decode(s+r.p,n-r.p,&residual)||residual.width!=W||residual.height!=H)goto done;}
    else {
        if(sparse_count){len=get(&r,4);payload=take(&r,len);if(!payload)goto done;counts=rfxl_unpack(payload,len,3,sparse_count);if(!counts)goto done;for(i=0;i<sparse_count;i++)changes+=counts[i];if(!changes)goto done;
            len=get(&r,4);payload=take(&r,len);if(!payload)goto done;deltas=rfxl_unpack(payload,len,3,changes);len=get(&r,4);payload=take(&r,len);if(!payload)goto done;sparse=indices(payload,len,changes,bits);if(!deltas||!sparse)goto done;}
        if(full_count){full=indices(s+r.p,n-r.p,full_count,bits);if(!full)goto done;}else if(r.p!=n)goto done;
    }
    for(i=0;i<count;i++) {
        bx=(unsigned)(i%columns)*16;by=(unsigned)(i/columns)*16;bw=W-bx<16?W-bx:16;bh=H-by<16?H-by:16;state=flags[i];
        if((type==6||state!=2)&&((int64_t)bx+dx[i]<0||(int64_t)by+dy[i]<0||(int64_t)bx+dx[i]+bw>W||(int64_t)by+dy[i]+bh>H))goto done;
        for(row=0;row<bh;row++)for(col=0;col<bw;col++) {
            size_t dest=((by+row)*W+bx+col)*4;
            if(type==7&&state==2){unsigned colour=full[full_pos++];if(colour>=colours)goto done;memcpy(next+dest,palette+colour*4,4);}
            else {size_t src=(size_t)(((int64_t)by+(int64_t)row+dy[i])*W+bx+(int64_t)col+dx[i])*4;for(j=0;j<4;j++)next[dest+j]=(uint8_t)(canvas[src+j]+(type==6?residual.rgba[dest+j]:0));}
        }
        if(type==7&&state==1){int position=-1;unsigned c;for(c=0;c<counts[block_pos];c++){unsigned colour=sparse[sparse_pos];position+=1+deltas[sparse_pos++];if((unsigned)position/16>=bh||(unsigned)position%16>=bw||colour>=colours)goto done;memcpy(next+((by+(unsigned)position/16)*W+bx+(unsigned)position%16)*4,palette+colour*4,4);}block_pos++;}
    }
    memcpy(canvas,next,(size_t)W*H*4);rc=0;
done:free(dx);free(dy);free(flags);free(next);free(a);free(b);free(counts);free(deltas);free(sparse);free(full);free(image.rgba);free(residual.rgba);return rc;
}
void rfxa_free(rfxa_animation *a) {size_t i;if(!a)return;free(a->rgba);if(a->audio)for(i=0;i<a->info.count;i++)free(a->audio[i].bytes);free(a->audio);memset(a,0,sizeof(*a));}
int rfxa_decode(const uint8_t *s,size_t n,rfxa_animation *out) {
    uint8_t *owned,*canvas;const uint8_t *raw;size_t size,frame_size,i,table,entry,off,len,audio;
    if(!out)return -1;memset(out,0,sizeof(*out));raw=unwrap(s,n,&owned,&size);
    if(inspect_raw(raw,size,&out->info)){free(owned);return -1;}
    frame_size=(size_t)out->info.width*out->info.height*4;if(out->info.count>LIMIT/frame_size){free(owned);return -1;}
    out->rgba=malloc(frame_size*out->info.count);out->audio=calloc(out->info.count,sizeof(*out->audio));canvas=calloc(frame_size,1);if(!out->rgba||!out->audio||!canvas)goto fail;
    table=u32(raw+28);entry=out->info.flags&1?12:8;
    for(i=0;i<out->info.count;i++) {
        off=u32(raw+table+i*entry);len=u32(raw+table+i*entry+4);audio=entry==12?u32(raw+table+i*entry+8):0;
        if(apply(raw+off,len,canvas,&out->info,(unsigned)i,raw+32,u16(raw+26)))goto fail;memcpy(out->rgba+i*frame_size,canvas,frame_size);
        if(audio){out->audio[i].bytes=malloc(audio);if(!out->audio[i].bytes)goto fail;memcpy(out->audio[i].bytes,raw+off+len,audio);out->audio[i].size=audio;}
    }free(canvas);free(owned);return 0;
fail:free(canvas);free(owned);rfxa_free(out);return -1;
}
int rfxa_encode(const rfxa_animation *a,unsigned quality,int compressed,uint8_t **data,size_t *size) {
    uint8_t **payloads=NULL,*out=NULL,*packed=NULL;size_t *lengths=NULL,total,frame_size,i,offset,packed_size;unsigned audio=0,entry;rfxl_image image;rau_info ai;
    if(!a||!data||!size||!a->rgba||quality>4||!a->info.count||!a->info.fps||!a->info.width||!a->info.height)return -1;*data=NULL;*size=0;
    frame_size=(size_t)a->info.width*a->info.height*4;if(a->info.count>LIMIT/frame_size)return -1;
    if(a->audio)for(i=0;i<a->info.count;i++)if(a->audio[i].size){if(rau_inspect(a->audio[i].bytes,a->audio[i].size,&ai)||ai.sample_rate!=a->info.audio_rate||ai.channels!=a->info.audio_channels||ai.bits!=a->info.audio_bits)return -1;audio=1;}
    entry=audio?12:8;total=32+(size_t)a->info.count*entry;payloads=calloc(a->info.count,sizeof(*payloads));lengths=calloc(a->info.count,sizeof(*lengths));if(!payloads||!lengths)goto fail;
    image.width=a->info.width;image.height=a->info.height;
    for(i=0;i<a->info.count;i++){image.rgba=a->rgba+i*frame_size;if(rfxl_encode(&image,(enum rfxl_quality)(quality==4?3:quality),&payloads[i],&lengths[i]))goto fail;if(lengths[i]+10>LIMIT-total)goto fail;total+=lengths[i]+10;if(audio){if(a->audio[i].size>LIMIT-total)goto fail;total+=a->audio[i].size;}}
    out=calloc(total,1);if(!out)goto fail;memcpy(out,"RFXA",4);out[4]=26;out[5]=(uint8_t)(quality==4?3:quality);out[6]=(uint8_t)audio;put(out+8,a->info.width,4);put(out+12,a->info.height,4);put(out+16,a->info.count,4);put(out+20,a->info.fps,2);put(out+28,32,4);
    if(audio){put(out+22,a->info.audio_rate,2);out[24]=a->info.audio_channels;out[25]=a->info.audio_bits;}
    offset=32+(size_t)a->info.count*entry;
    for(i=0;i<a->info.count;i++){put(out+32+i*entry,offset,4);put(out+36+i*entry,lengths[i]+10,4);if(audio)put(out+40+i*entry,a->audio[i].size,4);put(out+offset+6,a->info.width,2);put(out+offset+8,a->info.height,2);memcpy(out+offset+10,payloads[i],lengths[i]);offset+=lengths[i]+10;if(audio){memcpy(out+offset,a->audio[i].bytes,a->audio[i].size);offset+=a->audio[i].size;}}
    if(compressed){packed=lexa_brief_lz_encode(out,total,&packed_size);if(!packed||packed_size>LIMIT-16)goto fail;free(out);out=calloc(packed_size+16,1);if(!out)goto fail;memcpy(out,"RFXZ",4);out[4]=26;out[5]=14;put(out+8,total,4);memcpy(out+16,packed,packed_size);total=packed_size+16;}
    *data=out;*size=total;out=NULL;
fail:if(payloads)for(i=0;i<a->info.count;i++)free(payloads[i]);free(payloads);free(lengths);free(packed);free(out);return *data?0:-1;
}