/*
  This file is part of CDO. CDO is a collection of Operators to
  manipulate and analyse Climate model Data.

  Copyright (C) 2003-2018 Uwe Schulzweida, <uwe.schulzweida AT mpimet.mpg.de>
  See COPYING file for copying and redistribution conditions.

  This program is free software; you can redistribute it and/or modify
  it under the terms of the GNU General Public License as published by
  the Free Software Foundation; version 2 of the License.

  This program is distributed in the hope that it will be useful,
  but WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  GNU General Public License for more details.
*/

#include <cdi.h>

#include "cdo_int.h"
#include "pstream_int.h"

static void
setAttributes(list_t *kvlist, int vlistID)
{
  enum
  {
    Undefined = -99
  };
  const int delim = '@';
  int nvars = vlistNvars(vlistID);
  int ngrids = vlistNgrids(vlistID);
  int nzaxis = vlistNzaxis(vlistID);
  int maxvars = nvars + ngrids * 2 + nzaxis;
  std::vector<int> varIDs(maxvars);

  int kvn = list_size(kvlist);
  std::vector<char *> wname(kvn, NULL);
  for (int i = 0; i < kvn; ++i) wname[i] = NULL;

  char name[CDI_MAX_NAME];
  char buffer[CDI_MAX_NAME];
  for (listNode_t *kvnode = kvlist->head; kvnode; kvnode = kvnode->next)
    {
      char *varname = NULL, *attname = NULL;
      keyValues_t *kv = *(keyValues_t **) kvnode->data;
      strcpy(buffer, kv->key);
      char *result = strrchr(buffer, delim);
      if (result == NULL)
        {
          attname = buffer;
        }
      else
        {
          attname = result + 1;
          *result = 0;
          varname = buffer;
        }

      if (*attname == 0) cdoAbort("Attribute name missing in >%s<!", kv->key);

      int nv = 0;
      int cdiID = Undefined;
      if (varname && *varname)
        {
          for (int idx = 0; idx < nvars; idx++)
            {
              vlistInqVarName(vlistID, idx, name);
              if (wildcardmatch(varname, name) == 0)
                {
                  cdiID = vlistID;
                  varIDs[nv++] = idx;
                }
            }

          if (cdiID == Undefined)
            {
              /*
              for ( int idx = 0; idx < ngrids; idx++ )
                {
                  int gridID = vlistGrid(vlistID, idx);
                  gridInqXname(gridID, name);
                  if ( wildcardmatch(varname, name) == 0 )
                    {
                      cdiID = gridID;
                      varIDs[nv++] = CDI_GLOBAL;
                    }
                  gridInqYname(gridID, name);
                  if ( wildcardmatch(varname, name) == 0 )
                    {
                      cdiID = gridID;
                      varIDs[nv++] = CDI_GLOBAL;
                    }
                }
              */
              for (int idx = 0; idx < nzaxis; idx++)
                {
                  int zaxisID = vlistZaxis(vlistID, idx);
                  zaxisInqName(zaxisID, name);
                  if (wildcardmatch(varname, name) == 0)
                    {
                      cdiID = zaxisID;
                      varIDs[nv++] = CDI_GLOBAL;
                    }
                }
            }

          if (cdiID == Undefined)
            {
              bool lwarn = true;
              for (int i = 0; i < kvn; ++i)
                {
                  if (wname[i] == NULL)
                    {
                      wname[i] = strdup(varname);
                      break;
                    }
                  if (STR_IS_EQ(wname[i], varname))
                    {
                      lwarn = false;
                      break;
                    }
                }
              if (lwarn)
                {
                  cdoWarning("Variable >%s< not found!", varname);
                }
            }
        }
      else
        {
          cdiID = vlistID;
          varIDs[nv++] = CDI_GLOBAL;
        }

      if (cdiID != Undefined && nv > 0)
        {
          const char *value = (kv->nvalues > 0) ? kv->values[0] : NULL;
          int nvalues = kv->nvalues;
          if (nvalues == 1 && !*value) nvalues = 0;
          int dtype = literals_find_datatype(nvalues, kv->values);

          for (int idx = 0; idx < nv; ++idx)
            {
              int varID = varIDs[idx];
              // if ( cdoVerbose ) printf("varID, cdiID, attname %d %d %s %d\n",
              // varID, cdiID, attname, (int)strlen(attname));
              if (dtype == CDI_DATATYPE_INT8 || dtype == CDI_DATATYPE_INT16 || dtype == CDI_DATATYPE_INT32)
                {
                  std::vector<int> ivals(nvalues);
                  for (int i = 0; i < nvalues; ++i) ivals[i] = literal_to_int(kv->values[i]);
                  cdiDefAttInt(cdiID, varID, attname, dtype, nvalues, ivals.data());
                }
              else if (dtype == CDI_DATATYPE_FLT32 || dtype == CDI_DATATYPE_FLT64)
                {
                  std::vector<double> dvals(nvalues);
                  for (int i = 0; i < nvalues; ++i) dvals[i] = literal_to_double(kv->values[i]);
                  cdiDefAttFlt(cdiID, varID, attname, dtype, nvalues, dvals.data());
                }
              else
                {
                  int len = (value && *value) ? (int) strlen(value) : 0;
                  int outlen = 0;
                  std::vector<char> outvalue(len);
                  for (int i = 0; i < len; ++i)
                    {
                      if (i > 0 && value[i - 1] == '\\' && value[i] == 'n')
                        outvalue[outlen - 1] = '\n';
                      else
                        outvalue[outlen++] = value[i];
                    }
                  cdiDefAttTxt(cdiID, varID, attname, outlen, outvalue.data());
                }
            }
        }
    }

  for (int i = 0; i < kvn; ++i)
    if (wname[i]) free(wname[i]);
}

