#include "runtime.h"
#include "json_numbers.h"
#include <sundials/sundials_config.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>

static void json_string(FILE *f, const char *value) {
    fputc('"',f);
    for (const unsigned char *p=(const unsigned char *)value; *p; p++) {
        if (*p=='"' || *p=='\\') { fputc('\\',f); fputc(*p,f); }
        else if (*p<32) fprintf(f,"\\u%04x",*p);
        else fputc(*p,f);
    }
    fputc('"',f);
}
static void vector(FILE *f, const double *value, int count) {
    fputc('[',f);
    for (int i=0;i<count;i++) fprintf(f,"%s%.17g",i?",":"",value[i]);
    fputc(']',f);
}
static int probe(void) {
    double t, y[NSTATES], dy[NSTATES], w[NOUTPUTS];
    while (scanf("%lf",&t)==1) {
        for (int i=0;i<NSTATES;i++) if (scanf("%lf",&y[i])!=1) return 64;
        int ok=model_eval(t,y,dy,w);
        printf("{\"success\":%s",ok?"true":"false");
        if (ok) { printf(",\"rhs\":"); vector(stdout,dy,NSTATES); printf(",\"outputs\":"); vector(stdout,w,NOUTPUTS); }
        printf("}\n");
    }
    return 0;
}
static int write_result_index(const char *path, long series_start, long series_end,
                              long result_bytes, size_t sample_count) {
    if(series_start<0 || series_end<series_start || result_bytes<series_end) return 0;
    FILE *index=fopen(path,"wb"); if(!index) return 0;
    int ok=fprintf(index,"{\"version\":1,\"seriesStart\":%ld,\"seriesEnd\":%ld,"
                         "\"resultBytes\":%ld,\"sampleCount\":%zu}\n",
                   series_start,series_end,result_bytes,sample_count)>=0;
    if(ferror(index)) ok=0;
    if(fclose(index)) ok=0;
    return ok;
}
static int write_result(NativeRun *r, const char *path, const char *index_path) {
    if(index_path && !strcmp(path,index_path)) return 0;
    double dy[NSTATES], final[NOUTPUTS];
    int final_ok=model_eval(r->final_time,r->final_state,dy,final);
    size_t length=r->count*NOUTPUTS;
    if (length>268435456u/sizeof(double)) return 0;
    double *values=length?malloc(length*sizeof(double)):NULL;
    if (length && !values) return 0;
    for (size_t i=0;i<r->count;i++) {
        if (!model_eval(r->times[i],r->states+i*NSTATES,dy,values+i*NOUTPUTS)) {
            free(values); return 0;
        }
    }
    FILE *f=fopen(path,"wb"); if (!f) { free(values); return 0; }
    fprintf(f,"{\"success\":%s,\"status\":",!r->status?"true":"false");
    json_string(f,r->status==0?"completed":r->status==1?"cancelled":"failed");
    fprintf(f,",\"message\":"); json_string(f,r->message);
    fprintf(f,",\"backend\":\"native-c\",\"method\":\"%s\",\"solver\":\"%s\",\"sundialsVersion\":\"%s\","
              "\"simulatedUntil\":%.17g,\"solveSeconds\":%.17g,\"solveCpuSeconds\":%.17g,"
              "\"nfev\":%lu,\"acceptedSteps\":%lu,\"rejectedSteps\":%lu,\"stateTransitions\":%lu,"
              "\"solverStarts\":%lu,\"njev\":%lu,\"nlu\":%lu,\"maxAcceptedStep\":%.17g,\"series\":{",
            r->options.bdf?"BDF":"RK45",r->options.bdf?"CVODE":"Dormand-Prince 5(4)",SUNDIALS_VERSION,
            r->final_time,r->solve_seconds,r->solve_cpu_seconds,r->nfev,r->accepted,r->rejected,
            r->events,r->starts,r->njev,r->nlu,r->max_accepted_step);
    /* Binary-mode positions delimit the complete series object, including
       both braces. The optional index avoids scanning or parsing its values. */
    long series_start=-1,series_end=-1,result_bytes=-1;
    if(index_path) {
        series_start=ftell(f);
        if(series_start>0) series_start--; else series_start=-1;
    }
    int output_ok=1;
    if (r->count) {
        fprintf(f,"\"time\":");
        output_ok=native_json_write_array(f,r->times,r->count,1);
        for (int j=0;j<NOUTPUTS && output_ok;j++) {
            fputc(',',f); json_string(f,model_output_keys[j]); fputc(':',f);
            output_ok=native_json_write_array(f,values+j,r->count,NOUTPUTS);
        }
    }
    fputc('}',f);
    if(index_path) series_end=ftell(f);
    fprintf(f,",\"final\":{");
    if (final_ok) for (int j=0;j<NOUTPUTS && output_ok;j++) {
        if (j) fputc(',',f);
        json_string(f,model_output_keys[j]); fputc(':',f);
        output_ok=native_json_write_number(f,final[j]);
    }
    fprintf(f,"},\"finalState\":");
    if (output_ok) output_ok=native_json_write_array(f,r->final_state,NSTATES,1);
    fprintf(f,"}\n"); int ok=output_ok && !ferror(f);
    if(index_path) {
        result_bytes=ftell(f);
        if(series_start<0 || series_end<0 || result_bytes<0) ok=0;
    }
    if (fclose(f)) ok=0;
    /* Publish the index only after the complete result was successfully
       flushed and closed. An index I/O failure is a failed result write. */
    if(ok && index_path) ok=write_result_index(index_path,series_start,series_end,result_bytes,r->count);
    free(values); return ok;
}

