// Copyright (c) FIRST and other WPILib contributors.
// Open Source Software; you can modify and/or share it under the terms of
// the WPILib BSD license file in the root directory of this project.

package org.wpilib.epilogue.processor;

import java.io.IOException;
import java.io.PrintWriter;
import java.util.Collection;
import java.util.List;
import java.util.Map;
import javax.annotation.processing.ProcessingEnvironment;
import javax.lang.model.element.TypeElement;
import javax.lang.model.type.DeclaredType;
import javax.lang.model.type.TypeMirror;
import org.wpilib.epilogue.EpilogueConfiguration;

/**
 * Generates the {@code Epilogue} file used as the main entry point to logging with Epilogue in a
 * robot program. {@code Epilogue} has instances of every generated logger class, a {@link
 * EpilogueConfiguration config} object, and (if the main robot class inherits from {@link
 * org.wpilib.framework.TimedRobot TimedRobot}) a {@code bind()} method to automatically add a
 * periodic logging call to the robot.
 */
public class EpilogueGenerator {
  private final ProcessingEnvironment m_processingEnv;
  private final Map<TypeMirror, DeclaredType> m_customLoggers;

  public EpilogueGenerator(
      ProcessingEnvironment processingEnv, Map<TypeMirror, DeclaredType> customLoggers) {
    this.m_processingEnv = processingEnv;
    this.m_customLoggers = customLoggers;
  }

  /**
   * Creates the Epilogue file, which is the main entry point for users to set up and interact with
   * the generated loggers.
   *
   * @param loggerClassNames the names of the generated logger classes. Each of these will be
   *     instantiated in a public static field on the Epilogue class.
   * @param mainRobotClasses the main robot classes. May be empty. Used to generate a {@code
   *     update()} method for an easy entry point.
   * @param timedRobotClasses the main robot classes that extend from TimedRobot. Used to generate a
   *     {@code bind()} method to add a callback hook to a TimedRobot to log itself.
   */
  @SuppressWarnings("checkstyle:LineLength") // Source code templates exceed the line length limit
  public void writeEpilogueFile(
      List<String> loggerClassNames,
      Collection<TypeElement> mainRobotClasses,
      Collection<TypeElement> timedRobotClasses) {
    try {
      var centralStore =
          m_processingEnv.getFiler().createSourceFile("org.wpilib.epilogue.Epilogue");

      try (var out = new PrintWriter(centralStore.openOutputStream())) {
        out.println("package org.wpilib.epilogue;");
        out.println();

        out.println("import static org.wpilib.units.Units.Seconds;");
        out.println();

        out.println("import org.wpilib.hardware.hal.HAL;");
        out.println();

        out.println("public final class Epilogue {");

        // Usage reporting
        out.println(
            """
              static {
                HAL.reportUsage("Epilogue", "");
              }
            """);

        out.println(
            "  private static final EpilogueConfiguration config = new EpilogueConfiguration();");
        out.println();

        loggerClassNames.forEach(
            clazz -> {
              // public static final com.example.FooLogger com_example_fooLogger =
              //   new com.example.FooLogger();
              String field = clazz.replace('.', '_');
              out.printf("  public static final %s %s = new %s();%n", clazz, field, clazz);
            });
        m_customLoggers.values().stream()
            .distinct()
            .forEach(
                loggerType -> {
                  var loggerTypeName = loggerType.toString();
                  out.printf(
                      "  public static final %s %s = new %s();%n",
                      loggerTypeName, loggerTypeName.replace('.', '_'), loggerTypeName);
                });
        out.println();

        out.println(
            """
              public static void configure(java.util.function.Consumer<EpilogueConfiguration> configurator) {
                configurator.accept(config);
              }

              public static EpilogueConfiguration getConfig() {
                return config;
              }
            """);

        out.println(
            """
              /**
               * Checks if data associated with a given importance level should be logged.
               */
              public static boolean shouldLog(Logged.Importance importance) {
                return importance.compareTo(config.minimumImportance) >= 0;
              }
            """
                .stripTrailing());

        for (TypeElement mainRobotClass : mainRobotClasses) {
          String robotClassName = mainRobotClass.getQualifiedName().toString();

          out.println();
          out.print(
              """
                /**
                 * Updates Epilogue. This must be called periodically in order for Epilogue to record
                 * new values. Alternatively, {@code bind()} can be used to update at an offset from
                 * the main robot loop.
                 */
              """);
          out.println("  public static void update(" + robotClassName + " robot) {");
          out.println("    long start = System.nanoTime();");
          out.println(
              "    "
                  + StringUtils.loggerFieldName(mainRobotClass)
                  + ".tryUpdate(config.backend.getNested(config.root), robot, config.errorHandler);");
          out.println(
              "    config.backend.log(\"Epilogue/Stats/Last Run\", (System.nanoTime() - start) / 1e6);");
          out.println("  }");
        }

        for (TypeElement timedRobotClass : timedRobotClasses) {
          String robotClassName = timedRobotClass.getQualifiedName().toString();

          out.println();
          out.print(
              """
                /**
                 * Binds Epilogue updates to a timed robot's update period. Log calls will be made at the
                 * same update rate as the robot's loop function, but will be offset by a full phase
                 * (for example, a 20ms update rate but 10ms offset from the main loop invocation) to
                 * help avoid high CPU loads. However, this does mean that any logged data that reads
                 * directly from sensors will be slightly different from data used in the main robot
                 * loop.
                 */
              """);
          out.println("  public static void bind(" + robotClassName + " robot) {");
          out.println("    if (config.loggingPeriod == null) {");
          out.println("      config.loggingPeriod = Seconds.of(robot.getPeriod());");
          out.println("    }");
          out.println("    if (config.loggingPeriodOffset == null) {");
          out.println("      config.loggingPeriodOffset = config.loggingPeriod.div(2);");
          out.println("    }");
          out.println();
          out.println("    robot.addPeriodic(() -> {");
          out.println("      update(robot);");
          out.println(
              "    }, config.loggingPeriod.in(Seconds), config.loggingPeriodOffset.in(Seconds));");
          out.println("  }");
        }

        out.println("}");
      }
    } catch (IOException e) {
      throw new RuntimeException(e);
    }
  }
}
