package com.ai.trackwalker.demo;

import com.ai.trackwalker.common.Ids;
import com.ai.trackwalker.common.Jsonb;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.stereotype.Component;
import org.springframework.transaction.annotation.Transactional;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.IOException;
import java.sql.Timestamp;
import java.time.Instant;
import java.util.Arrays;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CopyOnWriteArrayList;

@Component
public class DemoRunEventPublisher {
    private static final long SSE_TIMEOUT_MS = 30 * 60 * 1000L;

    private final JdbcTemplate jdbc;
    private final Map<String, CopyOnWriteArrayList<SseEmitter>> emitters = new ConcurrentHashMap<>();

    public DemoRunEventPublisher(JdbcTemplate jdbc) {
        this.jdbc = jdbc;
    }

    public SseEmitter subscribe(String runId) {
        SseEmitter emitter = new SseEmitter(SSE_TIMEOUT_MS);
        emitters.computeIfAbsent(runId, ignored -> new CopyOnWriteArrayList<>()).add(emitter);
        emitter.onCompletion(() -> removeEmitter(runId, emitter));
        emitter.onTimeout(() -> removeEmitter(runId, emitter));
        emitter.onError(error -> removeEmitter(runId, emitter));

        for (Map<String, Object> event : history(runId)) {
            send(emitter, event);
        }
        return emitter;
    }

    @Transactional
    public Map<String, Object> emit(
            String runId,
            DemoStep step,
            String status,
            int progress,
            String title,
            String message,
            String entityType,
            String entityId,
            Map<String, Object> metrics,
            Map<String, Object> visualPayload
    ) {
        int attempt = latestAttempt(runId, step.key());
        String eventType = defaultEventType(runId, step, status);
        int stepProgress = calculateStepProgress(step, progress, status);
        return emitInternal(runId, step, status, progress, stepProgress, eventType, attempt, title, message,
                entityType, entityId, metrics, visualPayload);
    }

    @Transactional
    public Map<String, Object> emitRetry(
            String runId,
            DemoStep step,
            int progress,
            String title,
            String message,
            String entityType,
            String entityId,
            Map<String, Object> metrics,
            Map<String, Object> visualPayload
    ) {
        int attempt = latestAttempt(runId, step.key()) + 1;
        return emitInternal(runId, step, "retrying", progress, 0, "step.retrying", attempt, title, message,
                entityType, entityId, metrics, visualPayload);
    }

    private Map<String, Object> emitInternal(
            String runId,
            DemoStep step,
            String status,
            int progress,
            int stepProgress,
            String eventType,
            int attempt,
            String title,
            String message,
            String entityType,
            String entityId,
            Map<String, Object> metrics,
            Map<String, Object> visualPayload
    ) {
        String eventId = Ids.next("dre");
        Instant now = Instant.now();
        Map<String, Object> safeMetrics = metrics == null ? Map.of() : metrics;
        Map<String, Object> safeVisualPayload = visualPayload == null ? Map.of() : visualPayload;
        long sequence = nextSequence(runId);

        jdbc.update(
                "insert into demo_run_events(id, run_id, step_key, step_name, status, progress, title, message, entity_type, entity_id, metrics, visual_payload, created_at, sequence_no, event_type, step_progress, attempt) values (?,?,?,?,?,?,?,?,?,?,?::jsonb,?::jsonb,?,?,?,?,?)",
                eventId,
                runId,
                step.key(),
                step.label(),
                status,
                progress,
                title,
                message,
                entityType,
                entityId,
                Jsonb.write(safeMetrics),
                Jsonb.write(safeVisualPayload),
                Timestamp.from(now),
                sequence,
                eventType,
                stepProgress,
                attempt
        );

        mergeStepSnapshot(runId, step, status, progress, stepProgress, eventType, attempt, title, message,
                entityType, entityId, safeMetrics, safeVisualPayload, now);

        String runStatus = runStatus(step, status);
        if ("failed".equals(runStatus)) {
            jdbc.update("update demo_runs set status=?, current_step=?, progress=?, completed_at=?, error_message=? where id=?",
                    runStatus, step.key(), progress, Timestamp.from(now), message, runId);
        } else if ("completed".equals(runStatus) || "cancelled".equals(runStatus)) {
            jdbc.update("update demo_runs set status=?, current_step=?, progress=?, completed_at=? where id=?",
                    runStatus, step.key(), progress, Timestamp.from(now), runId);
        } else {
            jdbc.update("update demo_runs set status=?, current_step=?, progress=?, completed_at=null, error_message=null where id=?",
                    runStatus, step.key(), progress, runId);
        }

        Map<String, Object> event = Map.ofEntries(
                Map.entry("event_id", eventId),
                Map.entry("run_id", runId),
                Map.entry("sequence", sequence),
                Map.entry("step_key", step.key()),
                Map.entry("step_name", step.label()),
                Map.entry("event_type", eventType),
                Map.entry("status", status),
                Map.entry("progress", progress),
                Map.entry("run_progress", progress),
                Map.entry("step_progress", stepProgress),
                Map.entry("attempt", attempt),
                Map.entry("title", title),
                Map.entry("message", message == null ? "" : message),
                Map.entry("entity_type", entityType == null ? "" : entityType),
                Map.entry("entity_id", entityId == null ? "" : entityId),
                Map.entry("metrics", safeMetrics),
                Map.entry("visual_payload", safeVisualPayload),
                Map.entry("created_at", now.toString())
        );
        broadcast(runId, event);
        return event;
    }

