package org.raddatz.familienarchiv.ocr;

import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.raddatz.familienarchiv.ocr.BatchOcrDTO;
import org.raddatz.familienarchiv.ocr.OcrStatusDTO;
import org.raddatz.familienarchiv.ocr.TrainingHistoryResponse;
import org.raddatz.familienarchiv.ocr.TrainingInfoResponse;
import org.raddatz.familienarchiv.ocr.TriggerOcrDTO;
import org.raddatz.familienarchiv.ocr.TriggerSenderTrainingDTO;
import org.raddatz.familienarchiv.user.AppUser;
import org.raddatz.familienarchiv.ocr.OcrJob;
import org.raddatz.familienarchiv.ocr.OcrTrainingRun;
import org.raddatz.familienarchiv.security.Permission;
import org.raddatz.familienarchiv.security.RequirePermission;
import org.raddatz.familienarchiv.ocr.OcrBatchService;
import org.raddatz.familienarchiv.ocr.OcrProgressService;
import org.raddatz.familienarchiv.ocr.OcrService;
import org.raddatz.familienarchiv.ocr.OcrTrainingService;
import org.raddatz.familienarchiv.ocr.SegmentationTrainingExportService;
import org.raddatz.familienarchiv.ocr.SenderModelService;
import org.raddatz.familienarchiv.ocr.TrainingDataExportService;
import org.raddatz.familienarchiv.user.UserService;
import org.springframework.http.HttpHeaders;
import org.springframework.http.HttpStatus;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.security.core.Authentication;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.servlet.mvc.method.annotation.StreamingResponseBody;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import jakarta.validation.Valid;
import java.util.Map;
import java.util.UUID;

@RestController
@RequiredArgsConstructor
@Slf4j
public class OcrController {

    private final OcrService ocrService;
    private final OcrBatchService ocrBatchService;
    private final OcrProgressService ocrProgressService;
    private final UserService userService;
    private final TrainingDataExportService trainingDataExportService;
    private final SegmentationTrainingExportService segmentationTrainingExportService;
    private final OcrTrainingService ocrTrainingService;
    private final SenderModelService senderModelService;

    @PostMapping("/api/documents/{documentId}/ocr")
    @ResponseStatus(HttpStatus.ACCEPTED)
    @RequirePermission(Permission.WRITE_ALL)
    public Map<String, UUID> triggerOcr(
            @PathVariable UUID documentId,
            @RequestBody TriggerOcrDTO dto,
            Authentication authentication) {
        UUID userId = resolveUserId(authentication);
        UUID jobId = ocrService.startOcr(documentId, dto.getScriptType(), userId,
                Boolean.TRUE.equals(dto.getUseExistingAnnotations()));
        return Map.of("jobId", jobId);
    }

    @PostMapping("/api/ocr/batch")
    @ResponseStatus(HttpStatus.ACCEPTED)
    @RequirePermission(Permission.ADMIN)
    public Map<String, UUID> triggerBatch(
            @RequestBody @Valid BatchOcrDTO dto,
            Authentication authentication) {
        UUID userId = resolveUserId(authentication);
        UUID jobId = ocrBatchService.startBatch(dto.getDocumentIds(), userId);
        return Map.of("jobId", jobId);
    }

    @GetMapping("/api/ocr/jobs/{jobId}")
    @RequirePermission(Permission.READ_ALL)
    public OcrJob getJobStatus(@PathVariable UUID jobId) {
        return ocrService.getJob(jobId);
    }

    @GetMapping(value = "/api/ocr/jobs/{jobId}/progress", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    @RequirePermission(Permission.READ_ALL)
    public SseEmitter streamProgress(@PathVariable UUID jobId) {
        ocrService.getJob(jobId);
        return ocrProgressService.register(jobId);
    }

    @GetMapping("/api/documents/{documentId}/ocr-status")
    @RequirePermission(Permission.READ_ALL)
    public OcrStatusDTO getDocumentOcrStatus(@PathVariable UUID documentId) {
        return ocrService.getDocumentOcrStatus(documentId);
    }

    @GetMapping("/api/ocr/training-data/export")
    @RequirePermission(Permission.ADMIN)
    public ResponseEntity<StreamingResponseBody> exportTrainingData() {
        if (trainingDataExportService.queryEligibleBlocks().isEmpty()) {
            return ResponseEntity.noContent().build();
        }
        StreamingResponseBody body = trainingDataExportService.exportToZip();
        return ResponseEntity.ok()
                .contentType(MediaType.parseMediaType("application/zip"))
                .header(HttpHeaders.CONTENT_DISPOSITION, "attachment; filename=\"training-data.zip\"")
                .body(body);
    }

    @GetMapping("/api/ocr/segmentation-training-data/export")
    @RequirePermission(Permission.ADMIN)
    public ResponseEntity<StreamingResponseBody> exportSegmentationTrainingData() {
        if (segmentationTrainingExportService.querySegmentationBlocks().isEmpty()) {
            return ResponseEntity.noContent().build();
        }
        StreamingResponseBody body = segmentationTrainingExportService.exportToZip();
        return ResponseEntity.ok()
                .contentType(MediaType.parseMediaType("application/zip"))
                .header(HttpHeaders.CONTENT_DISPOSITION, "attachment; filename=\"segmentation-data.zip\"")
                .body(body);
    }

    @PostMapping("/api/ocr/train")
    @ResponseStatus(HttpStatus.CREATED)
    @RequirePermission(Permission.ADMIN)
    public OcrTrainingRun triggerTraining(Authentication authentication) {
        UUID userId = resolveUserId(authentication);
        return ocrTrainingService.triggerTraining(userId);
    }

    @PostMapping("/api/ocr/segtrain")
    @ResponseStatus(HttpStatus.CREATED)
    @RequirePermission(Permission.ADMIN)
    public OcrTrainingRun triggerSegTraining(Authentication authentication) {
        UUID userId = resolveUserId(authentication);
        return ocrTrainingService.triggerSegTraining(userId);
    }

    @GetMapping("/api/ocr/training-info")
    @RequirePermission(Permission.ADMIN)
    public TrainingInfoResponse getTrainingInfo() {
        return ocrTrainingService.getTrainingInfo();
    }

    @GetMapping("/api/ocr/training-info/global")
    @RequirePermission(Permission.ADMIN)
    public TrainingHistoryResponse getGlobalTrainingHistory() {
        return ocrTrainingService.getGlobalTrainingHistory();
    }

    @GetMapping("/api/ocr/training-info/{personId}")
    @RequirePermission(Permission.ADMIN)
    public TrainingHistoryResponse getSenderTrainingHistory(@PathVariable UUID personId) {
        return ocrTrainingService.getSenderTrainingHistory(personId);
    }

    @PostMapping("/api/ocr/train-sender")
    @ResponseStatus(HttpStatus.ACCEPTED)
    @RequirePermission(Permission.ADMIN)
    public OcrTrainingRun triggerSenderTraining(@Valid @RequestBody TriggerSenderTrainingDTO dto) {
        return senderModelService.triggerManualSenderTraining(dto.personId());
    }

    private UUID resolveUserId(Authentication authentication) {
        if (authentication == null || !authentication.isAuthenticated()) return null;
        try {
            AppUser user = userService.findByEmail(authentication.getName());
            return user != null ? user.getId() : null;
        } catch (Exception e) {
            log.warn("Failed to resolve user ID for authentication: {}", authentication.getName(), e);
            return null;
        }
    }
}
