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

  Copyright (C) 2003-2020 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_options.h"
#include "process_int.h"
#include "param_conversion.h"
#include "parse_literals.h"
#include "pmlist.h"
#include "util_wildcards.h"

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

  const int kvn = kvlist.size();
  std::vector<char *> wname(kvn, nullptr);
  for (int i = 0; i < kvn; ++i) wname[i] = nullptr;

  char name[CDI_MAX_NAME];
  char buffer[CDI_MAX_NAME];
  for (const auto &kv : kvlist)
    {
      char *varname = nullptr, *attname = nullptr;
      strcpy(buffer, kv.key.c_str());
      char *const result = strrchr(buffer, delim);
      if (result == nullptr)
        {
          attname = buffer;
        }
      else
        {
          attname = result + 1;
          *result = 0;
          varname = buffer;
        }

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

      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)
            {
              int length = CDI_MAX_NAME;
              /*
              for ( int idx = 0; idx < ngrids; idx++ )
                {
                  int gridID = vlistGrid(vlistID, idx);
                  length = CDI_MAX_NAME;
                  cdiInqKeyString(gridID, CDI_XAXIS, CDI_KEY_NAME, name, &length);
                  if (wildcardmatch(varname, name) == 0)
                    {
                      cdiID = gridID;
                      varIDs[nv++] = CDI_GLOBAL;
                    }
                  length = CDI_MAX_NAME;
                  cdiInqKeyString(gridID, CDI_YAXIS, CDI_KEY_NAME, name, &length);
                  if (wildcardmatch(varname, name) == 0)
                    {
                      cdiID = gridID;
                      varIDs[nv++] = CDI_GLOBAL;
                    }
                }
              */
              for (int idx = 0; idx < nzaxis; idx++)
                {
                  const auto zaxisID = vlistZaxis(vlistID, idx);
                  length = CDI_MAX_NAME;
                  cdiInqKeyString(zaxisID, CDI_GLOBAL, CDI_KEY_NAME, name, &length);
                  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] == nullptr)
                    {
                      wname[i] = strdup(varname);
                      break;
                    }
                  if (cstrIsEqual(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 auto &values = kv.values;
          const auto &value = kv.values[0];
          int nvalues = kv.nvalues;
          if (nvalues == 1 && value.empty()) nvalues = 0;

          const auto dtype = literalsFindDatatype(nvalues, values);

          for (int idx = 0; idx < nv; ++idx)
            {
              const auto varID = varIDs[idx];
              // if ( Options::cdoVerbose ) printf("varID, cdiID, attname %d %d %s %d\n", varID, cdiID, attname,
              // (int)strlen(attname));
              if (nvalues == 0)
                {
                  cdiDelAtt(cdiID, varID, attname);
                }
              else
                {
                  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] = literal2int(values[i]);
                      cdiDefAttInt(cdiID, varID, attname, dtype, nvalues, ivals.data());
                    }
                  else if (dtype == CDI_DATATYPE_FLT32 || dtype == CDI_DATATYPE_FLT64)
                    {
                      Varray<double> dvals(nvalues);
                      for (int i = 0; i < nvalues; ++i) dvals[i] = literal2double(values[i]);
                      cdiDefAttFlt(cdiID, varID, attname, dtype, nvalues, dvals.data());
                    }
                  else
                    {
                      if (nvalues > 1) cdoAbort("Multidimensional string attributes not supported! %s=\"%s\"", attname, values[1].c_str());
                      const auto len = (int) value.size();
                      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");

  const auto lcopy = unchangedRecord();

  const auto operatorID = cdoOperatorID();

  operatorInputArg(cdoOperatorEnter(operatorID));

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

  PMList pmlist;
  KVList kvlist;
  kvlist.name = "SETATTRIBUTE";
  if (kvlist.parseArguments(natts, cdoGetOperArgv()) != 0) cdoAbort("Parse error!");
  if (Options::cdoVerbose) kvlist.print();

  auto pkvlist = &kvlist;
  if (natts == 1)
    {
      KeyValues &kv = kvlist.front();
      if (kv.key == "FILE")
        {
          if (Options::cdoVerbose) cdoPrint("Reading attributes from: %s", kv.values[0].c_str());
          auto filename = parameter2word(kv.values[0].c_str());
          auto fp = fopen(filename, "r");
          if (fp == nullptr) cdoAbort("Open failed on: %s\n", filename);
          pmlist.readNamelist(fp, filename);
          pkvlist = &pmlist.front();
          if (pkvlist == nullptr) cdoAbort("Parse error!");
          fclose(fp);
          if (Options::cdoVerbose) pkvlist->print();
        }
    }

  const auto streamID1 = cdoOpenRead(0);

  const auto vlistID1 = cdoStreamInqVlist(streamID1);
  const auto vlistID2 = vlistDuplicate(vlistID1);

  setAttributes(*pkvlist, vlistID2);

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

  const auto streamID2 = cdoOpenWrite(1);
  cdoDefVlist(streamID2, vlistID2);

  Varray<double> array;
  if (!lcopy)
    {
      auto gridsizemax = vlistGridsizeMax(vlistID1);
      if (vlistNumber(vlistID1) != CDI_REAL) gridsizemax *= 2;
      array.resize(gridsizemax);
    }

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

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

          if (lcopy)
            {
              cdoCopyRecord(streamID2, streamID1);
            }
          else
            {
              size_t nmiss;
              cdoReadRecord(streamID1, array.data(), &nmiss);
              cdoWriteRecord(streamID2, array.data(), nmiss);
            }
        }

      tsID++;
    }

  cdoStreamClose(streamID1);
  cdoStreamClose(streamID2);

  cdoFinish();

  return nullptr;
}