    public List<Map<String, Object>> history(String runId) {
        return jdbc.query("""
                        select id as event_id,
                               run_id,
                               sequence_no,
                               step_key,
                               step_name,
                               event_type,
                               status,
                               progress,
                               step_progress,
                               attempt,
                               title,
                               message,
                               entity_type,
                               entity_id,
                               metrics::text as metrics,
                               visual_payload::text as visual_payload,
                               created_at
                        from demo_run_events
                        where run_id=?
                        order by sequence_no
                        """,
                (rs, rowNum) -> Map.ofEntries(
                        Map.entry("event_id", rs.getString("event_id")),
                        Map.entry("run_id", rs.getString("run_id")),
                        Map.entry("sequence", rs.getLong("sequence_no")),
                        Map.entry("step_key", rs.getString("step_key")),
                        Map.entry("step_name", rs.getString("step_name")),
                        Map.entry("event_type", rs.getString("event_type")),
                        Map.entry("status", rs.getString("status")),
                        Map.entry("progress", rs.getInt("progress")),
                        Map.entry("run_progress", rs.getInt("progress")),
                        Map.entry("step_progress", rs.getInt("step_progress")),
                        Map.entry("attempt", rs.getInt("attempt")),
                        Map.entry("title", rs.getString("title")),
                        Map.entry("message", rs.getString("message") == null ? "" : rs.getString("message")),
                        Map.entry("entity_type", rs.getString("entity_type") == null ? "" : rs.getString("entity_type")),
                        Map.entry("entity_id", rs.getString("entity_id") == null ? "" : rs.getString("entity_id")),
                        Map.entry("metrics", Jsonb.map(rs.getString("metrics"))),
                        Map.entry("visual_payload", Jsonb.map(rs.getString("visual_payload"))),
                        Map.entry("created_at", rs.getTimestamp("created_at").toInstant().toString())
                ),
                runId
        );
    }

    private long nextSequence(String runId) {
        Long value = jdbc.queryForObject(
                "update demo_runs set next_sequence=next_sequence+1 where id=? returning next_sequence",
                Long.class,
                runId
        );
        return value == null ? 1 : value;
    }

    private int latestAttempt(String runId, String stepKey) {
        Integer value = jdbc.queryForObject(
                "select coalesce(max(attempt), 1) from demo_run_events where run_id=? and step_key=?",
                Integer.class,
                runId,
                stepKey
        );
        return value == null ? 1 : value;
    }

    private String defaultEventType(String runId, DemoStep step, String status) {
        if ("completed".equals(status)) {
            return "step.completed";
        }
        if ("failed".equals(status)) {
            return "step.failed";
        }
        if ("cancelled".equals(status)) {
            return "step.cancelled";
        }
        if ("retrying".equals(status)) {
            return "step.retrying";
        }
        Integer count = jdbc.queryForObject(
                "select count(*) from demo_run_events where run_id=? and step_key=?",
                Integer.class,
                runId,
                step.key()
        );
        return count == null || count == 0 ? "step.started" : "step.progress";
    }