void *
Setattribute(void *process)
{
  int nrecs;
  int varID, levelID;

  cdoInitialize(process);

  cdoOperatorAdd("setattribute", 0, 0, "attributes");

  bool lcopy = UNCHANGED_RECORD;

  int operatorID = cdoOperatorID();

  operatorInputArg(cdoOperatorEnter(operatorID));

  int natts = operatorArgc();
  if (natts == 0) cdoAbort("Parameter missing!");

  list_t *pmlist = NULL;
  list_t *kvlist = kvlist_new("SETATTRIBUTES");
  if (kvlist_parse_cmdline(kvlist, natts, operatorArgv()) != 0) cdoAbort("Parse error!");
  if (cdoVerbose) kvlist_print(kvlist);

  if (natts == 1)
    {
      keyValues_t *kv = *(keyValues_t **) kvlist->head->data;
      if (STR_IS_EQ(kv->key, "FILE"))
        {
          if (cdoVerbose) cdoPrint("Reading attributes from: %s", kv->values[0]);
          const char *filename = parameter2word(kv->values[0]);
          FILE *fp = fopen(filename, "r");
          if (fp == NULL) cdoAbort("Open failed on: %s\n", filename);

          pmlist = namelist_to_pmlist(fp, filename);
          if (pmlist == NULL) cdoAbort("Parse error!");
          list_destroy(kvlist);
          kvlist = *(list_t **) pmlist->head->data;
          if (kvlist == NULL) cdoAbort("Parse error!");
          fclose(fp);
          if (cdoVerbose) kvlist_print(kvlist);
        }
    }

  int streamID1 = cdoStreamOpenRead(cdoStreamName(0));

  int vlistID1 = cdoStreamInqVlist(streamID1);
  int vlistID2 = vlistDuplicate(vlistID1);

  setAttributes(kvlist, vlistID2);

  if (pmlist)
    list_destroy(pmlist);
  else
    list_destroy(kvlist);

  int taxisID1 = vlistInqTaxis(vlistID1);
  int taxisID2 = taxisDuplicate(taxisID1);
  vlistDefTaxis(vlistID2, taxisID2);

  int streamID2 = cdoStreamOpenWrite(cdoStreamName(1), cdoFiletype());
  pstreamDefVlist(streamID2, vlistID2);

  std::vector<double> array;
  if (!lcopy)
    {
      size_t gridsize = (size_t) vlistGridsizeMax(vlistID1);
      if (vlistNumber(vlistID1) != CDI_REAL) gridsize *= 2;
      array.resize(gridsize);
    }

  int tsID = 0;
  while ((nrecs = cdoStreamInqTimestep(streamID1, tsID)))
    {
      taxisCopyTimestep(taxisID2, taxisID1);
      pstreamDefTimestep(streamID2, tsID);

      for (int recID = 0; recID < nrecs; recID++)
        {
          pstreamInqRecord(streamID1, &varID, &levelID);
          pstreamDefRecord(streamID2, varID, levelID);

          if (lcopy)
            {
              pstreamCopyRecord(streamID2, streamID1);
            }
          else
            {
              size_t nmiss;
              pstreamReadRecord(streamID1, array.data(), &nmiss);
              pstreamWriteRecord(streamID2, array.data(), nmiss);
            }
        }

      tsID++;
    }

  pstreamClose(streamID1);
  pstreamClose(streamID2);

  cdoFinish();

  return 0;
}
