package org.raddatz.familienarchiv.ocr;

import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.ExtendWith;
import org.mockito.junit.jupiter.MockitoExtension;
import org.raddatz.familienarchiv.ocr.OcrTrainingRun;
import org.raddatz.familienarchiv.ocr.SenderModel;
import org.raddatz.familienarchiv.ocr.TrainingStatus;
import org.raddatz.familienarchiv.exception.DomainException;
import org.raddatz.familienarchiv.exception.ErrorCode;
import org.raddatz.familienarchiv.person.Person;
import org.raddatz.familienarchiv.ocr.OcrTrainingRunRepository;
import org.raddatz.familienarchiv.ocr.SenderModelRepository;
import org.raddatz.familienarchiv.person.PersonService;

import org.raddatz.familienarchiv.ocr.SenderModelService;
import org.raddatz.familienarchiv.document.transcription.TranscriptionBlockQueryService;
import org.springframework.test.util.ReflectionTestUtils;
import org.springframework.transaction.support.TransactionCallback;
import org.springframework.transaction.support.TransactionTemplate;

import java.util.Optional;
import java.util.UUID;

import static org.assertj.core.api.Assertions.assertThat;
import static org.mockito.ArgumentMatchers.*;
import static org.mockito.Mockito.*;

@ExtendWith(MockitoExtension.class)
class SenderModelServiceTest {

    SenderModelRepository senderModelRepository;
    TranscriptionBlockQueryService transcriptionBlockQueryService;
    OcrTrainingRunRepository trainingRunRepository;
    OcrClient ocrClient;
    TransactionTemplate txTemplate;
    TrainingDataExportService trainingDataExportService;
    PersonService personService;
    SenderModelService selfProxy;
    SenderModelService service;

    UUID personId = UUID.randomUUID();

    @BeforeEach
    void setUp() {
        senderModelRepository = mock(SenderModelRepository.class);
        transcriptionBlockQueryService = mock(TranscriptionBlockQueryService.class);
        trainingRunRepository = mock(OcrTrainingRunRepository.class);
        ocrClient = mock(OcrClient.class);
        txTemplate = mock(TransactionTemplate.class);
        trainingDataExportService = mock(TrainingDataExportService.class);
        personService = mock(PersonService.class);
        selfProxy = mock(SenderModelService.class);

        // Execute transaction callbacks inline so unit tests run without a real DataSource.
        // lenient: not every test hits the txTemplate path, but the setup is shared.
        lenient().when(txTemplate.execute(any())).thenAnswer(inv -> {
            TransactionCallback<?> callback = inv.getArgument(0);
            return callback.doInTransaction(null);
        });

        service = new SenderModelService(senderModelRepository, transcriptionBlockQueryService,
                trainingRunRepository, ocrClient, txTemplate, trainingDataExportService, personService);
        ReflectionTestUtils.setField(service, "self", selfProxy);
        ReflectionTestUtils.setField(service, "activationThreshold", 100);
        ReflectionTestUtils.setField(service, "retrainDelta", 50);
    }

    // ─── getAllSenderModels ───────────────────────────────────────────────────

    @Test
    void getAllSenderModels_returnsAllModelsFromRepository() {
        SenderModel model = SenderModel.builder()
                .id(UUID.randomUUID()).personId(personId).correctedLinesAtTraining(100).build();
        when(senderModelRepository.findAll()).thenReturn(java.util.List.of(model));

        java.util.List<SenderModel> result = service.getAllSenderModels();

        assertThat(result).hasSize(1);
        assertThat(result.get(0).getPersonId()).isEqualTo(personId);
    }

    // ─── runSenderTraining ────────────────────────────────────────────────────

    @Test
    void runSenderTraining_queriesBlockCountForPerson() {
        when(transcriptionBlockQueryService.countManualKurrentBlocksByPerson(personId)).thenReturn(42L);
        // triggerSenderTraining needs a RUNNING row — return empty to abort early
        when(trainingRunRepository.findFirstByPersonIdAndStatus(personId, TrainingStatus.RUNNING))
                .thenReturn(Optional.empty());

        try {
            service.runSenderTraining(personId);
        } catch (Exception ignored) {
            // triggerSenderTraining will throw when no RUNNING row found
        }

        verify(transcriptionBlockQueryService).countManualKurrentBlocksByPerson(personId);
    }

