package org.raddatz.familienarchiv.ocr;

import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.raddatz.familienarchiv.ocr.TrainingHistoryResponse;
import org.raddatz.familienarchiv.ocr.TrainingInfoResponse;
import org.raddatz.familienarchiv.exception.DomainException;
import org.raddatz.familienarchiv.ocr.OcrTrainingRun;
import org.raddatz.familienarchiv.person.Person;
import org.raddatz.familienarchiv.ocr.SenderModel;
import org.raddatz.familienarchiv.ocr.TrainingStatus;
import org.raddatz.familienarchiv.document.transcription.TranscriptionBlock;
import org.raddatz.familienarchiv.ocr.OcrTrainingRunRepository;
import org.raddatz.familienarchiv.person.PersonService;

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

import java.time.Instant;
import java.util.List;
import java.util.Optional;
import java.util.UUID;

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

class OcrTrainingServiceTest {

    OcrTrainingRunRepository runRepository;
    TrainingDataExportService exportService;
    SegmentationTrainingExportService segExportService;
    OcrClient ocrClient;
    OcrHealthClient healthClient;
    TranscriptionBlockQueryService transcriptionBlockQueryService;
    TransactionTemplate txTemplate;
    PersonService personService;
    SenderModelService senderModelService;
    OcrTrainingService service;

    @BeforeEach
    void setUp() {
        runRepository = mock(OcrTrainingRunRepository.class);
        exportService = mock(TrainingDataExportService.class);
        segExportService = mock(SegmentationTrainingExportService.class);
        ocrClient = mock(OcrClient.class);
        healthClient = mock(OcrHealthClient.class);
        transcriptionBlockQueryService = mock(TranscriptionBlockQueryService.class);
        txTemplate = mock(TransactionTemplate.class);
        personService = mock(PersonService.class);
        senderModelService = mock(SenderModelService.class);

        // Execute transaction callbacks inline so unit tests run without a real DataSource
        when(txTemplate.execute(any())).thenAnswer(inv -> {
            TransactionCallback<?> callback = inv.getArgument(0);
            return callback.doInTransaction(null);
        });

        service = new OcrTrainingService(runRepository, exportService, segExportService, ocrClient, healthClient, transcriptionBlockQueryService, txTemplate, personService, senderModelService);

        when(transcriptionBlockQueryService.count()).thenReturn(0L);
        when(runRepository.findTop20ByOrderByCreatedAtDesc()).thenReturn(List.of());
        when(segExportService.querySegmentationBlocks()).thenReturn(List.of());
        when(senderModelService.getAllSenderModels()).thenReturn(List.of());
    }

    // ─── Concurrent guard ─────────────────────────────────────────────────────

    @Test
    void triggerTraining_throws409_whenRunningRunExists() {
        when(runRepository.findFirstByStatus(TrainingStatus.RUNNING))
                .thenReturn(Optional.of(OcrTrainingRun.builder()
                        .id(UUID.randomUUID()).status(TrainingStatus.RUNNING)
                        .blockCount(5).documentCount(2).modelName("german_kurrent").build()));

        assertThatThrownBy(() -> service.triggerTraining(null))
                .isInstanceOf(DomainException.class)
                .extracting("status")
                .satisfies(s -> assertThat(s.toString()).contains("409"));
    }

    // ─── Threshold guard ──────────────────────────────────────────────────────

    @Test
    void triggerTraining_throws422_whenFewerThan5Blocks() {
        when(runRepository.findFirstByStatus(TrainingStatus.RUNNING)).thenReturn(Optional.empty());
        when(exportService.queryEligibleBlocks()).thenReturn(List.of(
                TranscriptionBlock.builder().id(UUID.randomUUID()).documentId(UUID.randomUUID()).build(),
                TranscriptionBlock.builder().id(UUID.randomUUID()).documentId(UUID.randomUUID()).build()
        ));

        assertThatThrownBy(() -> service.triggerTraining(null))
                .isInstanceOf(DomainException.class);
    }

    // ─── Happy path ───────────────────────────────────────────────────────────

