using Nalu.Web.Models;

namespace Nalu.Web.Services;

public class ExtractionPipeline
{
    private readonly ValidatorLoader _loader;
    private readonly DeterministicLayer _deterministic;
    private readonly LlmExtractionService _llm;
    private readonly PostProcessorRegistry _postProcessors;
    private readonly EnrichmentService _enrichment;
    private readonly SuggestionBuilder _suggestions;
    private readonly CacheService _cache;

    public ExtractionPipeline(
        ValidatorLoader loader,
        DeterministicLayer deterministic,
        LlmExtractionService llm,
        PostProcessorRegistry postProcessors,
        EnrichmentService enrichment,
        SuggestionBuilder suggestions,
        CacheService cache)
    {
        _loader = loader;
        _deterministic = deterministic;
        _llm = llm;
        _postProcessors = postProcessors;
        _enrichment = enrichment;
        _suggestions = suggestions;
        _cache = cache;
    }

    public async Task<ExtractionResponse> ExecuteAsync(
        string validatorId,
        ExtractionRequest request,
        CancellationToken ct = default)
    {
        if (_cache.TryGet(validatorId, request, out var cached) && cached is not null)
            return cached;

        var validator = _loader.Load(validatorId);
        var det = _deterministic.Evaluate(validator, request.UserInput, request.Language);

        ExtractionResponse response = det.Outcome switch
        {
            DeterministicOutcome.Accepted =>
                await BuildFromValueAsync(validator, request, det.ExtractedValue,
                    highConfidence: true, ct),

            DeterministicOutcome.Rejected =>
                BuildFailed(validator, request, wasInvalidated: false),

            DeterministicOutcome.ConstraintFailed =>
                BuildFailed(validator, request, wasInvalidated: true),

            _ => await BuildFromLlmAsync(validator, request, ct)
        };

        _cache.Set(validatorId, request, response);
        return response;
    }

    private async Task<ExtractionResponse> BuildFromValueAsync(
        ValidatorDefinition validator,
        ExtractionRequest request,
        string? rawValue,
        bool highConfidence,
        CancellationToken ct)
    {
        var procOutput = _postProcessors.Apply(validator.PostProcessors, rawValue);

        if (!procOutput.IsValid)
        {
            var sug = _suggestions.Build(validator, request, null,
                obtained: false, certain: false, wasInvalidated: true);
            return new ExtractionResponse
            {
                Obtained = false, ExtractedValue = null, Confidence = "low",
                Certain = false, SuggestionToAgent = sug
            };
        }

        var enrichResult = await _enrichment.EnrichAsync(validator, procOutput.Value, ct);

        if (enrichResult.WasInvalidated)
        {
            var sug = _suggestions.Build(validator, request, null,
                obtained: false, certain: false, wasInvalidated: true);
            return new ExtractionResponse
            {
                Obtained = false, ExtractedValue = null, Confidence = "low",
                Certain = false, SuggestionToAgent = sug
            };
        }

        var value = enrichResult.Value;
        var obtained = value is not null;
        var certain = obtained && highConfidence && !procOutput.WasCorrected;

        var suggestion = _suggestions.Build(validator, request, value,
            obtained, certain,
            wasInvalidated: false,
            wasCorrected: procOutput.WasCorrected,
            originalValue: procOutput.OriginalValue,
            suggestionKeyOverride: procOutput.SuggestionKeyOverride);

        var confidence = obtained
            ? (certain ? (highConfidence ? "high" : "medium") : "medium")
            : "low";

        var valueFormat = value is null ? null
            : value.TrimStart().StartsWith('{') ? "object"
            : "scalar";

        return new ExtractionResponse
        {
            Obtained = obtained,
            ExtractedValue = value,
            Confidence = confidence,
            Certain = certain,
            SuggestionToAgent = suggestion,
            ValueFormat = valueFormat
        };
    }

    private ExtractionResponse BuildFailed(
        ValidatorDefinition validator,
        ExtractionRequest request,
        bool wasInvalidated)
    {
        var suggestion = _suggestions.Build(validator, request, null,
            obtained: false, certain: false, wasInvalidated: wasInvalidated);

        return new ExtractionResponse
        {
            Obtained = false, ExtractedValue = null, Confidence = "low",
            Certain = false, SuggestionToAgent = suggestion
        };
    }

    private async Task<ExtractionResponse> BuildFromLlmAsync(
        ValidatorDefinition validator,
        ExtractionRequest request,
        CancellationToken ct)
    {
        var llm = await _llm.ExtractAsync(validator, request, ct);

        return await BuildFromValueAsync(
            validator, request, llm.ExtractedValue,
            highConfidence: llm.Certain, ct);
    }
}
