from Function import *

def includes():
    return """
#include <map>
#include <string>
#include <vector>
#include "ctx.h"
#include "llm.h"
#include "memory_Context.h"
#include "main.h"
"""

def replaceAnd(s):
    return s.replace(" and", " &&")

def replaceAppend(s):
    return s.replace(".append(", ".push_back(")

def replaceComment(s):
    return s.replace("#", "//")

def replaceLen(s):
    posLen = s.find("len(")
    posEnd = s.find(")", posLen)
    if (
      posLen == -1 or
      posEnd == -1
    ):
      return s
    before = s[:posLen]
    name = s[posLen + len("len("):posEnd]
    after = s[posEnd + len(")"):]
    return f"{before}{name}.size(){after}"

def replaceOr(s):
    return s.replace(" or", " ||")

def replaceTrue(s):
    return s.replace("True", "true")

def translateParameter(s):
    # name: type -> type name
    parts = s.split(": ")
    indent = len(s) - len(s.lstrip())
    name = parts[0].lstrip()
    t = translateType(parts[1])
    indentation = "".join(" " * indent)
    return f"{indentation}{t} {name}"

def translateStatement(s, state):
    indent = len(s) - len(s.lstrip())
    indentation = "".join(" " * indent)
    ss = s.lstrip()
    posColon = ss.find(": ")
    posComma = ss.find(", ")
    posComment = ss.find("# ")
    posCtx = ss.find("c.")
    posEqual = ss.find(" = ")
    posFor = ss.find("for ")
    posIn = ss.find(" in ")
    posRange = ss.find("range(")
    posRangeEnd = ss.find("):")
    posClosingScope = ss.find("#}")
    posOpenSquareBracket = ss.find("[")

    # # -> //
    if posComment != -1:
      return replaceComment(s)

    # #} -> }
    if posClosingScope != -1:
      return f"{indentation}}}"

    # for name in range(x, y): -> for (auto name = x; name < y; ++name) {
    if (
        posFor >= 0 and
        posIn >= 0 and
        posRange >= 0
    ):
      name = ss[posFor + len("for "):posIn]
      x = ss[posRange + len("range("):posComma]
      y = ss[posComma + len(", "):posRangeEnd]
      return f"{indentation}for (auto {name} = {x}; {name} < {y}; ++{name}) {{"

    # name: type = value -> type name = value
    if (
        posColon >= 0 and
        posEqual >= 0
    ):
      name = ss[:posColon]
      type = ss[posColon + len(": "):posEqual]
      t = translateType(type)
      value = ss[posEqual + len(" = "):]
      return f"{indentation}{t} {name} = {value};"

    # name = value -> auto name = value
    if (
        posCtx == -1 and
        posColon == -1 and
        posOpenSquareBracket == -1 and
        posEqual >= 0
    ):
      name = ss[:posEqual]
      # Skip prepending 'auto' each time variable is assigned,
      # only do it the first time
      if name not in state.varNames:
          state.varNames[name] = True
          value = ss[posEqual + len(" = "):]
          return f"{indentation}auto {name} = {value};"

    # Keep "if ("
    if ss == "if (":
        state.isIf = True
        return s

    # Keep "if not ("
    if ss == "if not (":
        state.isIfNot = True
        return f"{indentation}if (!("

    # ): -> }
    if ss == "):":
        # if
        if state.isIf:
            state.isIf = False
            return f"{indentation}) {{"
        # if not
        state.isIfNot = False
        return f"{indentation})) {{"

    ending = ";"
    if state.isIf or state.isIfNot:
        ending = ""

    # Unknown.
    return f"{s}{ending}"

def translateType(s):
    # dict[X, Y] -> std::map<X, Y>
    if s.startswith("dict["):
        kv = s[len("dict["):-len("]")]
        parts = kv.split(", ")
        return f"std::map<{parts[0]}, {parts[1]}>"
    # str -> std::string
    if s == "str":
        return "std::string"

    # Unknown. Return as is.
    return s

class CPP:
    def __init__(self, fn):
        self.fn = fn
        self.isIf = False
        self.isIfNot = False
        self.varNames = {}

    def translate(self):
        returnType = translateType(self.fn.returnType)

        # Parameters.
        params = []
        for i in range(0, len(self.fn.parameters)):
            p = translateParameter(self.fn.parameters[i])
            # Make Context passed by reference.
            #if "Context" in p:
            #    p = p.replace("Context", "Context&")
            params.append(p)
        strparams = "\n".join(params)
        if (len(strparams) > 0):
            strparams += "\n"

        # Statements.
        sts = []
        for i in range(0, len(self.fn.statements)):
          s = translateStatement(self.fn.statements[i], self)
          s = replaceAnd(s)
          s = replaceOr(s)
          s = replaceAppend(s)
          # Replace len twice to account for double invocation.
          s = replaceLen(s)
          s = replaceLen(s)
          s = replaceTrue(s)
          sts.append(s)
        strstatements = "\n".join(sts)

        return f"""{returnType} {self.fn.name}(
{strparams}) {{
{strstatements}
}}
"""