    @Test
    void triggerTraining_createsRunWithCorrectCounts_andMarksDone() throws Exception {
        when(runRepository.findFirstByStatus(TrainingStatus.RUNNING)).thenReturn(Optional.empty());

        UUID docA = UUID.randomUUID();
        UUID docB = UUID.randomUUID();
        List<TranscriptionBlock> blocks = List.of(
                TranscriptionBlock.builder().id(UUID.randomUUID()).documentId(docA).build(),
                TranscriptionBlock.builder().id(UUID.randomUUID()).documentId(docA).build(),
                TranscriptionBlock.builder().id(UUID.randomUUID()).documentId(docA).build(),
                TranscriptionBlock.builder().id(UUID.randomUUID()).documentId(docA).build(),
                TranscriptionBlock.builder().id(UUID.randomUUID()).documentId(docB).build()
        );
        when(exportService.queryEligibleBlocks()).thenReturn(blocks);
        when(exportService.exportToZip()).thenReturn(out -> {});
        when(ocrClient.trainModel(any())).thenReturn(new OcrClient.TrainingResult(0.05, 0.95, 0.05, 3));

        OcrTrainingRun saved = OcrTrainingRun.builder()
                .id(UUID.randomUUID()).status(TrainingStatus.RUNNING)
                .blockCount(5).documentCount(2).modelName("german_kurrent").build();
        when(runRepository.save(any())).thenReturn(saved);

        service.triggerTraining(null);

        // Verify run created with correct counts and then updated to DONE
        verify(runRepository, times(2)).save(argThat(run ->
                run.getBlockCount() == 5 || run.getStatus() == TrainingStatus.DONE));
    }

    // ─── Failure path ─────────────────────────────────────────────────────────

    @Test
    void triggerTraining_marksRunFailed_whenOcrClientThrows() throws Exception {
        when(runRepository.findFirstByStatus(TrainingStatus.RUNNING)).thenReturn(Optional.empty());

        UUID docA = UUID.randomUUID();
        List<TranscriptionBlock> blocks = List.of(
                TranscriptionBlock.builder().id(UUID.randomUUID()).documentId(docA).build(),
                TranscriptionBlock.builder().id(UUID.randomUUID()).documentId(docA).build(),
                TranscriptionBlock.builder().id(UUID.randomUUID()).documentId(docA).build(),
                TranscriptionBlock.builder().id(UUID.randomUUID()).documentId(docA).build(),
                TranscriptionBlock.builder().id(UUID.randomUUID()).documentId(docA).build()
        );
        when(exportService.queryEligibleBlocks()).thenReturn(blocks);
        when(exportService.exportToZip()).thenReturn(out -> {});
        when(ocrClient.trainModel(any())).thenThrow(new RuntimeException("OCR service timeout"));

        OcrTrainingRun saved = OcrTrainingRun.builder()
                .id(UUID.randomUUID()).status(TrainingStatus.RUNNING)
                .blockCount(5).documentCount(1).modelName("german_kurrent").build();
        when(runRepository.save(any())).thenReturn(saved);

        service.triggerTraining(null);

        verify(runRepository, atLeastOnce()).save(argThat(run ->
                run.getStatus() == TrainingStatus.FAILED && run.getErrorMessage() != null));
    }

    // ─── triggerSegTraining ───────────────────────────────────────────────────

    @Test
    void triggerSegTraining_throws409_whenRunningRunExists() {
        when(runRepository.findFirstByStatus(TrainingStatus.RUNNING))
                .thenReturn(Optional.of(OcrTrainingRun.builder()
                        .id(UUID.randomUUID()).status(TrainingStatus.RUNNING)
                        .blockCount(5).documentCount(2).modelName("blla").build()));

        assertThatThrownBy(() -> service.triggerSegTraining(null))
                .isInstanceOf(DomainException.class)
                .extracting("status")
                .satisfies(s -> assertThat(s.toString()).contains("409"));
    }

    @Test
    void triggerSegTraining_throws422_whenFewerThan5Segments() {
        when(runRepository.findFirstByStatus(TrainingStatus.RUNNING)).thenReturn(Optional.empty());
        when(segExportService.querySegmentationBlocks()).thenReturn(List.of(
                TranscriptionBlock.builder().id(UUID.randomUUID()).documentId(UUID.randomUUID()).build(),
                TranscriptionBlock.builder().id(UUID.randomUUID()).documentId(UUID.randomUUID()).build()
        ));

        assertThatThrownBy(() -> service.triggerSegTraining(null))
                .isInstanceOf(DomainException.class);
    }

