package com.zhengmeng.ocrplatform.worker;

import com.zhengmeng.ocrplatform.document.DocumentStorageService;
import com.zhengmeng.ocrplatform.document.OcrDocumentEntity;
import com.zhengmeng.ocrplatform.document.OcrDocumentRepository;
import com.zhengmeng.ocrplatform.engine.OcrEngineResult;
import com.zhengmeng.ocrplatform.engine.OcrEngineAdapter;
import com.zhengmeng.ocrplatform.engine.RecognizedKeyValue;
import com.zhengmeng.ocrplatform.engine.RecognizedTextBlock;
import com.zhengmeng.ocrplatform.extraction.OcrKeyValueEntity;
import com.zhengmeng.ocrplatform.extraction.OcrKeyValueRepository;
import com.zhengmeng.ocrplatform.recognition.OcrRawResultEntity;
import com.zhengmeng.ocrplatform.recognition.OcrRawResultRepository;
import com.zhengmeng.ocrplatform.recognition.OcrTextBlockEntity;
import com.zhengmeng.ocrplatform.recognition.OcrTextBlockRepository;
import com.zhengmeng.ocrplatform.task.OcrTaskEntity;
import com.zhengmeng.ocrplatform.task.OcrTaskRepository;
import com.zhengmeng.ocrplatform.task.TaskEventService;
import com.zhengmeng.ocrplatform.task.TaskStatus;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;

import java.time.Duration;
import java.time.LocalDateTime;
import java.util.List;
import java.util.Optional;

@Service
public class OcrTaskExecutionService {
    private final OcrTaskRepository taskRepository;
    private final OcrDocumentRepository documentRepository;
    private final DocumentStorageService documentStorageService;
    private final OcrRawResultRepository rawResultRepository;
    private final OcrTextBlockRepository textBlockRepository;
    private final OcrKeyValueRepository keyValueRepository;
    private final TaskEventService taskEventService;
    private final WorkerProperties properties;

    public OcrTaskExecutionService(OcrTaskRepository taskRepository,
                                   OcrDocumentRepository documentRepository,
                                   DocumentStorageService documentStorageService,
                                   OcrRawResultRepository rawResultRepository,
                                   OcrTextBlockRepository textBlockRepository,
                                   OcrKeyValueRepository keyValueRepository,
                                   TaskEventService taskEventService,
                                   WorkerProperties properties) {
        this.taskRepository = taskRepository;
        this.documentRepository = documentRepository;
        this.documentStorageService = documentStorageService;
        this.rawResultRepository = rawResultRepository;
        this.textBlockRepository = textBlockRepository;
        this.keyValueRepository = keyValueRepository;
        this.taskEventService = taskEventService;
        this.properties = properties;
    }

    @Transactional
    public Optional<TaskExecutionContext> claim(String taskId) {
        LocalDateTime now = LocalDateTime.now();
        int changed = taskRepository.markRecognizing(taskId, TaskStatus.QUEUED, TaskStatus.RECOGNIZING, now);
        if (changed == 0) {
            return Optional.empty();
        }

        taskEventService.record(
                taskId,
                "TASK_RECOGNIZING",
                TaskStatus.QUEUED,
                TaskStatus.RECOGNIZING,
                "OCR worker claimed task",
                null,
                now
        );

        List<OcrDocumentEntity> documents = documentRepository.findByTaskIdOrderByIdAsc(taskId);
        if (documents.isEmpty()) {
            throw new IllegalStateException("Task has no document: " + taskId);
        }

        OcrDocumentEntity document = documents.getFirst();
        OcrTaskEntity task = taskRepository.findByTaskId(taskId)
                .orElseThrow(() -> new IllegalStateException("Task not found: " + taskId));

        return Optional.of(new TaskExecutionContext(
                taskId,
                document,
                documentStorageService.resolve(document.getObjectPath()),
                task.getAttemptCount() + 1
        ));
    }

    @Transactional
    public void complete(String taskId,
                         int runNo,
                         OcrEngineResult result,
                         String requestPayloadJson,
                         Long elapsedMs) {
        LocalDateTime now = LocalDateTime.now();

        int changed = taskRepository.markCompleted(taskId, TaskStatus.RECOGNIZING, TaskStatus.COMPLETED, now);
        if (changed == 0) {
            taskEventService.record(
                    taskId,
                    "TASK_RESULT_DISCARDED",
                    null,
                    null,
                    "OCR completion skipped because task state changed before persistence",
                    null,
                    now
            );
            return;
        }

        OcrRawResultEntity rawResult = new OcrRawResultEntity();
        rawResult.setTaskId(taskId);
        rawResult.setRunNo(runNo);
        rawResult.setEngineCode(result.engineCode());
        rawResult.setEngineVersion(result.engineVersion());
        rawResult.setResultType("OCR_ENGINE_RAW_JSON");
        rawResult.setRawJson(result.rawJson());
        rawResult.setRequestPayloadJson(requestPayloadJson);
        rawResult.setElapsedMs(elapsedMs);
        rawResult.setCreatedAt(now);
        rawResultRepository.save(rawResult);

        for (RecognizedTextBlock textBlock : result.textBlocks()) {
            OcrTextBlockEntity entity = new OcrTextBlockEntity();
            entity.setTaskId(taskId);
            entity.setRunNo(runNo);
            entity.setPageNo(textBlock.pageNo());
            entity.setBlockType(textBlock.blockType());
            entity.setTextContent(textBlock.textContent());
            entity.setConfidence(textBlock.confidence());
            entity.setBboxJson(textBlock.bboxJson());
            entity.setReadingOrder(textBlock.readingOrder());
            entity.setSourceEngine(result.engineCode());
            entity.setCreatedAt(now);
            textBlockRepository.save(entity);
        }

        for (RecognizedKeyValue keyValue : result.keyValues()) {
            OcrKeyValueEntity entity = new OcrKeyValueEntity();
            entity.setTaskId(taskId);
            entity.setRunNo(runNo);
            entity.setFieldKey(keyValue.fieldKey());
            entity.setFieldName(keyValue.fieldName());
            entity.setValueText(keyValue.valueText());
            entity.setNormalizedValue(keyValue.normalizedValue());
            entity.setConfidence(keyValue.confidence());
            entity.setReviewStatus("PENDING");
            entity.setPageNo(keyValue.pageNo());
            entity.setSourceEngine(result.engineCode());
            entity.setEvidenceJson(keyValue.evidenceJson());
            entity.setCreatedAt(now);
            entity.setUpdatedAt(now);
            keyValueRepository.save(entity);
        }

        taskEventService.record(
                taskId,
                "TASK_COMPLETED",
                TaskStatus.RECOGNIZING,
                TaskStatus.COMPLETED,
                "OCR worker completed task",
                null,
                now
        );
    }