    // ─── Activation threshold ─────────────────────────────────────────────────

    @Test
    void checkAndTriggerTraining_doesNothing_belowActivationThreshold() {
        when(transcriptionBlockQueryService.countManualKurrentBlocksByPerson(personId)).thenReturn(99L);
        when(senderModelRepository.findByPersonId(personId)).thenReturn(Optional.empty());

        SenderModelService spy = spy(service);
        spy.checkAndTriggerTraining(personId);

        verify(spy, never()).runOrQueueSenderTraining(any(), anyInt());
    }

    @Test
    void checkAndTriggerTraining_triggersTraining_atActivationThreshold() {
        when(transcriptionBlockQueryService.countManualKurrentBlocksByPerson(personId)).thenReturn(100L);
        when(senderModelRepository.findByPersonId(personId)).thenReturn(Optional.empty());

        SenderModelService spy = spy(service);
        doReturn(false).when(spy).runOrQueueSenderTraining(personId, 100);

        spy.checkAndTriggerTraining(personId);

        verify(spy).runOrQueueSenderTraining(personId, 100);
    }

    // ─── Retrain delta ────────────────────────────────────────────────────────

    @Test
    void checkAndTriggerTraining_doesNothing_belowRetrainDelta() {
        SenderModel existing = SenderModel.builder().personId(personId)
                .correctedLinesAtTraining(100).build();
        when(senderModelRepository.findByPersonId(personId)).thenReturn(Optional.of(existing));
        when(transcriptionBlockQueryService.countManualKurrentBlocksByPerson(personId)).thenReturn(149L);

        SenderModelService spy = spy(service);
        spy.checkAndTriggerTraining(personId);

        verify(spy, never()).runOrQueueSenderTraining(any(), anyInt());
    }

    @Test
    void checkAndTriggerTraining_triggersRetrain_atRetrainDelta() {
        SenderModel existing = SenderModel.builder().personId(personId)
                .correctedLinesAtTraining(100).build();
        when(senderModelRepository.findByPersonId(personId)).thenReturn(Optional.of(existing));
        when(transcriptionBlockQueryService.countManualKurrentBlocksByPerson(personId)).thenReturn(150L);

        SenderModelService spy = spy(service);
        doReturn(false).when(spy).runOrQueueSenderTraining(personId, 150);

        spy.checkAndTriggerTraining(personId);

        verify(spy).runOrQueueSenderTraining(personId, 150);
    }

    // ─── runNow flag ──────────────────────────────────────────────────────────

    @Test
    void checkAndTriggerTraining_callsTrigger_whenRunNow() {
        when(transcriptionBlockQueryService.countManualKurrentBlocksByPerson(personId)).thenReturn(100L);
        when(senderModelRepository.findByPersonId(personId)).thenReturn(Optional.empty());

        SenderModelService spy = spy(service);
        doReturn(true).when(spy).runOrQueueSenderTraining(personId, 100);
        doNothing().when(spy).triggerSenderTraining(personId, 100);

        spy.checkAndTriggerTraining(personId);

        verify(spy).triggerSenderTraining(personId, 100);
    }

    @Test
    void checkAndTriggerTraining_doesNotCallTrigger_whenQueued() {
        when(transcriptionBlockQueryService.countManualKurrentBlocksByPerson(personId)).thenReturn(100L);
        when(senderModelRepository.findByPersonId(personId)).thenReturn(Optional.empty());

        SenderModelService spy = spy(service);
        doReturn(false).when(spy).runOrQueueSenderTraining(personId, 100);

        spy.checkAndTriggerTraining(personId);

        verify(spy, never()).triggerSenderTraining(any(), anyInt());
    }

    // ─── runOrQueueSenderTraining ─────────────────────────────────────────────

    @Test
    void runOrQueueSenderTraining_returnsTrue_whenIdle() {
        when(trainingRunRepository.existsByPersonIdAndStatus(personId, TrainingStatus.QUEUED)).thenReturn(false);
        when(trainingRunRepository.findFirstByStatus(TrainingStatus.RUNNING)).thenReturn(Optional.empty());

        boolean result = service.runOrQueueSenderTraining(personId, 120);

        assertThat(result).isTrue();
        verify(trainingRunRepository, never()).save(argThat(r -> r.getStatus() == TrainingStatus.QUEUED));
    }