    private int calculateStepProgress(DemoStep step, int runProgress, String status) {
        if ("completed".equals(status)) {
            return 100;
        }
        List<DemoStep> steps = Arrays.stream(DemoStep.values())
                .filter(item -> item != DemoStep.RUN_CANCELLED && item != DemoStep.RUN_FAILED)
                .toList();
        int index = steps.indexOf(step);
        if (index < 0) {
            return 100;
        }
        int previousProgress = index == 0 ? 0 : steps.get(index - 1).progress();
        int span = Math.max(step.progress() - previousProgress, 1);
        return Math.max(0, Math.min(99, (int) Math.round((runProgress - previousProgress) * 100.0 / span)));
    }

    private void mergeStepSnapshot(
            String runId,
            DemoStep step,
            String status,
            int runProgress,
            int stepProgress,
            String eventType,
            int attempt,
            String title,
            String message,
            String entityType,
            String entityId,
            Map<String, Object> metrics,
            Map<String, Object> visualPayload,
            Instant now
    ) {
        String json = jdbc.queryForObject("select step_snapshots::text from demo_runs where id=?", String.class, runId);
        Map<String, Object> snapshots = new LinkedHashMap<>(Jsonb.map(json));
        Map<String, Object> existing = objectMap(snapshots.get(step.key()));
        Map<String, Object> mergedVisual = deepMerge(objectMap(existing.get("visual_payload")), visualPayload);

        Map<String, Object> snapshot = new LinkedHashMap<>();
        snapshot.put("step_key", step.key());
        snapshot.put("step_name", step.label());
        snapshot.put("status", status);
        snapshot.put("event_type", eventType);
        snapshot.put("run_progress", runProgress);
        snapshot.put("step_progress", stepProgress);
        snapshot.put("attempt", attempt);
        snapshot.put("title", title);
        snapshot.put("message", message == null ? "" : message);
        snapshot.put("entity_type", entityType == null ? "" : entityType);
        snapshot.put("entity_id", entityId == null ? "" : entityId);
        snapshot.put("metrics", metrics);
        snapshot.put("visual_payload", mergedVisual);
        snapshot.put("updated_at", now.toString());
        snapshots.put(step.key(), snapshot);
        jdbc.update("update demo_runs set step_snapshots=?::jsonb where id=?", Jsonb.write(snapshots), runId);
    }

    private Map<String, Object> objectMap(Object value) {
        if (value instanceof Map<?, ?> map) {
            Map<String, Object> result = new LinkedHashMap<>();
            map.forEach((key, item) -> result.put(String.valueOf(key), item));
            return result;
        }
        return new LinkedHashMap<>();
    }

    private Map<String, Object> deepMerge(Map<String, Object> base, Map<String, Object> patch) {
        Map<String, Object> merged = new LinkedHashMap<>(base);
        patch.forEach((key, value) -> {
            if (value instanceof Map<?, ?> && merged.get(key) instanceof Map<?, ?>) {
                merged.put(key, deepMerge(objectMap(merged.get(key)), objectMap(value)));
            } else {
                merged.put(key, value);
            }
        });
        return merged;
    }

    private void broadcast(String runId, Map<String, Object> event) {
        for (SseEmitter emitter : emitters.getOrDefault(runId, new CopyOnWriteArrayList<>())) {
            send(emitter, event);
        }
    }

    private void send(SseEmitter emitter, Map<String, Object> event) {
        try {
            emitter.send(SseEmitter.event()
                    .id(String.valueOf(event.get("event_id")))
                    .name("demo-run-event")
                    .data(event));
        } catch (IOException | IllegalStateException ignored) {
            emitters.values().forEach(list -> list.remove(emitter));
        }
    }

    private void removeEmitter(String runId, SseEmitter emitter) {
        CopyOnWriteArrayList<SseEmitter> runEmitters = emitters.get(runId);
        if (runEmitters != null) {
            runEmitters.remove(emitter);
        }
    }

    private static String runStatus(DemoStep step, String eventStatus) {
        if (step == DemoStep.RUN_COMPLETED) {
            return "completed";
        }
        if (step == DemoStep.RUN_FAILED || "failed".equals(eventStatus)) {
            return "failed";
        }
        if (step == DemoStep.RUN_CANCELLED || "cancelled".equals(eventStatus)) {
            return "cancelled";
        }
        return "running";
    }
}