int main(int argc, char **argv) {
    NativeRun r={0};
    r.options=(NativeOptions){0,10,.02,.001,1e-6,300,0,1,NULL};
    const char *output="result.json", *index_path=NULL;
    for (int i=1;i<argc;i++) {
        const char *arg=argv[i];
        if (!strcmp(arg,"--probe")) return probe();
        if (!strcmp(arg,"--init")) {
            double y[NSTATES]; if (!model_init(y)) return 2;
            vector(stdout,y,NSTATES); fputc('\n',stdout); return 0;
        }
        if (!strcmp(arg,"--solve-only")) { r.options.record_samples=0; continue; }
        if (i+1==argc) return 64;
        const char *value=argv[++i];
        if (!strcmp(arg,"--output")) output=value;
        else if (!strcmp(arg,"--result-index")) index_path=value;
        else if (!strcmp(arg,"--cancel-file")) r.options.cancel_path=value;
        else if (!strcmp(arg,"--method")) {
            if (strcmp(value,"RK45") && strcmp(value,"BDF")) return 64;
            r.options.bdf=!strcmp(value,"BDF");
        } else {
            char *end; double number=strtod(value,&end);
            if (*end || !isfinite(number)) return 64;
            if (!strcmp(arg,"--start")) r.options.start=number;
            else if (!strcmp(arg,"--stop")) r.options.stop=number;
            else if (!strcmp(arg,"--sample-step")) r.options.sample_step=number;
            else if (!strcmp(arg,"--max-step")) r.options.max_step=number;
            else if (!strcmp(arg,"--rtol")) r.options.rtol=number;
            else if (!strcmp(arg,"--timeout")) r.options.timeout=number;
            else return 64;
        }
    }
    if (!(r.options.stop>r.options.start && r.options.sample_step>0 && r.options.max_step>0 &&
          r.options.rtol>0 && r.options.timeout>0) ||
        r.options.start+r.options.sample_step==r.options.start ||
        r.options.start+r.options.max_step==r.options.start ||
        !isfinite(r.options.stop-r.options.start)) return 64;
    if (r.options.record_samples &&
        ((r.options.stop-r.options.start)/r.options.sample_step>1000000 ||
         ((r.options.stop-r.options.start)/r.options.sample_step+1024)*(NSTATES+NOUTPUTS)*sizeof(double)>268435456)) return 64;
    native_solve(&r);
    int saved=write_result(&r,output,index_path);
    int code=saved?(r.status==2?2:0):3;
    native_run_free(&r); return code;
}