    @Test
    void runOrQueueSenderTraining_savesQueued_whenRunning() {
        when(trainingRunRepository.existsByPersonIdAndStatus(personId, TrainingStatus.QUEUED)).thenReturn(false);
        when(trainingRunRepository.findFirstByStatus(TrainingStatus.RUNNING)).thenReturn(
                Optional.of(OcrTrainingRun.builder().id(UUID.randomUUID()).status(TrainingStatus.RUNNING)
                        .blockCount(5).documentCount(1).modelName("german_kurrent").build()));
        when(transcriptionBlockQueryService.countManualKurrentBlocksByPerson(personId)).thenReturn(120L);
        when(trainingRunRepository.save(any())).thenAnswer(inv -> inv.getArgument(0));

        boolean result = service.runOrQueueSenderTraining(personId, 120);

        assertThat(result).isFalse();
        verify(trainingRunRepository).save(argThat(r ->
                r.getStatus() == TrainingStatus.QUEUED && personId.equals(r.getPersonId())));
    }

    @Test
    void runOrQueueSenderTraining_returnsFalse_whenAlreadyQueued() {
        when(trainingRunRepository.existsByPersonIdAndStatus(personId, TrainingStatus.QUEUED)).thenReturn(true);

        boolean result = service.runOrQueueSenderTraining(personId, 120);

        assertThat(result).isFalse();
        verify(trainingRunRepository, never()).save(any());
    }

    @Test
    void runOrQueueSenderTraining_createsRunningRow_atomically() {
        // RUNNING row must be created inside the same @Transactional boundary as the idle check,
        // eliminating the race window between the check and a separate triggerSenderTraining call.
        when(trainingRunRepository.existsByPersonIdAndStatus(personId, TrainingStatus.QUEUED)).thenReturn(false);
        when(trainingRunRepository.findFirstByStatus(TrainingStatus.RUNNING)).thenReturn(Optional.empty());
        when(transcriptionBlockQueryService.countManualKurrentBlocksByPerson(personId)).thenReturn(120L);
        when(trainingRunRepository.save(any())).thenAnswer(inv -> {
            OcrTrainingRun r = inv.getArgument(0);
            if (r.getId() == null) r.setId(UUID.randomUUID());
            return r;
        });

        boolean result = service.runOrQueueSenderTraining(personId, 120);

        assertThat(result).isTrue();
        verify(trainingRunRepository).save(argThat(r ->
                r.getStatus() == TrainingStatus.RUNNING && personId.equals(r.getPersonId())));
    }

    // ─── triggerSenderTraining — happy path ──────────────────────────────────

    @Test
    void triggerSenderTraining_savesModelRecord_onSuccess() throws Exception {
        OcrTrainingRun runningRun = OcrTrainingRun.builder()
                .id(UUID.randomUUID()).status(TrainingStatus.RUNNING)
                .personId(personId).blockCount(10).documentCount(0)
                .modelName("sender_" + personId).build();
        when(trainingRunRepository.findFirstByPersonIdAndStatus(personId, TrainingStatus.RUNNING))
                .thenReturn(Optional.of(runningRun));
        when(trainingDataExportService.exportForSender(personId)).thenReturn(out -> {});
        when(ocrClient.trainSenderModel(any(), any()))
                .thenReturn(new OcrClient.TrainingResult(0.08, 0.92, 0.08, 5));
        when(senderModelRepository.findByPersonId(personId)).thenReturn(Optional.empty());
        when(senderModelRepository.save(any())).thenAnswer(inv -> inv.getArgument(0));
        when(trainingRunRepository.save(any())).thenAnswer(inv -> inv.getArgument(0));
        when(trainingRunRepository.findFirstByStatusOrderByCreatedAtAsc(TrainingStatus.QUEUED))
                .thenReturn(Optional.empty());

        service.triggerSenderTraining(personId, 120);

        verify(senderModelRepository).save(argThat(m ->
                ("/app/models/sender_" + personId + ".mlmodel").equals(m.getModelPath())
                        && Double.valueOf(0.08).equals(m.getCer())
                        && personId.equals(m.getPersonId())));
        verify(trainingRunRepository, atLeastOnce()).save(argThat(r ->
                r.getStatus() == TrainingStatus.DONE
                        && Double.valueOf(0.08).equals(r.getCer())
                        && r.getCompletedAt() != null));
    }

