/**
 * @file data_generation/llama/generate_brewery.cpp
 * @brief Builds brewery prompts with regional context, performs retry-based
 * inference, and validates structured JSON output for brewery records.
 */

#include <spdlog/spdlog.h>

#include <array>
#include <stdexcept>
#include <string>

#include "data_generation/llama_generator.h"
#include "data_generation/llama_generator_helpers.h"

namespace {

std::string ExtractFinalJsonPayload(std::string raw_response) {
   auto trim = [](std::string_view text) -> std::string_view {
      const std::size_t first = text.find_first_not_of(" \t\n\r");
      if (first == std::string_view::npos) {
         return {};
      }

      const std::size_t last = text.find_last_not_of(" \t\n\r");
      return text.substr(first, last - first + 1);
   };

   static const std::array<std::string_view, 6> separator_tokens = {
       "<|think|>", "<think|>",   "<|turn|>",
       "<turn|>",   "<channel|>", "<|channel|>"};

   std::size_t separator_pos = std::string::npos;
   std::size_t separator_length = 0;
   for (const std::string_view token : separator_tokens) {
      const std::size_t candidate_pos = raw_response.rfind(token);
      if (candidate_pos != std::string::npos &&
          (separator_pos == std::string::npos ||
           candidate_pos > separator_pos)) {
         separator_pos = candidate_pos;
         separator_length = token.size();
      }
   }

   if (separator_pos != std::string::npos) {
      raw_response.erase(0, separator_pos + separator_length);
   }

   const std::string_view trimmed = trim(raw_response);
   std::string json_candidate =
       ExtractLastJsonObjectPublic(std::string(trimmed));
   if (!json_candidate.empty()) {
      return ExtractLastJsonObjectPublic(std::string(trimmed));
   }

   return std::string(trimmed);
}

}  // namespace

BreweryResult LlamaGenerator::GenerateBrewery(
    const BreweryLocation& location, const std::string& region_context) {
   /**
    * Preprocess and truncate region context to manageable size
    */
   const std::string safe_region_context =
       PrepareRegionContextPublic(region_context);

   /**
    * Load brewery system prompt from file
    * Falls back to minimal inline prompt if file not found
    */
   const std::string system_prompt =
       LoadBrewerySystemPrompt("prompts/system.md");

   /**
    * User prompt: provides geographic context to guide generation towards
    * culturally appropriate and locally-inspired brewery attributes
    */
   std::string prompt =
       "Write a brewery name and place-specific long description for a craft "
       "brewery in ";
   prompt.append(location.city_name);
   if (!location.country_name.empty()) {
      prompt.append(", ");
      prompt.append(location.country_name);
   }
   if (safe_region_context.empty()) {
      prompt.append(".");
   } else {
      prompt.append(". Regional context: ");
      prompt.append(safe_region_context);
   }

   /**
    * Store location context for retry prompts (without repeating full context)
    */
   std::string retry_location = "Location: ";
   retry_location.append(location.city_name);
   if (!location.country_name.empty()) {
      retry_location.append(", ");
      retry_location.append(location.country_name);
   }

   /**
    * RETRY LOOP with validation and error correction
    * Attempts to generate valid brewery data up to 3 times, with feedback-based
    * refinement
    */
   const int max_attempts = 3;
   std::string raw;
   std::string last_error;

   // Limit output length to keep it concise and focused
   constexpr int max_tokens = 1052;
   for (int attempt = 0; attempt < max_attempts; ++attempt) {
      // Generate brewery data from LLM
      raw = Infer(system_prompt, prompt, max_tokens);
      spdlog::debug("LlamaGenerator: raw output (attempt {}): {}", attempt + 1,
                    raw);

      // Validate output: parse JSON and check required fields

      std::string name;
      std::string description;
      const std::string json_only = ExtractFinalJsonPayload(raw);
      const std::string validation_error =
          ValidateBreweryJsonPublic(json_only, name, description);
      if (validation_error.empty()) {
         // Success: return parsed brewery data
         return {std::move(name), std::move(description)};
      }

      // Validation failed: log error and prepare corrective feedback

      last_error = validation_error;
      spdlog::warn("LlamaGenerator: malformed brewery JSON (attempt {}): {}",
                   attempt + 1, validation_error);

      // Update prompt with error details to guide LLM toward correct output.
      // For retries, use a compact prompt format to avoid exceeding token
      // limits.
      prompt =
          "Your previous response was invalid. Error: " + validation_error +
          "\nReturn ONLY valid JSON with exactly these keys: "
          "{\"name\": \"<brewery name>\", "
          "\"description\": \"<single-paragraph description>\"}."
          "\nDo not include markdown, comments, extra keys, or literal "
          "placeholder values.";
      prompt += "\n\n";
      prompt += retry_location;
   }

   // All retry attempts exhausted: log failure and throw exception
   spdlog::error(
       "LlamaGenerator: malformed brewery response after {} attempts: "
       "{}",
       max_attempts, last_error.empty() ? raw : last_error);
   throw std::runtime_error("LlamaGenerator: malformed brewery response");
}