    @Test
    void triggerSegTraining_createsRunWithBlla_andMarksDoneWithCer() throws Exception {
        when(runRepository.findFirstByStatus(TrainingStatus.RUNNING)).thenReturn(Optional.empty());

        UUID docA = UUID.randomUUID();
        UUID docB = UUID.randomUUID();
        List<TranscriptionBlock> segs = List.of(
                TranscriptionBlock.builder().id(UUID.randomUUID()).documentId(docA).build(),
                TranscriptionBlock.builder().id(UUID.randomUUID()).documentId(docA).build(),
                TranscriptionBlock.builder().id(UUID.randomUUID()).documentId(docA).build(),
                TranscriptionBlock.builder().id(UUID.randomUUID()).documentId(docA).build(),
                TranscriptionBlock.builder().id(UUID.randomUUID()).documentId(docB).build()
        );
        when(segExportService.querySegmentationBlocks()).thenReturn(segs);
        when(segExportService.exportToZip()).thenReturn(out -> {});
        when(ocrClient.segtrainModel(any())).thenReturn(new OcrClient.TrainingResult(null, 0.92, 0.08, 5));

        OcrTrainingRun saved = OcrTrainingRun.builder()
                .id(UUID.randomUUID()).status(TrainingStatus.RUNNING)
                .blockCount(5).documentCount(2).modelName("blla").build();
        when(runRepository.save(any())).thenReturn(saved);

        service.triggerSegTraining(null);

        verify(runRepository, atLeastOnce()).save(argThat(run ->
                run.getStatus() == TrainingStatus.DONE
                        && "blla".equals(run.getModelName())
                        && run.getCer() != null));
    }

    @Test
    void triggerSegTraining_marksRunFailed_whenOcrClientThrows() throws Exception {
        when(runRepository.findFirstByStatus(TrainingStatus.RUNNING)).thenReturn(Optional.empty());

        UUID docA = UUID.randomUUID();
        List<TranscriptionBlock> segs = List.of(
                TranscriptionBlock.builder().id(UUID.randomUUID()).documentId(docA).build(),
                TranscriptionBlock.builder().id(UUID.randomUUID()).documentId(docA).build(),
                TranscriptionBlock.builder().id(UUID.randomUUID()).documentId(docA).build(),
                TranscriptionBlock.builder().id(UUID.randomUUID()).documentId(docA).build(),
                TranscriptionBlock.builder().id(UUID.randomUUID()).documentId(docA).build()
        );
        when(segExportService.querySegmentationBlocks()).thenReturn(segs);
        when(segExportService.exportToZip()).thenReturn(out -> {});
        when(ocrClient.segtrainModel(any())).thenThrow(new RuntimeException("seg timeout"));

        OcrTrainingRun saved = OcrTrainingRun.builder()
                .id(UUID.randomUUID()).status(TrainingStatus.RUNNING)
                .blockCount(5).documentCount(1).modelName("blla").build();
        when(runRepository.save(any())).thenReturn(saved);

        service.triggerSegTraining(null);

        verify(runRepository, atLeastOnce()).save(argThat(run ->
                run.getStatus() == TrainingStatus.FAILED && run.getErrorMessage() != null));
    }

    // ─── getTrainingInfo: batch person name resolution ────────────────────────

    @Test
    void getTrainingInfo_resolves_person_names_from_all_senderModels_in_batch() {
        UUID personA = UUID.randomUUID();
        UUID personB = UUID.randomUUID();
        List<SenderModel> models = List.of(
                SenderModel.builder().id(UUID.randomUUID()).personId(personA).correctedLinesAtTraining(100).build(),
                SenderModel.builder().id(UUID.randomUUID()).personId(personB).correctedLinesAtTraining(80).build()
        );
        when(senderModelService.getAllSenderModels()).thenReturn(models);
        when(exportService.queryEligibleBlocks()).thenReturn(List.of());

        Person pa = Person.builder().id(personA).firstName("Anna").lastName("Müller").build();
        Person pb = Person.builder().id(personB).firstName("Karl").lastName("Bauer").build();
        when(personService.getAllById(argThat(ids -> ids.size() == 2 && ids.containsAll(List.of(personA, personB)))))
                .thenReturn(List.of(pa, pb));
        when(healthClient.isHealthy()).thenReturn(true);

        TrainingInfoResponse info = service.getTrainingInfo();

        verify(personService, never()).getById(any());
        verify(personService, times(1)).getAllById(any());
        assertThat(info.personNames()).containsKey(personA.toString());
        assertThat(info.personNames()).containsKey(personB.toString());
    }

