package com.agent.management.rag.structured;

import com.agent.management.model.entity.DataSource;
import com.agent.management.model.entity.RagStructuredSchemaProfileEntity;
import com.agent.management.rag.capability.SqlSchemaSnapshot;
import com.agent.management.repository.RagStructuredSchemaProfileRepository;
import com.agent.management.service.DataSourceService;
import com.agent.management.service.DynamicJdbcService;
import com.agent.management.service.SchemaExplorerService;
import com.fasterxml.jackson.databind.ObjectMapper;
import jakarta.annotation.PreDestroy;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import java.nio.charset.StandardCharsets;
import java.security.MessageDigest;
import java.sql.Connection;
import java.time.Duration;
import java.time.Instant;
import java.util.ArrayList;
import java.util.HexFormat;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;

@Service
public class StructuredSchemaProfileService {
    public static final String ALGORITHM_VERSION = "logical-rel-v2";
    private static final Duration RELATION_REVERIFY_INTERVAL = Duration.ofHours(24);

    private final DataSourceService dataSources;
    private final DynamicJdbcService jdbc;
    private final SchemaExplorerService explorer;
    private final StructuredRelationshipDiscoveryService discovery;
    private final RagStructuredSchemaProfileRepository repository;
    private final ObjectMapper objectMapper;
    private final ExecutorService scanExecutor;
    private final Duration firstScanBudget;
    private final Map<CacheKey, StructuredSchemaProfile> cache = new ConcurrentHashMap<>();
    private final Map<CacheKey, CompletableFuture<StructuredSchemaProfile>> scans = new ConcurrentHashMap<>();

    @Autowired
    public StructuredSchemaProfileService(DataSourceService dataSources, DynamicJdbcService jdbc,
                                          SchemaExplorerService explorer,
                                          StructuredRelationshipDiscoveryService discovery,
                                          RagStructuredSchemaProfileRepository repository,
                                          ObjectMapper objectMapper) {
        this(dataSources, jdbc, explorer, discovery, repository, objectMapper,
                Executors.newFixedThreadPool(2, runnable -> {
                    Thread thread = new Thread(runnable, "structured-profile-scan");
                    thread.setDaemon(true);
                    return thread;
                }), Duration.ofSeconds(3));
    }

    StructuredSchemaProfileService(DataSourceService dataSources, DynamicJdbcService jdbc,
                                   SchemaExplorerService explorer,
                                   StructuredRelationshipDiscoveryService discovery,
                                   RagStructuredSchemaProfileRepository repository,
                                   ObjectMapper objectMapper, ExecutorService scanExecutor,
                                   Duration firstScanBudget) {
        this.dataSources = dataSources;
        this.jdbc = jdbc;
        this.explorer = explorer;
        this.discovery = discovery;
        this.repository = repository;
        this.objectMapper = objectMapper;
        this.scanExecutor = scanExecutor;
        this.firstScanBudget = firstScanBudget;
    }

    public StructuredSchemaProfile getProfile(Long dataSourceId) {
        DataSource source = dataSources.getDecrypted(dataSourceId);
        if (source == null || !dataSourceId.equals(source.getId())) {
            throw new IllegalArgumentException("structured datasource does not match requested dataSourceId");
        }
        Identity identity = identity(source);
        Baseline baseline = baseline(dataSourceId, source, identity);
        CacheKey key = new CacheKey(dataSourceId, identity.sourceFingerprint(), identity.catalog(), identity.schema(),
                baseline.schemaFingerprint(), ALGORITHM_VERSION);

        StructuredSchemaProfile memory = cache.get(key);
        if (fresh(memory)) return memory;

        Optional<StructuredSchemaProfile> persisted = loadPersisted(key);
        if (persisted.isPresent() && fresh(persisted.get())) {
            cache.put(key, persisted.get());
            return persisted.get();
        }

        CompletableFuture<StructuredSchemaProfile> scan = scans.computeIfAbsent(key, ignored -> startScan(key, baseline));
        try {
            return scan.get(firstScanBudget.toMillis(), TimeUnit.MILLISECONDS);
        } catch (TimeoutException error) {
            List<String> warnings = new ArrayList<>(baseline.warnings());
            warnings.add("logical relationship scan did not complete within query budget; using physical schema only");
            return fallbackProfile(key, baseline, warnings);
        } catch (Exception error) {
            List<String> warnings = new ArrayList<>(baseline.warnings());
            warnings.add("logical relationship scan failed; using physical schema only: " + shortMessage(error));
            return fallbackProfile(key, baseline, warnings);
        }
    }

