package com.zhengmeng.ocrplatform.engine;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.http.MediaType;
import org.springframework.http.client.SimpleClientHttpRequestFactory;
import org.springframework.stereotype.Component;
import org.springframework.web.client.RestClient;
import org.springframework.web.client.RestClientException;
import org.springframework.web.client.RestClientResponseException;

import java.io.IOException;
import java.math.BigDecimal;
import java.nio.file.Files;
import java.util.ArrayList;
import java.util.Base64;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;

@Component
@ConditionalOnProperty(prefix = "ocr-platform.engine.paddle-layout", name = "enabled", havingValue = "true")
public class PaddleLayoutParsingEngineAdapter implements OcrEngineAdapter {
    private final EngineProperties properties;
    private final ObjectMapper objectMapper = new ObjectMapper();
    private final RestClient restClient;

    public PaddleLayoutParsingEngineAdapter(EngineProperties properties) {
        this.properties = properties;
        SimpleClientHttpRequestFactory requestFactory = new SimpleClientHttpRequestFactory();
        requestFactory.setConnectTimeout(properties.paddleLayout().timeout());
        requestFactory.setReadTimeout(properties.paddleLayout().timeout());
        this.restClient = RestClient.builder()
                .requestFactory(requestFactory)
                .build();
    }

    @Override
    public String engineCode() {
        return "paddle-layout";
    }

    @Override
    public String engineVersion() {
        return "PaddleOCR layout-parsing";
    }

    @Override
    public OcrEngineResult recognize(OcrEngineRequest request) {
        try {
            String responseBody = restClient.post()
                    .uri(properties.paddleLayout().endpoint())
                    .contentType(MediaType.APPLICATION_JSON)
                    .body(requestPayload(request))
                    .retrieve()
                    .body(String.class);
            if (responseBody == null || responseBody.isBlank()) {
                throw new HttpOcrEngineException("PaddleOCR layout service returned an empty response body");
            }
            return parseResponse(responseBody);
        } catch (RestClientResponseException ex) {
            throw new HttpOcrEngineException("PaddleOCR layout service returned HTTP " + ex.getStatusCode().value() + ": " + ex.getResponseBodyAsString(), ex);
        } catch (RestClientException ex) {
            throw new HttpOcrEngineException("Failed to call PaddleOCR layout service", ex);
        } catch (IOException ex) {
            throw new HttpOcrEngineException("Failed to read PaddleOCR request file or parse response", ex);
        }
    }

    private Map<String, Object> requestPayload(OcrEngineRequest request) throws IOException {
        Map<String, Object> payload = new LinkedHashMap<>();
        payload.put("file", Base64.getEncoder().encodeToString(Files.readAllBytes(request.filePath())));
        payload.put("fileType", fileType(request));
        return payload;
    }

    private int fileType(OcrEngineRequest request) {
        String contentType = request.contentType();
        if (contentType != null && contentType.equalsIgnoreCase("application/pdf")) {
            return 0;
        }
        String filename = request.originalFilename();
        if (filename != null && filename.toLowerCase().endsWith(".pdf")) {
            return 0;
        }
        return 1;
    }

    private OcrEngineResult parseResponse(String responseBody) throws IOException {
        JsonNode root = objectMapper.readTree(responseBody);
        int errorCode = root.path("errorCode").asInt(0);
        if (errorCode != 0) {
            throw new HttpOcrEngineException("PaddleOCR layout service error " + errorCode + ": " + root.path("errorMsg").asText(""));
        }

        JsonNode result = root.path("result");
        return new OcrEngineResult(
                engineCode(),
                engineVersion(),
                responseBody,
                parseTextBlocks(result),
                List.of()
        );
    }

    private List<RecognizedTextBlock> parseTextBlocks(JsonNode result) {
        List<RecognizedTextBlock> blocks = new ArrayList<>();
        JsonNode layoutResults = result.path("layoutParsingResults");
        int order = 1;
        for (JsonNode layoutResult : layoutResults) {
            JsonNode pruned = layoutResult.path("prunedResult");
            JsonNode overallOcr = pruned.path("overall_ocr_res");
            JsonNode texts = overallOcr.path("rec_texts");
            JsonNode scores = overallOcr.path("rec_scores");
            JsonNode boxes = firstPresent(overallOcr, "rec_boxes", "dt_boxes", "rec_polys");
            for (int i = 0; i < texts.size(); i++) {
                blocks.add(new RecognizedTextBlock(
                        pageNo(layoutResult),
                        "TEXT",
                        texts.path(i).asText(""),
                        decimal(scores.path(i)),
                        jsonOrNull(boxes.path(i)),
                        order++
                ));
            }

            JsonNode tables = pruned.path("table_res_list");
            for (JsonNode table : tables) {
                String html = table.path("pred_html").asText("");
                if (!html.isBlank()) {
                    blocks.add(new RecognizedTextBlock(
                            pageNo(layoutResult),
                            "TABLE_HTML",
                            html,
                            null,
                            jsonOrNull(table.path("cell_box_list")),
                            order++
                    ));
                }
            }
        }
        return blocks;
    }

    private JsonNode firstPresent(JsonNode node, String... fieldNames) {
        for (String fieldName : fieldNames) {
            JsonNode value = node.path(fieldName);
            if (!value.isMissingNode() && !value.isNull()) {
                return value;
            }
        }
        return objectMapper.createArrayNode();
    }

    private int pageNo(JsonNode layoutResult) {
        int pageNo = layoutResult.path("pageNo").asInt(0);
        return pageNo <= 0 ? 1 : pageNo;
    }

    private BigDecimal decimal(JsonNode node) {
        if (node == null || node.isMissingNode() || node.isNull() || !node.isNumber()) {
            return null;
        }
        return node.decimalValue();
    }

    private String jsonOrNull(JsonNode node) {
        if (node == null || node.isMissingNode() || node.isNull()) {
            return null;
        }
        return node.toString();
    }
}