    @Test
    void getTrainingInfo_includesSenderModels_inResponse() {
        UUID personId = UUID.randomUUID();
        SenderModel model = SenderModel.builder()
                .id(UUID.randomUUID()).personId(personId).correctedLinesAtTraining(120).build();
        when(senderModelService.getAllSenderModels()).thenReturn(List.of(model));
        when(exportService.queryEligibleBlocks()).thenReturn(List.of());
        when(healthClient.isHealthy()).thenReturn(false);

        TrainingInfoResponse info = service.getTrainingInfo();

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

    // ─── Orphan recovery ──────────────────────────────────────────────────────

    @Test
    void recoverOrphanedRuns_marksRunFailed_whenOlderThanOneHour() {
        OcrTrainingRun orphan = OcrTrainingRun.builder()
                .id(UUID.randomUUID()).status(TrainingStatus.RUNNING)
                .blockCount(5).documentCount(1).modelName("german_kurrent")
                .createdAt(Instant.now().minusSeconds(7200))
                .build();
        when(runRepository.findFirstByStatus(TrainingStatus.RUNNING)).thenReturn(Optional.of(orphan));
        when(runRepository.save(any())).thenReturn(orphan);

        service.recoverOrphanedRuns();

        verify(runRepository).save(argThat(run ->
                run.getStatus() == TrainingStatus.FAILED
                        && run.getErrorMessage().contains("Abgebrochen")));
    }

    @Test
    void recoverOrphanedRuns_doesNothing_whenRunIsRecent() {
        OcrTrainingRun recent = OcrTrainingRun.builder()
                .id(UUID.randomUUID()).status(TrainingStatus.RUNNING)
                .blockCount(5).documentCount(1).modelName("german_kurrent")
                .createdAt(Instant.now().minusSeconds(60))
                .build();
        when(runRepository.findFirstByStatus(TrainingStatus.RUNNING)).thenReturn(Optional.of(recent));

        service.recoverOrphanedRuns();

        verify(runRepository, never()).save(any());
    }

    // ─── getGlobalTrainingHistory ─────────────────────────────────────────────

    @Test
    void getGlobalTrainingHistory_returnsOnlyGlobalRuns() {
        List<OcrTrainingRun> globalRuns = List.of(
                OcrTrainingRun.builder().id(UUID.randomUUID()).status(TrainingStatus.DONE)
                        .blockCount(10).documentCount(2).modelName("german_kurrent").build()
        );
        when(runRepository.findByPersonIdIsNullOrderByCreatedAtDesc()).thenReturn(globalRuns);

        TrainingHistoryResponse result = service.getGlobalTrainingHistory();

        assertThat(result.runs()).hasSize(1);
        assertThat(result.personNames()).isEmpty();
        verify(runRepository).findByPersonIdIsNullOrderByCreatedAtDesc();
    }

    // ─── getSenderTrainingHistory ─────────────────────────────────────────────

    @Test
    void getSenderTrainingHistory_includesPersonName_andRuns() {
        UUID personId = UUID.randomUUID();
        Person person = Person.builder().id(personId).firstName("Anna").lastName("Müller").build();
        when(personService.getById(personId)).thenReturn(person);

        List<OcrTrainingRun> senderRuns = List.of(
                OcrTrainingRun.builder().id(UUID.randomUUID()).status(TrainingStatus.DONE)
                        .personId(personId).blockCount(5).documentCount(1).modelName("sender_x").build()
        );
        when(runRepository.findByPersonIdOrderByCreatedAtDesc(personId)).thenReturn(senderRuns);

        TrainingHistoryResponse result = service.getSenderTrainingHistory(personId);

        assertThat(result.runs()).hasSize(1);
        assertThat(result.personNames()).containsKey(personId.toString());
    }

    @Test
    void getSenderTrainingHistory_propagates404_whenPersonNotFound() {
        UUID unknownId = UUID.randomUUID();
        when(personService.getById(unknownId))
                .thenThrow(DomainException.notFound(org.raddatz.familienarchiv.exception.ErrorCode.PERSON_NOT_FOUND, "not found"));

        assertThatThrownBy(() -> service.getSenderTrainingHistory(unknownId))
                .isInstanceOf(DomainException.class);
    }
}
