using AiAgent.Core.Infrastructure;
using Mscc.GenerativeAI;
using System.Text.Json;
using static AiAgent.Core.Constant.ProgramConstants;
using static AiAgent.Core.Helpers.GenericHelper;

namespace AiAgent.Core.Agents.Base;

public class BaseAgent()
{
    /// <summary>
    /// Gets the default configuration for a generative model.
    /// </summary>
    public static GenerativeModel GetBaseModel => new()
    {
        Model = GeminiFlashLatest, //GeminiPro25,
        UseGoogleSearch = true,
        UseGrounding = false,
        UseCodeExecution = false,
        UseJsonMode = false,
        Timeout = TimeSpan.FromMinutes(60)
    };

    /// <summary>
    /// Gets the default base configuration for generation settings.
    /// </summary>
    public static GenerationConfig GetBaseConfig => new()
    {
        Temperature = 0.0f,
        CandidateCount = 1,
        ThinkingConfig = new ThinkingConfig
        {
            ThinkingBudget = -1, //24576, //32000,
            IncludeThoughts = true,
        }, 
        StopSequences= ["\r\n\\\r\n\\\r\n\\\r\n\\"]

    };

    /// <summary>
    /// Validates the existence of required files and determines whether the operation can proceed.
    /// </summary>
    /// <remarks>This method performs a series of checks to ensure that the necessary files are in the correct
    /// state before proceeding with the operation.  It logs messages to indicate the reason for any failure or skipped
    /// execution.</remarks>
    /// <param name="name">The name of the operation, used for logging purposes.</param>
    /// <param name="briefFile">The path to the prompt file. The operation will only proceed if this file exists.</param>
    /// <param name="outputFile">The path to the output file. The operation will only proceed if this file does not already exist.</param>
    /// <param name="supportingDocuments">An optional list of paths to supporting input files. The operation will only proceed if all specified files
    /// exist.</param>
    /// <returns>A <see cref="Results"/> value indicating the outcome of the validation: <list type="bullet">
    /// <item><description><see cref="Results.Failed"/> if the prompt file or any required supporting file does not
    /// exist.</description></item> <item><description><see cref="Results.NotExecuted"/> if the output file already
    /// exists.</description></item> <item><description><see cref="Results.Success"/> if all required conditions are met
    /// and the operation can proceed.</description></item> </list></returns>
    public static Results CheckDocuments(string name, string? briefFile, string? outputFile, List<string>? supportingDocuments = null)
    {
        Log($"\n\n===========================================================================");
        Log($"Starting the {name}");
        Log($"---------------------------------\n");

        // Check the prompt, we only work if the prompt exists
        if (briefFile != null)
            if (File.Exists(briefFile) is false)
            {
                Log($"- Not doing anything as the input file does not exist: {briefFile}");
                return Results.Failed;
            }

        // Check the output, we only work if the output does not exist
        if (outputFile != null)
            if (File.Exists(outputFile))
            {
                Log($"- Not doing anything as the output file already exists: {outputFile}");
                return Results.NotExecuted;
            }

        // We only work if the supporting documents exist
        if (supportingDocuments != null)
            foreach (var file in supportingDocuments)
                if (File.Exists(file) is false)
                {
                    Log($"- Not doing anything as the supporting document does not exist: {file}");
                    return Results.Failed;
                }

        // Now we can work
        return Results.Success;
    }

