package com.agent.management.rag.structured;

import com.agent.management.rag.bridge.RagAiBridgeClient;
import com.agent.management.rag.model.RagQuery;
import lombok.RequiredArgsConstructor;
import org.springframework.stereotype.Service;

import java.util.ArrayList;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;

@Service
@RequiredArgsConstructor
public class VannaSqlGenerationService implements SqlGenerationService {
    private final RagAiBridgeClient bridge;
    private final StructuredQueryContextService contexts;

    @Override
    public Optional<String> generateSql(RagQuery query, Long datasourceId) {
        return generate(query, datasourceId).map(GeneratedSql::sql);
    }

    public Optional<GeneratedSql> generate(RagQuery query, Long datasourceId) {
        if (!allowed(query)) return Optional.empty();
        StructuredQueryContext context = contexts.build(query, datasourceId);
        return generate(query, datasourceId, context);
    }

    public Optional<GeneratedSql> generate(RagQuery query, Long datasourceId, StructuredQueryContext context) {
        if (!allowed(query)) return Optional.empty();
        if (!datasourceId.equals(context.dataSourceId())) {
            throw new IllegalStateException("structured query context belongs to another datasource");
        }
        Map<String, Object> request = request(query, datasourceId, context);
        Map<String, Object> response = bridge.textToSql(request);
        Object sql = response.get("sql");
        if (sql == null || String.valueOf(sql).isBlank()) return Optional.empty();
        List<String> warnings = new ArrayList<>(context.warnings());
        Object bridgeWarnings = response.get("warnings");
        if (bridgeWarnings instanceof List<?> list) list.stream().map(String::valueOf).forEach(warnings::add);
        return Optional.of(new GeneratedSql(String.valueOf(sql), context, warnings));
    }

    public Optional<String> repair(RagQuery query, Long datasourceId, String failedSql, String error,
                                   StructuredQueryContext context) {
        Map<String, Object> request = new LinkedHashMap<>();
        request.put("language", "SQL");
        request.put("question", query.getQuery());
        request.put("query", failedSql);
        request.put("error", error);
        request.put("schemaText", context.minimalDdl());
        request.put("dialect", context.dialect());
        request.put("selectedTables", context.selectedTables());
        request.put("foreignKeys", context.foreignKeys());
        request.put("relationships", context.relationships());
        request.put("sampledValues", context.sampledValues());
        request.put("referenceRows", context.referenceRows());
        request.put("maxRows", limit(query));
        request.put("maxDepth", 1);
        Object fixed = bridge.repair(request).get("query");
        return fixed == null || String.valueOf(fixed).isBlank()
                ? Optional.empty() : Optional.of(String.valueOf(fixed));
    }

    public Optional<String> repair(RagQuery query, Long datasourceId, String failedSql, String error) {
        return repair(query, datasourceId, failedSql, error, contexts.build(query, datasourceId));
    }

    static boolean allowed(RagQuery query) {
        if (query.getFilters() == null) return false;
        return Boolean.TRUE.equals(query.getFilters().get("allowTextToSql"))
                || "AUTO_GENERATE".equals(String.valueOf(query.getFilters().get("retrievalMode")));
    }

    private static Map<String, Object> request(RagQuery query, Long datasourceId,
                                                StructuredQueryContext context) {
        Map<String, Object> request = new LinkedHashMap<>();
        request.put("query", query.getQuery());
        request.put("datasourceId", datasourceId);
        request.put("dialect", context.dialect());
        request.put("schemaVersion", context.schemaVersion());
        request.put("schemaFingerprint", context.schemaFingerprint());
        request.put("profileVersion", context.profileVersion());
        request.put("ddl", context.minimalDdl());
        request.put("selectedTables", context.selectedTables());
        request.put("foreignKeys", context.foreignKeys());
        request.put("relationships", context.relationships());
        request.put("sampledValues", context.sampledValues());
        request.put("referenceRows", context.referenceRows());
        request.put("entityMentions", context.entityMentions());
        request.put("contextWarnings", context.warnings());
        request.put("documentation", generationRules(query));
        request.put("maxRows", limit(query));
        return request;
    }

    private static int limit(RagQuery query) {
        Object configured = query.getFilters() == null ? null : query.getFilters().get("maxRows");
        int value = configured instanceof Number number ? number.intValue()
                : (query.getTopK() != null ? query.getTopK() : 50);
        return Math.max(1, Math.min(value, 1000));
    }

    private static String generationRules(RagQuery query) {
        if (query.getFilters() == null || !(query.getFilters().get("generationRules") instanceof List<?> rules)) {
            return "";
        }
        return rules.stream().map(String::valueOf).collect(java.util.stream.Collectors.joining("\n- ", "- ", ""));
    }

    public record GeneratedSql(String sql, StructuredQueryContext context, List<String> warnings) {
        public GeneratedSql { warnings = List.copyOf(warnings == null ? List.of() : warnings); }
    }
}