    public void prewarm(Long dataSourceId) {
        CompletableFuture.runAsync(() -> {
            try { getProfile(dataSourceId); } catch (Exception ignored) { }
        });
    }

    public void refreshAsync(Long dataSourceId) {
        cache.keySet().removeIf(key -> dataSourceId.equals(key.dataSourceId()));
        CompletableFuture.runAsync(() -> {
            try {
                DataSource source = dataSources.getDecrypted(dataSourceId);
                Identity identity = identity(source);
                Baseline baseline = baseline(dataSourceId, source, identity);
                CacheKey key = new CacheKey(dataSourceId, identity.sourceFingerprint(), identity.catalog(),
                        identity.schema(), baseline.schemaFingerprint(), ALGORITHM_VERSION);
                scans.computeIfAbsent(key, ignored -> startScan(key, baseline));
            } catch (Exception ignored) { }
        });
    }

    public Map<String, Object> status(Long dataSourceId) {
        boolean scanning = scans.keySet().stream().anyMatch(key -> dataSourceId.equals(key.dataSourceId()));
        Optional<RagStructuredSchemaProfileEntity> stored =
                repository.findFirstByDataSourceIdOrderByScannedAtDesc(dataSourceId);
        if (stored.isEmpty()) {
            return Map.of("dataSourceId", dataSourceId, "scanStatus", scanning ? "SCANNING" : "NOT_SCANNED");
        }
        RagStructuredSchemaProfileEntity entity = stored.get();
        Map<String, Object> result = new LinkedHashMap<>();
        result.put("dataSourceId", dataSourceId);
        result.put("scanStatus", scanning ? "SCANNING" : entity.getScanStatus());
        result.put("schemaFingerprint", entity.getSchemaFingerprint());
        result.put("algorithmVersion", entity.getAlgorithmVersion());
        result.put("scannedAt", entity.getScannedAt());
        result.put("lastError", entity.getLastError() == null ? "" : entity.getLastError());
        try {
            StructuredSchemaProfile profile = objectMapper.readValue(
                    entity.getProfileJson(), StructuredSchemaProfile.class);
            result.put("profileVersion", profile.profileVersion());
            result.put("tableCount", profile.schemaSnapshot().tables().size());
            result.put("physicalRelationships", count(profile, StructuredRelationship.RelationshipType.PHYSICAL_FK));
            result.put("inferredRelationships", count(profile, StructuredRelationship.RelationshipType.INFERRED));
            result.put("conditionalRelationships", count(profile, StructuredRelationship.RelationshipType.CONDITIONAL));
            result.put("ambiguousRelationships", count(profile, StructuredRelationship.RelationshipType.AMBIGUOUS));
            result.put("dialectSupported", "mysql".equalsIgnoreCase(profile.dialect()));
            result.put("metricFields", metricFields(profile));
            result.put("fieldCommentCoverage", fieldCommentCoverage(profile));
            result.put("recommendedQuestionTypes", recommendedQuestionTypes(profile));
            result.put("warnings", profile.warnings());
        } catch (Exception error) {
            result.put("warnings", List.of("stored profile JSON is unreadable"));
        }
        return result;
    }

    @jakarta.transaction.Transactional
    public void deleteProfiles(Long dataSourceId) {
        cache.keySet().removeIf(key -> dataSourceId.equals(key.dataSourceId()));
        repository.deleteByDataSourceId(dataSourceId);
    }

