package com.zhengmeng.ocrplatform.engine;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.zhengmeng.ocrplatform.llm.LlmChatBotConfig;
import com.zhengmeng.ocrplatform.llm.LlmModelService;
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.pp-chat-ocr", name = "enabled", havingValue = "true")
public class PpChatOcrHttpEngineAdapter implements OcrEngineAdapter {
    private final EngineProperties properties;
    private final LlmModelService llmModelService;
    private final ObjectMapper objectMapper = new ObjectMapper();
    private final RestClient restClient;

    public PpChatOcrHttpEngineAdapter(EngineProperties properties, LlmModelService llmModelService) {
        this.properties = properties;
        this.llmModelService = llmModelService;
        SimpleClientHttpRequestFactory requestFactory = new SimpleClientHttpRequestFactory();
        requestFactory.setConnectTimeout(properties.ppChatOcr().timeout());
        requestFactory.setReadTimeout(properties.ppChatOcr().timeout());
        this.restClient = RestClient.builder()
                .requestFactory(requestFactory)
                .build();
    }

    @Override
    public String engineCode() {
        return "pp-chatocr";
    }

    @Override
    public String engineVersion() {
        return "PP-ChatOCR visual-chat";
    }

    @Override
    public OcrEngineResult recognize(OcrEngineRequest request) {
        try {
            String visualBody = post(properties.ppChatOcr().visualEndpoint(), visualPayload(request));
            JsonNode visualRoot = readSuccessfulResponse(visualBody, "PP-ChatOCR visual");
            JsonNode visualResult = visualRoot.path("result");

            String chatBody = null;
            JsonNode chatResult = objectMapper.createObjectNode();
            if (!properties.ppChatOcr().extractionKeys().isEmpty()) {
                chatBody = post(properties.ppChatOcr().chatEndpoint(), chatPayload(visualResult));
                JsonNode chatRoot = readSuccessfulResponse(chatBody, "PP-ChatOCR chat");
                chatResult = chatRoot.path("result").path("chatResult");
            }

            Map<String, Object> raw = new LinkedHashMap<>();
            raw.put("visual", objectMapper.readTree(visualBody));
            if (chatBody == null) {
                raw.put("chatSkipped", "No extraction keys configured");
            } else {
                raw.put("chat", objectMapper.readTree(chatBody));
            }
            String rawJson = objectMapper.writeValueAsString(raw);
            return new OcrEngineResult(
                    engineCode(),
                    engineVersion(),
                    rawJson,
                    parseTextBlocks(visualResult),
                    parseKeyValues(chatResult)
            );
        } catch (RestClientResponseException ex) {
            throw new HttpOcrEngineException("PP-ChatOCR service returned HTTP " + ex.getStatusCode().value() + ": " + ex.getResponseBodyAsString(), ex);
        } catch (RestClientException ex) {
            throw new HttpOcrEngineException("Failed to call PP-ChatOCR service", ex);
        } catch (IOException ex) {
            throw new HttpOcrEngineException("Failed to read PP-ChatOCR request file or parse response", ex);
        }
    }

    private String post(java.net.URI endpoint, Map<String, Object> payload) {
        String responseBody = restClient.post()
                .uri(endpoint)
                .contentType(MediaType.APPLICATION_JSON)
                .body(payload)
                .retrieve()
                .body(String.class);
        if (responseBody == null || responseBody.isBlank()) {
            throw new HttpOcrEngineException("PP-ChatOCR service returned an empty response body");
        }
        return responseBody;
    }

    private Map<String, Object> visualPayload(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));
        payload.put("useDocOrientationClassify", true);
        payload.put("useDocUnwarping", true);
        payload.put("useTextlineOrientation", true);
        payload.put("useSealRecognition", true);
        payload.put("useTableRecognition", true);
        return payload;
    }

    private Map<String, Object> chatPayload(JsonNode visualResult) {
        Map<String, Object> payload = new LinkedHashMap<>();
        payload.put("keyList", properties.ppChatOcr().extractionKeys());
        payload.put("visualInfo", objectMapper.convertValue(visualResult.path("visualInfo"), List.class));
        payload.put("useVectorRetrieval", false);
        payload.put("minCharacters", 500);
        llmModelService.currentChatBotConfig().ifPresent(config -> {
            payload.put("chatBotConfig", config.toChatBotConfig());
            payload.put("chatBotModelId", config.modelId());
        });
        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 JsonNode readSuccessfulResponse(String body, String label) throws IOException {
        JsonNode root = objectMapper.readTree(body);
        int errorCode = root.path("errorCode").asInt(0);
        if (errorCode != 0) {
            throw new HttpOcrEngineException(label + " service error " + errorCode + ": " + root.path("errorMsg").asText(""));
        }
        return root;
    }

    private List<RecognizedTextBlock> parseTextBlocks(JsonNode visualResult) {
        List<RecognizedTextBlock> blocks = new ArrayList<>();
        JsonNode layoutResults = visualResult.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 List<RecognizedKeyValue> parseKeyValues(JsonNode chatResult) {
        if (chatResult == null || chatResult.isMissingNode() || chatResult.isNull() || chatResult.isEmpty()) {
            return List.of();
        }

        JsonNode candidates = firstPresent(chatResult, "keyValues", "kvPairs", "fields");
        if (candidates.isArray()) {
            List<RecognizedKeyValue> values = new ArrayList<>();
            for (JsonNode item : candidates) {
                String key = text(item, "fieldKey", text(item, "key", ""));
                String value = text(item, "valueText", text(item, "value", ""));
                values.add(new RecognizedKeyValue(
                        key,
                        text(item, "fieldName", text(item, "name", key)),
                        value,
                        text(item, "normalizedValue", value),
                        decimal(item.path("confidence")),
                        intOrDefault(item.path("pageNo"), 1),
                        jsonOrNull(item.path("evidence"))
                ));
            }
            return values;
        }

        if (chatResult.isObject()) {
            List<RecognizedKeyValue> values = new ArrayList<>();
            chatResult.fields().forEachRemaining(entry -> {
                JsonNode node = entry.getValue();
                String value = node.isValueNode() ? node.asText() : node.toString();
                values.add(new RecognizedKeyValue(
                        entry.getKey(),
                        entry.getKey(),
                        value,
                        value,
                        null,
                        1,
                        "{\"source\":\"pp-chatocr-chat\"}"
                ));
            });
            return values;
        }

        return List.of();
    }

    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 Integer intOrDefault(JsonNode node, int defaultValue) {
        return node == null || !node.canConvertToInt() ? defaultValue : node.asInt();
    }

    private String text(JsonNode node, String name, String defaultValue) {
        JsonNode value = node == null ? null : node.get(name);
        if (value == null || value.isNull()) {
            return defaultValue;
        }
        return value.isValueNode() ? value.asText() : value.toString();
    }

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