    /// <summary>
    /// Executes a generative model query using the specified configuration, input prompt, and optional input data
    /// files,  and writes the generated output to a file.
    /// </summary>
    /// <remarks>This method attempts to execute the query using multiple API keys in sequence. If all
    /// attempts fail, the method  returns <see cref="Results.Failed"/>. Any errors encountered during execution are
    /// logged.</remarks>
    /// <param name="model">The generative model to use for the query. The model's API key will be set during execution.</param>
    /// <param name="confi">The configuration settings for the generation process.</param>
    /// <param name="name">The name associated with the query, used for identification purposes.</param>
    /// <param name="promptFile">The path to the file containing the input prompt for the generative model.</param>
    /// <param name="outputFile">The path to the file where the generated output will be written.</param>
    /// <param name="inputDataFiles">An optional list of file paths containing additional input data for the query. Can be <see langword="null"/>.</param>
    /// <returns>A <see cref="Results"/> value indicating the outcome of the operation.  Returns <see cref="Results.Success"/> if
    /// the query completes successfully and the output is written to the file;  otherwise, returns <see
    /// cref="Results.Failed"/>.</returns>
    public static async Task<Results> RunAsync(GenerativeModel model, GenerationConfig confi, string apiKey, string name, string promptFile, string outputFile, List<string>? inputDataFiles = null, string[]? jsonMandatoryFields = null)
    {
        // Make sure that everything is in place in order to do some work
        var results = CheckDocuments(name, promptFile, outputFile, inputDataFiles);

        // Anything but success is necessary
        if (results != Results.Success)
            return results;

        for (int tries = 0; tries < 5; tries++)
        {
            // Get the model core name
            Log($"- Starting Gemini Query with model: {model.Model}\n");

            // Get the prompt
            string prompt = await File.ReadAllTextAsync(promptFile);

            // Query
            var response = await GeminiApiClient.Query(prompt, name, model, confi, apiKey, inputDataFiles);

            // Check for failure
            if (response.Contains("\"status\": \"FAILURE\""))
            {
                Log($"-> Gemini Query failed:\n\n{response}");
                if (response.Contains("\"reason_code\": \"RESEARCH_TOOL_UNAVAILABLE\""))
                {
                    await Task.Delay(5000);
                    continue;
                }
            }

            // Check for json tags
            if (jsonMandatoryFields is not null)
            {
                // Isolate the JSON part
                if (response.Contains("```json"))
                    response = response.Split("```json", StringSplitOptions.TrimEntries).Last();

                if (response.Contains("```"))
                    response = response.Split("```", StringSplitOptions.TrimEntries).First();
                

                try
                {
                    var options = new JsonDocumentOptions
                    {
                        AllowTrailingCommas = true,
                        CommentHandling = JsonCommentHandling.Skip
                    };

                    JsonDocument doc = JsonDocument.Parse(response, options);
                    JsonElement root = doc.RootElement;

                    // Parse the JSON to make sure it is valid and contains the mandatory fields
                    foreach (string field in jsonMandatoryFields)
                        if (!root.TryGetProperty(field, out _))
                            throw new Exception($"Mandatory field '{field}' not found in the JSON response.");

                    Log($"\n\nThe response is a valid JSON with all mandatory fields: {string.Join(", ", jsonMandatoryFields)}");
                }
                catch (Exception ex)
                {
                    Log($"- JSON Parsing failed: {ex.Message}");
                    await Task.Delay(120000);
                    continue;
                }
            }

            // Write the output
            if (response.Length > 0)
            {
                File.WriteAllText(outputFile, response);
                return Results.Success;
            }

            Log($"- Gemini Query attempt {tries + 1} failed, trying again...");
            await Task.Delay(2000);
        }

        throw new Exception("All attempts to run the Gemini Query have failed.");
        return Results.Failed;
    }

    /// <summary>
    /// Executes the average worker process using the specified input and output files, and an optional list of
    /// additional data files.
    /// </summary>
    /// <remarks>This method validates the provided files and configuration before executing the worker
    /// process. If the validation fails, the method returns the corresponding error result without performing the
    /// operation.</remarks>
    /// <param name="name">The name of the worker process. This is used to identify the operation.</param>
    /// <param name="promptFile">The path to the prompt file that contains the input instructions for the worker.</param>
    /// <param name="outputFile">The path to the output file where the results of the operation will be written.</param>
    /// <param name="inputDataFiles">An optional list of paths to additional input data files required for the operation. Can be <see
    /// langword="null"/> if no additional files are needed.</param>
    /// <returns>A <see cref="Results"/> value indicating the outcome of the operation. Returns <see cref="Results.Success"/> if
    /// the operation completes successfully; otherwise, returns an appropriate error code.</returns>
    public static async Task<Results> RunAverageWorker(string name, string promptFile, string outputFile, List<string>? inputDataFiles = null, bool jsonOutput = true)
    {
        // Make sure that everything is in place in order to do some work
        var results = CheckDocuments(name, promptFile, outputFile, inputDataFiles);

        // Anything but success is necessary
        if (results != Results.Success)
            return results;

        // Launch the average worker with the default average base configuration
        var response = await RunAsync(GetBaseModel, GetBaseConfig, ApiKeys.GetNextFreeKey(), name, promptFile, outputFile, inputDataFiles);

        return response;
    }