    private CompletableFuture<StructuredSchemaProfile> startScan(CacheKey key, Baseline baseline) {
        CompletableFuture<StructuredSchemaProfile> future = CompletableFuture.supplyAsync(() -> {
            StructuredRelationshipDiscoveryService.DiscoveryResult discovered = discovery.discover(
                    key.dataSourceId(), key.catalog(), key.schema(), baseline.snapshot(), baseline.columnsByTable());
            List<String> warnings = new ArrayList<>(baseline.warnings());
            warnings.addAll(discovered.warnings());
            Instant scannedAt = Instant.now();
            return new StructuredSchemaProfile(key.dataSourceId(), baseline.dialect(), key.catalog(), key.schema(),
                    key.sourceFingerprint(), key.schemaFingerprint(), key.algorithmVersion(),
                    profileVersion(key.schemaFingerprint(), scannedAt), "SUCCESS", scannedAt,
                    baseline.snapshot(), baseline.columnsByTable(), discovered.relationships(), warnings);
        }, scanExecutor);
        future.whenComplete((profile, error) -> {
            try {
                if (error == null) {
                    cache.put(key, profile);
                    saveSuccessful(profile);
                } else {
                    saveFailureIfAbsent(key, error);
                }
            } finally {
                scans.remove(key, future);
            }
        });
        return future;
    }

    private Baseline baseline(Long dataSourceId, DataSource source, Identity identity) {
        List<String> warnings = new ArrayList<>();
        Map<String, List<SchemaExplorerService.ColumnSpec>> columnsByTable = new LinkedHashMap<>();
        List<SqlSchemaSnapshot.TableSchema> tables = new ArrayList<>();
        for (SchemaExplorerService.TableSummary table : explorer.listTables(dataSourceId, emptyToNull(identity.schema()))) {
            List<SchemaExplorerService.ColumnSpec> columns = explorer.describeTable(
                    dataSourceId, emptyToNull(identity.schema()), table.name());
            columnsByTable.put(table.name(), columns);
            tables.add(new SqlSchemaSnapshot.TableSchema(table.name(), table.remarks(), columns.stream()
                    .map(column -> new SqlSchemaSnapshot.ColumnSchema(column.name(), column.type(), column.remarks()))
                    .toList()));
        }
        if (tables.isEmpty()) throw new IllegalArgumentException("datasource schema contains no visible tables");
        List<SqlSchemaSnapshot.ForeignKey> foreignKeys = new ArrayList<>();
        for (SqlSchemaSnapshot.TableSchema table : tables) {
            try {
                explorer.listForeignKeys(dataSourceId, emptyToNull(identity.schema()), table.name()).stream()
                        .map(key -> new SqlSchemaSnapshot.ForeignKey(key.name(), key.sourceTable(), key.sourceColumn(),
                                key.targetTable(), key.targetColumn())).forEach(foreignKeys::add);
            } catch (Exception error) {
                warnings.add("foreign keys for " + table.name() + " unavailable: " + shortMessage(error));
            }
        }
        if (foreignKeys.isEmpty()) warnings.add("database metadata exposes no physical foreign keys");
        SqlSchemaSnapshot snapshot = new SqlSchemaSnapshot(tables, foreignKeys);
        return new Baseline(dialect(source.getType()), snapshot, columnsByTable,
                schemaFingerprint(snapshot, columnsByTable), warnings);
    }

    private Identity identity(DataSource source) {
        String catalog = "", schema = source.getSchemaName() == null ? "" : source.getSchemaName();
        try (Connection connection = jdbc.getConnection(source)) {
            catalog = value(connection.getCatalog());
            if (schema.isBlank()) {
                try { schema = value(connection.getSchema()); } catch (Exception ignored) { schema = ""; }
            }
        } catch (Exception error) {
            throw new IllegalArgumentException("cannot resolve requested datasource identity: " + shortMessage(error), error);
        }
        String fingerprint = hash(value(source.getType()) + "|" + value(source.getJdbcUrl()) + "|"
                + value(source.getUsername()) + "|" + catalog + "|" + schema);
        return new Identity(catalog, schema, fingerprint);
    }

