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.kafka.core.KafkaTemplate;
import org.springframework.stereotype.Component;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

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

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

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

    public ModelArtifactEventPublisher(JdbcTemplate jdbc, KafkaTemplate<String, Object> kafka) {
        this.jdbc = jdbc;
        this.kafka = kafka;
    }

    public Map<String, Object> publish(
            String jobId,
            String stage,
            String status,
            double progress,
            String message,
            Map<String, Object> payload
    ) {
        Long sequence = jdbc.queryForObject(
                "select coalesce(max(sequence_no),0)+1 from model_artifact_job_events where job_id=?",
                Long.class, jobId);
        long next = sequence == null ? 1L : sequence;
        Instant now = Instant.now();
        String eventId = Ids.next("mae");
        jdbc.update("""
                insert into model_artifact_job_events(
                    id, job_id, sequence_no, stage, status, progress, message, payload, created_at
                ) values (?,?,?,?,?,?,?,?::jsonb,?)
                """, eventId, jobId, next, stage, status, progress, message,
                Jsonb.write(payload == null ? Map.of() : payload), Timestamp.from(now));
        jdbc.update("""
                update model_artifact_jobs set status=?, progress=?, current_stage=?, updated_at=? where id=?
                """, status, progress, stage, Timestamp.from(now), jobId);
        Map<String, Object> event = Map.of(
                "id", eventId,
                "job_id", jobId,
                "sequence_no", next,
                "stage", stage,
                "status", status,
                "progress", progress,
                "message", message,
                "payload", payload == null ? Map.of() : payload,
                "created_at", now.toString()
        );
        try {
            kafka.send("model.artifact.job.progress", jobId, event);
        } catch (Exception ignored) {
            // Database and SSE remain authoritative when Kafka is temporarily unavailable.
        }
        for (SseEmitter emitter : subscribers.getOrDefault(jobId, new CopyOnWriteArrayList<>())) {
            send(emitter, event);
        }
        return event;
    }

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

    public List<Map<String, Object>> history(String jobId, long afterSequence) {
        List<Map<String, Object>> rows = jdbc.queryForList("""
                select id, job_id, sequence_no, stage, status, progress, message,
                       payload::text as payload, created_at
                from model_artifact_job_events where job_id=? and sequence_no>?
                order by sequence_no
                """, jobId, afterSequence);
        rows.forEach(row -> row.put("payload", Jsonb.map(String.valueOf(row.get("payload")))));
        return rows;
    }

    private void send(SseEmitter emitter, Map<String, Object> event) {
        try {
            emitter.send(SseEmitter.event().name("artifact-job").data(event));
        } catch (IOException e) {
            emitter.complete();
        }
    }

    private void remove(String jobId, SseEmitter emitter) {
        CopyOnWriteArrayList<SseEmitter> list = subscribers.get(jobId);
        if (list != null) {
            list.remove(emitter);
            if (list.isEmpty()) {
                subscribers.remove(jobId);
            }
        }
    }
}