    /// <summary>
    /// Executes the average worker process using the specified input and output files, and an optional list of
    /// additional data files.
    /// </summary>
    /// <remarks>This method validates the provided files and configuration before executing the worker
    /// process. If the validation fails, the method returns the corresponding error result without performing the
    /// operation.</remarks>
    /// <param name="name">The name of the worker process. This is used to identify the operation.</param>
    /// <param name="promptFile">The path to the prompt file that contains the input instructions for the worker.</param>
    /// <param name="outputFile">The path to the output file where the results of the operation will be written.</param>
    /// <param name="inputDataFiles">An optional list of paths to additional input data files required for the operation. Can be <see
    /// langword="null"/> if no additional files are needed.</param>
    /// <returns>A <see cref="Results"/> value indicating the outcome of the operation. Returns <see cref="Results.Success"/> if
    /// the operation completes successfully; otherwise, returns an appropriate error code.</returns>
    public static async Task<Results> RunAverageOfflineWorker(string name, string promptFile, string outputFile, List<string>? inputDataFiles = null, bool jsonOutput = true)
    {
        // Make sure that everything is in place in order to do some work
        var results = CheckDocuments(name, promptFile, outputFile, inputDataFiles);

        // Anything but success is necessary
        if (results != Results.Success)
            return results;

        GenerativeModel model = GetBaseModel;
        model.UseGoogleSearch = false;
        model.UseGrounding = false;

        // Launch the average worker with the default average base configuration
        var response = await RunAsync(model, GetBaseConfig, ApiKeys.GetNextFreeKey(), name, promptFile, outputFile, inputDataFiles);

        return response;
    }

    /// <summary>
    /// Executes the Clever Offline Worker process using the specified parameters.
    /// </summary>
    /// <remarks>This method performs a series of checks to ensure that all required files and configurations
    /// are in place before executing the worker process. It uses a generative model with a predefined configuration to
    /// process the input and generate the output.</remarks>
    /// <param name="name">The name of the task or operation to be executed.</param>
    /// <param name="promptFile">The path to the prompt file that contains the instructions for the worker.</param>
    /// <param name="outputFile">The path to the file where the output of the operation will be saved.</param>
    /// <param name="inputDataFiles">An optional list of input data file paths to be used during the operation. Can be <see langword="null"/> if no
    /// input files are required.</param>
    /// <param name="jsonOutput">A value indicating whether the output should be formatted as JSON. The default value is <see langword="true"/>.</param>
    /// <returns>A <see cref="Results"/> enumeration value indicating the outcome of the operation. Returns <see
    /// cref="Results.Success"/> if the operation completes successfully; otherwise, returns an appropriate error code.</returns>
    public static async Task<Results> RunCleverOfflineWorker(string name, string promptFile, string outputFile, List<string>? inputDataFiles = null, bool jsonOutput = true)
    {
        // Make sure that everything is in place in order to do some work
        var results = CheckDocuments(name, promptFile, outputFile, inputDataFiles);

        // Anything but success is necessary
        if (results != Results.Success)
            return results;

        GenerativeModel model = GetBaseModel;
        model.Model = GeminiPro25;
        model.UseGoogleSearch = false;
        model.UseGrounding = false;

        // Launch the average worker with the default average base configuration
        var response = await RunAsync(model, GetBaseConfig, ApiKeys.GetNextFreeKey(), name, promptFile, outputFile, inputDataFiles);

        return response;
    }
    
    /// <summary>
    /// Executes a generative worker task using the specified configuration and input files.
    /// </summary>
    /// <remarks>This method validates the provided files and configuration before executing the task. It uses
    /// a generative model with default settings, including Google Search and grounding enabled. Ensure that all
    /// required files are accessible and properly formatted before calling this method.</remarks>
    /// <param name="name">The name of the task to be executed. This is used to identify the operation.</param>
    /// <param name="promptFile">The path to the prompt file that contains the instructions or input for the generative model.</param>
    /// <param name="outputFile">The path to the file where the output of the operation will be saved.</param>
    /// <param name="inputDataFiles">An optional list of paths to additional input data files required for the operation. Can be <see
    /// langword="null"/> if no additional input files are needed.</param>
    /// <param name="jsonOutput">A value indicating whether the output should be formatted as JSON. The default value is <see langword="true"/>.</param>
    /// <returns>A <see cref="Results"/> enumeration value indicating the outcome of the operation. Returns <see
    /// cref="Results.Success"/> if the operation completes successfully; otherwise, returns an appropriate error
    /// result.</returns>
    public static async Task<Results> RunCleverWorker(string name, string promptFile, string outputFile, List<string>? inputDataFiles = null, bool jsonOutput = true)
    {
        // Make sure that everything is in place in order to do some work
        var results = CheckDocuments(name, promptFile, outputFile, inputDataFiles);

        // Anything but success is necessary
        if (results != Results.Success)
            return results;

        GenerativeModel model = GetBaseModel;
        model.Model = GeminiPro25;
        model.UseGoogleSearch = true;
        model.UseGrounding = true;

        // Launch the average worker with the default average base configuration
        var response = await RunAsync(GetBaseModel, GetBaseConfig, ApiKeys.GetNextFreeKey(), name, promptFile, outputFile, inputDataFiles);

        return response;
    }
}