    private Optional<StructuredSchemaProfile> loadPersisted(CacheKey key) {
        return repository.findByDataSourceIdAndCatalogNameAndSchemaName(
                        key.dataSourceId(), key.catalog(), key.schema())
                .filter(entity -> key.sourceFingerprint().equals(entity.getSourceFingerprint()))
                .filter(entity -> key.schemaFingerprint().equals(entity.getSchemaFingerprint()))
                .filter(entity -> key.algorithmVersion().equals(entity.getAlgorithmVersion()))
                .filter(entity -> "SUCCESS".equals(entity.getScanStatus()))
                .flatMap(entity -> {
                    try { return Optional.of(objectMapper.readValue(entity.getProfileJson(), StructuredSchemaProfile.class)); }
                    catch (Exception ignored) { return Optional.empty(); }
                });
    }

    private void saveSuccessful(StructuredSchemaProfile profile) {
        try {
            RagStructuredSchemaProfileEntity entity = repository
                    .findByDataSourceIdAndCatalogNameAndSchemaName(
                            profile.dataSourceId(), profile.catalog(), profile.schema())
                    .orElseGet(RagStructuredSchemaProfileEntity::new);
            entity.setDataSourceId(profile.dataSourceId());
            entity.setCatalogName(profile.catalog());
            entity.setSchemaName(profile.schema());
            entity.setProfileJson(objectMapper.writeValueAsString(profile));
            entity.setSourceFingerprint(profile.sourceFingerprint());
            entity.setSchemaFingerprint(profile.schemaFingerprint());
            entity.setAlgorithmVersion(profile.algorithmVersion());
            entity.setScanStatus("SUCCESS");
            entity.setScannedAt(profile.scannedAt());
            entity.setLastError(null);
            repository.save(entity);
        } catch (Exception ignored) {
            // 持久化失败不能影响本次只读查询；内存画像仍可继续使用。
        }
    }

    private void saveFailureIfAbsent(CacheKey key, Throwable error) {
        try {
            if (repository.findByDataSourceIdAndCatalogNameAndSchemaName(
                    key.dataSourceId(), key.catalog(), key.schema()).isPresent()) return;
            RagStructuredSchemaProfileEntity entity = new RagStructuredSchemaProfileEntity();
            entity.setDataSourceId(key.dataSourceId());
            entity.setCatalogName(key.catalog());
            entity.setSchemaName(key.schema());
            entity.setProfileJson("{}");
            entity.setSourceFingerprint(key.sourceFingerprint());
            entity.setSchemaFingerprint(key.schemaFingerprint());
            entity.setAlgorithmVersion(key.algorithmVersion());
            entity.setScanStatus("FAILED");
            entity.setScannedAt(Instant.now());
            entity.setLastError(shortMessage(error));
            repository.save(entity);
        } catch (Exception ignored) {
            // 失败状态写入失败也不能阻断业务查询。
        }
    }

    private static StructuredSchemaProfile fallbackProfile(CacheKey key, Baseline baseline, List<String> warnings) {
        List<StructuredRelationship> physical = baseline.snapshot().foreignKeys().stream()
                .map(foreignKey -> new StructuredRelationship(StructuredRelationship.RelationshipType.PHYSICAL_FK,
                        foreignKey.sourceTable(), foreignKey.sourceColumn(), foreignKey.targetTable(),
                        foreignKey.targetColumn(), 1, List.of(),
                        List.of("JDBC imported foreign key: " + String.valueOf(foreignKey.name())), List.of()))
                .toList();
        Instant now = Instant.now();
        return new StructuredSchemaProfile(key.dataSourceId(), baseline.dialect(), key.catalog(), key.schema(),
                key.sourceFingerprint(), key.schemaFingerprint(), key.algorithmVersion(),
                profileVersion(key.schemaFingerprint(), now), "FALLBACK", now,
                baseline.snapshot(), baseline.columnsByTable(), physical, warnings);
    }

    private static boolean fresh(StructuredSchemaProfile profile) {
        return profile != null && "SUCCESS".equals(profile.scanStatus()) && profile.scannedAt() != null
                && profile.scannedAt().isAfter(Instant.now().minus(RELATION_REVERIFY_INTERVAL));
    }

    private static String profileVersion(String schemaFingerprint, Instant scannedAt) {
        return ALGORITHM_VERSION + ":" + schemaFingerprint.substring(0, 12) + ":" + scannedAt.toEpochMilli();
    }