    // ─── triggerSenderTraining — failure path ────────────────────────────────

    @Test
    void triggerSenderTraining_marksRunFailed_whenOcrClientThrows() throws Exception {
        OcrTrainingRun runningRun = OcrTrainingRun.builder()
                .id(UUID.randomUUID()).status(TrainingStatus.RUNNING)
                .personId(personId).blockCount(10).documentCount(0)
                .modelName("sender_" + personId).build();
        when(trainingRunRepository.findFirstByPersonIdAndStatus(personId, TrainingStatus.RUNNING))
                .thenReturn(Optional.of(runningRun));
        when(trainingDataExportService.exportForSender(personId)).thenReturn(out -> {});
        when(ocrClient.trainSenderModel(any(), any()))
                .thenThrow(new RuntimeException("OCR service timeout"));
        when(trainingRunRepository.save(any())).thenAnswer(inv -> inv.getArgument(0));
        when(trainingRunRepository.findFirstByStatusOrderByCreatedAtAsc(TrainingStatus.QUEUED))
                .thenReturn(Optional.empty());

        service.triggerSenderTraining(personId, 120);

        verify(trainingRunRepository, atLeastOnce()).save(argThat(r ->
                r.getStatus() == TrainingStatus.FAILED
                        && r.getErrorMessage() != null
                        && r.getCompletedAt() != null));
        verify(senderModelRepository, never()).save(any());
    }

    @Test
    void triggerSenderTraining_throwsDomainException_whenRunningRowMissingAfterDispatch() {
        when(trainingRunRepository.findFirstByPersonIdAndStatus(personId, TrainingStatus.RUNNING))
                .thenReturn(Optional.empty());

        org.assertj.core.api.Assertions.assertThatThrownBy(
                () -> service.triggerSenderTraining(personId, 0))
                .isInstanceOf(DomainException.class);
    }

    // ─── triggerSenderTraining — queue promotion ──────────────────────────────

    // ─── triggerManualSenderTraining ──────────────────────────────────────────

    @Test
    void triggerManualSenderTraining_returnsRunningRun_whenIdle() {
        when(personService.getById(personId)).thenReturn(Person.builder().id(personId).build());
        when(transcriptionBlockQueryService.countManualKurrentBlocksByPerson(personId)).thenReturn(0L);
        when(trainingRunRepository.existsByPersonIdAndStatus(personId, TrainingStatus.QUEUED)).thenReturn(false);
        when(trainingRunRepository.findFirstByStatus(TrainingStatus.RUNNING)).thenReturn(Optional.empty());
        OcrTrainingRun runningRun = OcrTrainingRun.builder()
                .id(UUID.randomUUID()).status(TrainingStatus.RUNNING)
                .personId(personId).blockCount(0).documentCount(0).modelName("sender_" + personId).build();
        when(trainingRunRepository.save(any())).thenReturn(runningRun);
        when(trainingRunRepository.findFirstByPersonIdAndStatus(personId, TrainingStatus.RUNNING))
                .thenReturn(Optional.of(runningRun));

        OcrTrainingRun result = service.triggerManualSenderTraining(personId);

        assertThat(result.getStatus()).isEqualTo(TrainingStatus.RUNNING);
        assertThat(result.getPersonId()).isEqualTo(personId);
    }

    @Test
    void triggerManualSenderTraining_returnsQueuedRun_whenAnotherRunning() {
        when(personService.getById(personId)).thenReturn(Person.builder().id(personId).build());
        when(transcriptionBlockQueryService.countManualKurrentBlocksByPerson(personId)).thenReturn(0L);
        when(trainingRunRepository.existsByPersonIdAndStatus(personId, TrainingStatus.QUEUED)).thenReturn(false);
        when(trainingRunRepository.findFirstByStatus(TrainingStatus.RUNNING)).thenReturn(
                Optional.of(OcrTrainingRun.builder().id(UUID.randomUUID()).status(TrainingStatus.RUNNING)
                        .blockCount(1).documentCount(0).modelName("german_kurrent").build()));
        OcrTrainingRun queuedRun = OcrTrainingRun.builder()
                .id(UUID.randomUUID()).status(TrainingStatus.QUEUED)
                .personId(personId).blockCount(0).documentCount(0).modelName("sender_" + personId).build();
        when(trainingRunRepository.save(any())).thenReturn(queuedRun);
        when(trainingRunRepository.findFirstByPersonIdAndStatus(personId, TrainingStatus.QUEUED))
                .thenReturn(Optional.of(queuedRun));

        OcrTrainingRun result = service.triggerManualSenderTraining(personId);

        assertThat(result.getStatus()).isEqualTo(TrainingStatus.QUEUED);
    }

