#include "atom_server.h"
#include "atom_mb.h"
#include <arpa/inet.h>
#include <errno.h>
#include <fcntl.h>
#include <string.h>
#include <sys/socket.h>
#include <unistd.h>
#define AS_POOL 4
#define AS_PEERS 8
#define AS_ITEMS 2000
#ifndef MSG_NOSIGNAL
#define MSG_NOSIGNAL 0
#endif
typedef struct { int fd;uint8_t rx[260],tx[260];size_t received,needed,sent,length;uint64_t deadline; } as_peer;
typedef struct {
 int used,enabled,suspended,listener,running,error,os_error;uint32_t generation,received,replied,exceptions,retry;
 uint32_t ip;uint16_t port,base[4],count[4],values[4][AS_ITEMS];uint64_t retry_at,heartbeat; uint8_t dirty[4][AS_ITEMS];
 as_peer peers[AS_PEERS];
} as_server;
static as_server as_slots[AS_POOL];
static uint16_t as_get16(const uint8_t *p){return (uint16_t)((p[0]<<8)|p[1]);}
static void as_put16(uint8_t *p,uint16_t v){p[0]=(uint8_t)(v>>8);p[1]=(uint8_t)v;}
static void as_peer_close(as_peer *p){if(p->fd>=0)close(p->fd);p->fd=-1;p->received=0;p->needed=7;p->sent=0;p->length=0;}
static void as_stop(as_server *s){if(s->listener>=0)close(s->listener);s->listener=-1;s->running=0;for(unsigned i=0;i<AS_PEERS;i++)as_peer_close(&s->peers[i]);}
static void as_snapshot(as_server *s,uint32_t h,uint32_t *o){
 memset(o,0,24*sizeof(*o));o[0]=h;o[1]=(uint32_t)s->running;o[2]=(uint32_t)s->error;o[3]=s->error?AM_SOCKET:AM_OK;o[4]=(uint32_t)s->os_error;
 for(unsigned i=0;i<AS_PEERS;i++)if(s->peers[i].fd>=0)o[5]++;
 o[6]=s->received;o[7]=s->replied;o[8]=s->exceptions;
}
static int as_nonblock(int fd){int flags=fcntl(fd,F_GETFL,0);return flags>=0 && fcntl(fd,F_SETFL,flags|O_NONBLOCK)==0;}
static int as_range(as_server *s,unsigned table,uint16_t addr,uint16_t count){return count && addr>=s->base[table] && (uint32_t)addr+count<=(uint32_t)s->base[table]+s->count[table];}
static void as_reply(as_server *s,as_peer *p){
 uint8_t f=p->rx[7],exception=0;uint16_t address=0,q=0;unsigned table=0,index=0;size_t n=9;
 if(p->needed<12){exception=(f==1 || f==2 || f==3 || f==4 || f==5 || f==6 || f==15 || f==16)?3:1;goto finish;}
 address=as_get16(p->rx+8);q=as_get16(p->rx+10);
 switch(f){
 case 1:case 2:case 3:case 4:
  table=f==1?0:f==2?1:f==3?2:3;
  if(p->needed!=12 || !q || q>(f<=2?2000:125)){exception=3;break;}
  if(!as_range(s,table,address,q)){exception=2;break;}
  index=address-s->base[table];p->tx[8]=(uint8_t)(f<=2?(q+7u)/8u:2u*q);n=9+p->tx[8];memset(p->tx+9,0,n-9);
  for(unsigned i=0;i<q;i++){if(f<=2){if(s->values[table][index+i])p->tx[9+i/8]|=(uint8_t)(1u<<(i%8));}else as_put16(p->tx+9+2*i,s->values[table][index+i]);}
  break;
 case 5:case 6:
  table=f==5?0:2;
  if(p->needed!=12 || (f==5 && q!=0 && q!=0xff00)){exception=3;break;}
  if(!as_range(s,table,address,1)){exception=2;break;}
  s->values[table][address-s->base[table]]=f==5?(q!=0):q;s->dirty[table][address-s->base[table]]=1;memcpy(p->tx+8,p->rx+8,4);n=12;break;
 case 15:case 16:
  table=f==15?0:2;
  if(!q || q>(f==15?1968:123) || p->needed<13 || p->rx[12]!=(f==15?(q+7u)/8u:2u*q) || p->needed!=(size_t)(13+p->rx[12])){exception=3;break;}
  if(!as_range(s,table,address,q)){exception=2;break;}
  index=address-s->base[table];
  for(unsigned i=0;i<q;i++){s->values[table][index+i]=f==15?(uint16_t)((p->rx[13+i/8]>>(i%8))&1):as_get16(p->rx+13+2*i);s->dirty[table][index+i]=1;}
  memcpy(p->tx+8,p->rx+8,4);n=12;break;
 default:exception=1;break;
 }
finish:
 s->received++;memcpy(p->tx,p->rx,4);p->tx[6]=p->rx[6];p->tx[7]=f;
 if(exception){p->tx[7]=(uint8_t)(f|0x80);p->tx[8]=exception;n=9;s->exceptions++;}
 as_put16(p->tx+4,(uint16_t)(n-6));p->length=n;p->sent=0;
}
static void as_tick(as_server *s,uint64_t now){
 if(!s->enabled)return;
 if(s->listener<0){
  if(now<s->retry_at)return;
  struct sockaddr_in a;int reuse=1;s->listener=socket(AF_INET,SOCK_STREAM,0);
  if(s->listener<0)goto listener_error;
  setsockopt(s->listener,SOL_SOCKET,SO_REUSEADDR,&reuse,sizeof(reuse));
  if(!as_nonblock(s->listener))goto listener_error;
  memset(&a,0,sizeof(a));a.sin_family=AF_INET;a.sin_port=htons(s->port);a.sin_addr.s_addr=htonl(s->ip);
  if(bind(s->listener,(struct sockaddr *)&a,sizeof(a))<0 || listen(s->listener,AS_PEERS)<0)goto listener_error;
  s->running=1;s->error=0;s->os_error=0;
 }
 /* One accept and one I/O action per peer per scan; fixed work/storage bounds. */
 int fd=accept(s->listener,NULL,NULL);
 if(fd>=0){unsigned i;for(i=0;i<AS_PEERS;i++)if(s->peers[i].fd<0)break;
  if(i==AS_PEERS || !as_nonblock(fd))close(fd);else{s->peers[i].fd=fd;s->peers[i].deadline=now+60000;}}
 else if(errno!=EAGAIN && errno!=EWOULDBLOCK && errno!=EINTR)goto listener_error;
 for(unsigned i=0;i<AS_PEERS;i++){
  as_peer *p=&s->peers[i];if(p->fd<0)continue;
  if(now>=p->deadline){as_peer_close(p);continue;}
  if(p->length){
   ssize_t n=send(p->fd,p->tx+p->sent,p->length-p->sent,MSG_NOSIGNAL);
   if(n>0){p->sent+=(size_t)n;if(p->sent==p->length){s->replied++;p->length=0;p->sent=0;p->received=0;p->needed=7;p->deadline=now+60000;}}
   else if(n==0 || (errno!=EAGAIN && errno!=EWOULDBLOCK && errno!=EINTR))as_peer_close(p);
  }else{
   ssize_t n=recv(p->fd,p->rx+p->received,p->needed-p->received,0);
   if(n>0){if(p->received==0)p->deadline=now+5000;p->received+=(size_t)n;
    if(p->received==7){uint16_t length=as_get16(p->rx+4);if(as_get16(p->rx+2)!=0 || length<2 || length>254){as_peer_close(p);continue;}p->needed=6u+length;}
    if(p->received==p->needed){as_reply(s,p);p->deadline=now+5000;}
   }else if(n==0 || (errno!=EAGAIN && errno!=EWOULDBLOCK && errno!=EINTR))as_peer_close(p);
  }
 }
 return;
listener_error:
 s->os_error=errno;s->error=1;as_stop(s);s->retry_at=now+s->retry;
}
uint32_t as_call(uint32_t cmd,uint32_t h,uint32_t ip,uint16_t port,uint8_t table,uint16_t addr,uint16_t q,uint16_t *data,uint32_t cap,uint32_t *out,uint64_t now,uint32_t retry){
 as_server *s;unsigned i;uint32_t result=0;
 if(cmd==AS_OPEN){
  memset(out,0,24*sizeof(*out));if(!port || !retry)return AM_ARGUMENT;
  for(i=0;i<AS_POOL;i++)if(!as_slots[i].used && as_slots[i].generation<0xffffffu)break;
  if(i==AS_POOL)return 257;
  s=&as_slots[i];uint32_t generation=s->generation+1;memset(s,0,sizeof(*s));s->generation=generation;s->used=1;s->listener=-1;s->ip=ip;s->port=port;s->retry=retry;s->heartbeat=now;
  for(unsigned j=0;j<AS_PEERS;j++){s->peers[j].fd=-1;as_peer_close(&s->peers[j]);}
  h=(generation<<8)|(i+65u);as_snapshot(s,h,out);return 0;
 }
 i=(h&255u);if(i<65 || i>=65+AS_POOL){memset(out,0,24*sizeof(*out));return 256;}i-=64;
 if( !as_slots[i-1].used || as_slots[i-1].generation!=(h>>8)){memset(out,0,24*sizeof(*out));return 256;}
 s=&as_slots[i-1];
 switch(cmd){
 case AS_CLOSE:as_stop(s);s->used=0;as_snapshot(s,0,out);return 0;
 case AS_ENABLE:s->enabled=1;s->suspended=0;s->heartbeat=now;break;
 case AS_DISABLE:s->suspended=0;s->enabled=0;as_stop(s);break;
 case AS_TICK:
  s->heartbeat=now;
  if(s->suspended){s->suspended=0;s->enabled=1;s->retry_at=now;}
  as_tick(s,now);
  break;
 case AS_PUSH:case AS_PULL:
  if(table>3 || !q || q>AS_ITEMS || (uint32_t)addr+q>65536u)result=AM_ARGUMENT;
  else if(!data || cap<q)result=258;
  else if(cmd==AS_PUSH){
   if(table<2){for(unsigned j=0;j<q;j++)if(data[j]>1){result=AM_ARGUMENT;break;}}
   if(!result){
    if(s->count[table] && (s->base[table]!=addr || s->count[table]!=q)) { result=AM_ARGUMENT;break; }
    s->base[table]=addr;s->count[table]=q;
    for(unsigned j=0;j<q;j++)if(!s->dirty[table][j])s->values[table][j]=data[j];
   }
  }else if(s->base[table]!=addr || s->count[table]!=q)result=AM_ARGUMENT;
  else {memcpy(data,s->values[table],q*sizeof(*data));memset(s->dirty[table],0,q);}
  break;
 default:result=259;break;
 }
 as_snapshot(s,h,out);return result;
}
void as_shutdown(void){for(unsigned i=0;i<AS_POOL;i++)if(as_slots[i].used){as_stop(&as_slots[i]);as_slots[i].used=0;}}