    private static long count(StructuredSchemaProfile profile, StructuredRelationship.RelationshipType type) {
        return profile.relationships().stream().filter(relationship -> relationship.type() == type).count();
    }

    private static List<String> metricFields(StructuredSchemaProfile profile) {
        List<String> result = new ArrayList<>();
        profile.columnSpecs().forEach((table, columns) -> columns.forEach(column -> {
            String name = column.name().toLowerCase();
            if (name.equals("score") || name.endsWith("_score") || name.contains("rank")
                    || name.contains("priority") || name.contains("rating")
                    || name.equals("weight") || name.equals("level")) {
                result.add(table + "." + column.name());
            }
        }));
        return result;
    }

    private static double fieldCommentCoverage(StructuredSchemaProfile profile) {
        long total = profile.columnSpecs().values().stream().mapToLong(List::size).sum();
        if (total == 0) return 0;
        long described = profile.columnSpecs().values().stream().flatMap(List::stream)
                .filter(column -> column.remarks() != null && !column.remarks().isBlank()).count();
        return Math.round((described * 1000.0 / total)) / 1000.0;
    }

    private static List<String> recommendedQuestionTypes(StructuredSchemaProfile profile) {
        List<String> result = new ArrayList<>();
        result.add("单表筛选/列表查询");
        if (!profile.relationships().isEmpty()) result.add("多表关联查询");
        if (!metricFields(profile).isEmpty()) result.add("评分/排名/优先级查询");
        if (!"mysql".equalsIgnoreCase(profile.dialect())) result.add("自动 Text2SQL 暂不支持该方言");
        return result;
    }

    private static String dialect(String type) {
        if (type == null) return "";
        return switch (type.toUpperCase()) {
            case "MYSQL" -> "mysql";
            case "POSTGRESQL", "PG" -> "postgresql";
            default -> type.toLowerCase();
        };
    }

    private static String hash(String value) {
        try {
            return HexFormat.of().formatHex(MessageDigest.getInstance("SHA-256")
                    .digest(value.getBytes(StandardCharsets.UTF_8)));
        } catch (Exception error) { throw new IllegalStateException(error); }
    }

    private static String schemaFingerprint(SqlSchemaSnapshot snapshot,
                                            Map<String, List<SchemaExplorerService.ColumnSpec>> columnsByTable) {
        StringBuilder canonical = new StringBuilder();
        snapshot.tables().forEach(table -> {
            canonical.append("T|").append(table.name()).append('|').append(value(table.description())).append('\n');
            columnsByTable.getOrDefault(table.name(), List.of()).forEach(column -> canonical
                    .append("C|").append(table.name()).append('|').append(column.name()).append('|')
                    .append(column.type()).append('|').append(column.size()).append('|').append(column.nullable())
                    .append('|').append(value(column.remarks())).append('\n'));
        });
        snapshot.foreignKeys().forEach(key -> canonical.append("F|").append(key.sourceTable()).append('|')
                .append(key.sourceColumn()).append('|').append(key.targetTable()).append('|')
                .append(key.targetColumn()).append('\n'));
        return hash(canonical.toString());
    }

    private static String value(String value) { return value == null ? "" : value; }
    private static String emptyToNull(String value) { return value == null || value.isBlank() ? null : value; }
    private static String shortMessage(Exception error) { String message = error.getMessage(); return message == null
            ? error.getClass().getSimpleName() : message.substring(0, Math.min(200, message.length())); }
    private static String shortMessage(Throwable error) { String message = error.getMessage(); return message == null
            ? error.getClass().getSimpleName() : message.substring(0, Math.min(200, message.length())); }

    @PreDestroy
    public void close() { scanExecutor.shutdownNow(); }

    private record Identity(String catalog, String schema, String sourceFingerprint) {}
    private record Baseline(String dialect, SqlSchemaSnapshot snapshot,
                            Map<String, List<SchemaExplorerService.ColumnSpec>> columnsByTable,
                            String schemaFingerprint, List<String> warnings) {}
    public record CacheKey(Long dataSourceId, String sourceFingerprint, String catalog, String schema,
                           String schemaFingerprint, String algorithmVersion) {}
}