    @Test
    void triggerManualSenderTraining_throws404_whenPersonNotFound() {
        when(personService.getById(personId))
                .thenThrow(DomainException.notFound(ErrorCode.PERSON_NOT_FOUND, "Person not found"));

        org.assertj.core.api.Assertions.assertThatThrownBy(
                () -> service.triggerManualSenderTraining(personId))
                .isInstanceOf(DomainException.class);
    }

    @Test
    void triggerManualSenderTraining_throwsDomainException_whenRunRowMissingAfterCreate() {
        when(personService.getById(personId)).thenReturn(Person.builder().id(personId).build());
        when(transcriptionBlockQueryService.countManualKurrentBlocksByPerson(personId)).thenReturn(0L);
        when(trainingRunRepository.existsByPersonIdAndStatus(personId, TrainingStatus.QUEUED)).thenReturn(false);
        when(trainingRunRepository.findFirstByStatus(TrainingStatus.RUNNING)).thenReturn(Optional.empty());
        OcrTrainingRun runningRun = OcrTrainingRun.builder()
                .id(UUID.randomUUID()).status(TrainingStatus.RUNNING)
                .personId(personId).blockCount(0).documentCount(0).modelName("sender_" + personId).build();
        when(trainingRunRepository.save(any())).thenReturn(runningRun);
        // Simulate the run row not being found after creation (defensive path)
        when(trainingRunRepository.findFirstByPersonIdAndStatus(personId, TrainingStatus.RUNNING))
                .thenReturn(Optional.empty());

        org.assertj.core.api.Assertions.assertThatThrownBy(
                () -> service.triggerManualSenderTraining(personId))
                .isInstanceOf(DomainException.class);
    }

    @Test
    void triggerSenderTraining_promotesNextQueued_afterCompletion() throws Exception {
        UUID nextPersonId = UUID.randomUUID();

        OcrTrainingRun runningRun = OcrTrainingRun.builder()
                .id(UUID.randomUUID()).status(TrainingStatus.RUNNING)
                .personId(personId).blockCount(10).documentCount(0)
                .modelName("sender_" + personId).build();
        when(trainingRunRepository.findFirstByPersonIdAndStatus(personId, TrainingStatus.RUNNING))
                .thenReturn(Optional.of(runningRun));
        when(trainingRunRepository.save(any())).thenAnswer(inv -> {
            OcrTrainingRun r = inv.getArgument(0);
            if (r.getId() == null) r.setId(UUID.randomUUID());
            return r;
        });
        when(trainingDataExportService.exportForSender(personId)).thenReturn(out -> {});
        when(ocrClient.trainSenderModel(any(), any()))
                .thenReturn(new OcrClient.TrainingResult(0.05, 0.95, 0.05, 3));
        when(senderModelRepository.findByPersonId(personId)).thenReturn(Optional.empty());

        OcrTrainingRun queued = OcrTrainingRun.builder()
                .id(UUID.randomUUID()).status(TrainingStatus.QUEUED)
                .personId(nextPersonId).blockCount(5).documentCount(0)
                .modelName("sender_" + nextPersonId).build();
        when(trainingRunRepository.findFirstByStatusOrderByCreatedAtAsc(TrainingStatus.QUEUED))
                .thenReturn(Optional.of(queued));
        when(transcriptionBlockQueryService.countManualKurrentBlocksByPerson(nextPersonId)).thenReturn(5L);

        SenderModelService spy = spy(service);
        // Stub the recursive call to stop the chain after one promotion
        doNothing().when(spy).triggerSenderTraining(eq(nextPersonId), anyInt());

        spy.triggerSenderTraining(personId, 10);

        // Queued run promoted to RUNNING
        verify(trainingRunRepository).save(argThat(r ->
                nextPersonId.equals(r.getPersonId()) && r.getStatus() == TrainingStatus.RUNNING));
        // Next training was kicked off
        verify(spy).triggerSenderTraining(eq(nextPersonId), anyInt());
    }
}
