package com.ai.trackwalker.model;

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.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.IOException;
import java.math.BigDecimal;
import java.sql.Timestamp;
import java.time.Instant;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CopyOnWriteArrayList;

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

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

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

    public SseEmitter subscribe(String taskId, long afterSequence) {
        SseEmitter emitter = new SseEmitter(SSE_TIMEOUT_MS);
        subscribers.computeIfAbsent(taskId, ignored -> new CopyOnWriteArrayList<>()).add(emitter);
        emitter.onCompletion(() -> remove(taskId, emitter));
        emitter.onTimeout(() -> remove(taskId, emitter));
        emitter.onError(error -> remove(taskId, emitter));
        try {
            for (Map<String, Object> event : events(taskId, afterSequence)) {
                send(emitter, event);
            }
            emitter.send(SseEmitter.event().name("connected").data(Map.of("task_id", taskId)));
        } catch (IOException e) {
            remove(taskId, emitter);
        }
        return emitter;
    }

    public synchronized Map<String, Object> publish(
            String taskId,
            String eventType,
            String stepKey,
            String status,
            double progress,
            String message,
            Map<String, Object> payload
    ) {
        Long sequence = jdbc.queryForObject(
                "select coalesce(max(sequence_no), 0) + 1 from model_test_events where test_task_id=?",
                Long.class,
                taskId
        );
        long resolvedSequence = sequence == null ? 1L : sequence;
        String id = Ids.next("mte");
        Instant now = Instant.now();
        jdbc.update("""
                        insert into model_test_events(id, test_task_id, sequence_no, event_type, step_key, status, progress, message, payload, created_at)
                        values (?,?,?,?,?,?,?, ?,?::jsonb,?)
                        """,
                id, taskId, resolvedSequence, eventType, stepKey, status,
                BigDecimal.valueOf(progress), message, Jsonb.write(payload == null ? Map.of() : payload), Timestamp.from(now));
        Map<String, Object> event = new HashMap<>();
        event.put("event_id", id);
        event.put("task_id", taskId);
        event.put("sequence_no", resolvedSequence);
        event.put("event_type", eventType);
        event.put("step_key", stepKey);
        event.put("status", status);
        event.put("progress", progress);
        event.put("message", message);
        event.put("payload", payload == null ? Map.of() : payload);
        event.put("created_at", now.toString());
        notifySubscribers(taskId, event);
        return event;
    }

    public List<Map<String, Object>> events(String taskId, long afterSequence) {
        List<Map<String, Object>> rows = jdbc.queryForList("""
                select id as event_id, test_task_id as task_id, sequence_no, event_type, step_key,
                       status, progress, message, payload::text as payload, created_at
                from model_test_events
                where test_task_id=? and sequence_no>?
                order by sequence_no
                """, taskId, afterSequence);
        List<Map<String, Object>> result = new ArrayList<>();
        for (Map<String, Object> source : rows) {
            Map<String, Object> row = new HashMap<>(source);
            row.put("payload", Jsonb.map(String.valueOf(row.get("payload"))));
            result.add(row);
        }
        return result;
    }

    public void complete(String taskId) {
        List<SseEmitter> emitters = subscribers.remove(taskId);
        if (emitters != null) {
            emitters.forEach(SseEmitter::complete);
        }
    }

    private void notifySubscribers(String taskId, Map<String, Object> event) {
        List<SseEmitter> emitters = subscribers.get(taskId);
        if (emitters == null) {
            return;
        }
        for (SseEmitter emitter : emitters) {
            try {
                send(emitter, event);
            } catch (IOException e) {
                remove(taskId, emitter);
            }
        }
    }

    private void send(SseEmitter emitter, Map<String, Object> event) throws IOException {
        emitter.send(SseEmitter.event()
                .id(String.valueOf(event.get("sequence_no")))
                .name(String.valueOf(event.get("event_type")))
                .data(event));
    }

    private void remove(String taskId, SseEmitter emitter) {
        List<SseEmitter> emitters = subscribers.get(taskId);
        if (emitters != null) {
            emitters.remove(emitter);
            if (emitters.isEmpty()) {
                subscribers.remove(taskId);
            }
        }
    }
}