    @Transactional
    public void fail(String taskId,
                     int runNo,
                     OcrEngineAdapter adapter,
                     Exception exception,
                     String requestPayloadJson,
                     Long elapsedMs) {
        LocalDateTime now = LocalDateTime.now();
        String message = exception.getMessage() == null ? exception.getClass().getName() : exception.getMessage();
        if (message.length() > 1000) {
            message = message.substring(0, 1000);
        }
        String errorCode = normalizeErrorCode(exception);
        String engineCode = adapter == null ? "unknown-engine" : adapter.engineCode();
        String engineVersion = adapter == null ? null : adapter.engineVersion();

        OcrTaskEntity task = taskRepository.findByTaskId(taskId)
                .orElseThrow(() -> new IllegalStateException("Task not found: " + taskId));
        if (task.getStatus() != TaskStatus.RECOGNIZING) {
            return;
        }

        OcrRawResultEntity rawResult = new OcrRawResultEntity();
        rawResult.setTaskId(taskId);
        rawResult.setRunNo(runNo);
        rawResult.setEngineCode(engineCode);
        rawResult.setEngineVersion(engineVersion);
        rawResult.setResultType("OCR_ENGINE_ERROR");
        rawResult.setRawJson(message);
        rawResult.setRequestPayloadJson(requestPayloadJson);
        rawResult.setElapsedMs(elapsedMs);
        rawResult.setErrorCode(errorCode);
        rawResult.setErrorMessage(message);
        rawResult.setCreatedAt(now);
        rawResultRepository.save(rawResult);

        if (task.getAttemptCount() + 1 <= properties.maxRetryAttempts()) {
            int retried = taskRepository.markRetryQueued(
                    taskId,
                    TaskStatus.RECOGNIZING,
                    TaskStatus.QUEUED,
                    "WORKER_ERROR",
                    message,
                    now
            );
            if (retried > 0) {
                taskEventService.record(
                        taskId,
                        "TASK_RETRY_SCHEDULED",
                        TaskStatus.RECOGNIZING,
                        TaskStatus.QUEUED,
                        "OCR task failed and will be retried",
                        null,
                        now
                );
            }
            return;
        }

        int failed = taskRepository.markFailed(
                taskId,
                TaskStatus.RECOGNIZING,
                TaskStatus.FAILED,
                "WORKER_ERROR",
                message,
                now
        );
        if (failed > 0) {
            taskEventService.record(
                    taskId,
                    "TASK_FAILED",
                    TaskStatus.RECOGNIZING,
                    TaskStatus.FAILED,
                message,
                null,
                now
            );
        }
    }

    @Transactional
    public int recoverTimedOutRecognizingTasks() {
        LocalDateTime now = LocalDateTime.now();
        LocalDateTime cutoff = now.minus(Duration.ofMillis(properties.recognizingTimeoutMs()));
        List<OcrTaskEntity> timedOutTasks = taskRepository
                .findTop50ByStatusAndStartedAtBeforeOrderByStartedAtAsc(TaskStatus.RECOGNIZING, cutoff);

        for (OcrTaskEntity task : timedOutTasks) {
            String taskId = task.getTaskId();
            if (task.getAttemptCount() + 1 <= properties.maxRetryAttempts()) {
                int changed = taskRepository.markRetryQueued(
                        taskId,
                        TaskStatus.RECOGNIZING,
                        TaskStatus.QUEUED,
                        "TASK_TIMEOUT",
                        "Task timed out in recognition state and was requeued",
                        now
                );
                if (changed > 0) {
                    taskEventService.record(
                            taskId,
                            "TASK_TIMEOUT_REQUEUED",
                            TaskStatus.RECOGNIZING,
                            TaskStatus.QUEUED,
                            "OCR task timed out and has been requeued",
                            null,
                            now
                    );
                }
                continue;
            }

            int changed = taskRepository.markFailed(
                    taskId,
                    TaskStatus.RECOGNIZING,
                    TaskStatus.FAILED,
                    "TASK_TIMEOUT",
                    "Task timed out and retry budget has been exhausted",
                    now
            );
            if (changed > 0) {
                taskEventService.record(
                        taskId,
                        "TASK_TIMEOUT_FAILED",
                        TaskStatus.RECOGNIZING,
                        TaskStatus.FAILED,
                        "OCR task timed out without completing",
                        null,
                        now
                );
            }
        }

        return timedOutTasks.size();
    }

    private String normalizeErrorCode(Exception exception) {
        if (exception == null) {
            return "TASK_PROCESSING_ERROR";
        }
        if (exception instanceof RuntimeException) {
            return "TASK_PROCESSING_ERROR";
        }
        return exception.getClass().getSimpleName();
    }

}