package org.raddatz.familienarchiv.ocr;

import com.fasterxml.jackson.annotation.JsonProperty;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.DeserializationFeature;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import lombok.extern.slf4j.Slf4j;
import org.raddatz.familienarchiv.ocr.ScriptType;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.core.ParameterizedTypeReference;
import org.springframework.core.io.ByteArrayResource;
import org.springframework.http.HttpEntity;
import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;
import org.springframework.http.client.JdkClientHttpRequestFactory;
import org.springframework.lang.Nullable;
import org.springframework.stereotype.Component;
import org.springframework.util.LinkedMultiValueMap;
import org.springframework.util.MultiValueMap;
import org.springframework.web.client.RestClient;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.net.URI;
import java.net.http.HttpClient;
import java.net.http.HttpRequest;
import java.net.http.HttpResponse;
import java.nio.charset.StandardCharsets;
import java.time.Duration;
import java.util.List;
import java.util.Map;
import java.util.function.Consumer;

@Component
@Slf4j
public class RestClientOcrClient implements OcrClient, OcrHealthClient {

    private static final ObjectMapper NDJSON_MAPPER = new ObjectMapper()
            .configure(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES, true);

    private final RestClient restClient;
    private final RestClient trainingRestClient;
    private final HttpClient streamingHttpClient;
    private final String baseUrl;
    private final String trainingToken;

    public RestClientOcrClient(
            @Value("${app.ocr.base-url:http://ocr-service:8000}") String baseUrl,
            @Value("${app.ocr.training-token:}") String trainingToken) {
        this.baseUrl = baseUrl;
        this.trainingToken = trainingToken;

        HttpClient httpClient = HttpClient.newBuilder()
                .version(HttpClient.Version.HTTP_1_1)
                .connectTimeout(Duration.ofSeconds(10))
                .build();
        JdkClientHttpRequestFactory requestFactory = new JdkClientHttpRequestFactory(httpClient);
        requestFactory.setReadTimeout(Duration.ofMinutes(10));

        this.restClient = RestClient.builder()
                .baseUrl(baseUrl)
                .requestFactory(requestFactory)
                .build();

        HttpClient trainingHttpClient = HttpClient.newBuilder()
                .version(HttpClient.Version.HTTP_1_1)
                .connectTimeout(Duration.ofSeconds(10))
                .build();
        JdkClientHttpRequestFactory trainingRequestFactory = new JdkClientHttpRequestFactory(trainingHttpClient);
        trainingRequestFactory.setReadTimeout(Duration.ofMinutes(10));
        this.trainingRestClient = RestClient.builder()
                .baseUrl(baseUrl)
                .requestFactory(trainingRequestFactory)
                .build();

        this.streamingHttpClient = HttpClient.newBuilder()
                .version(HttpClient.Version.HTTP_1_1)
                .connectTimeout(Duration.ofSeconds(10))
                .build();
    }

    @Override
    public List<OcrBlockResult> extractBlocks(String pdfUrl, ScriptType scriptType) {
        Map<String, String> body = Map.of(
                "pdfUrl", pdfUrl,
                "scriptType", scriptType.name(),
                "language", "de");

        List<OcrBlockJson> response = restClient.post()
                .uri("/ocr")
                .contentType(MediaType.APPLICATION_JSON)
                .body(body)
                .retrieve()
                .body(new ParameterizedTypeReference<>() {});

        if (response == null) return List.of();

        return response.stream()
                .map(OcrBlockJson::toResult)
                .toList();
    }

    private RestClient.RequestBodySpec addTrainingAuth(RestClient.RequestBodySpec spec) {
        if (trainingToken != null && !trainingToken.isBlank()) {
            return spec.header("X-Training-Token", trainingToken);
        }
        return spec;
    }

    @Override
    public OcrClient.TrainingResult trainModel(byte[] trainingDataZip) {
        ByteArrayResource zipResource = new ByteArrayResource(trainingDataZip) {
            @Override
            public String getFilename() { return "training-data.zip"; }
        };

        MultiValueMap<String, Object> body = new LinkedMultiValueMap<>();
        HttpHeaders partHeaders = new HttpHeaders();
        partHeaders.setContentType(MediaType.parseMediaType("application/zip"));
        body.add("file", new HttpEntity<>(zipResource, partHeaders));

        TrainingResultJson result = addTrainingAuth(
                trainingRestClient.post()
                        .uri("/train")
                        .contentType(MediaType.MULTIPART_FORM_DATA))
                .body(body)
                .retrieve()
                .body(TrainingResultJson.class);

        if (result == null) return new OcrClient.TrainingResult(null, null, null, null);
        return new OcrClient.TrainingResult(result.loss(), result.accuracy(), result.cer(), result.epochs());
    }

    @Override
    public OcrClient.TrainingResult segtrainModel(byte[] trainingDataZip) {
        ByteArrayResource zipResource = new ByteArrayResource(trainingDataZip) {
            @Override
            public String getFilename() { return "segmentation-data.zip"; }
        };

        MultiValueMap<String, Object> body = new LinkedMultiValueMap<>();
        HttpHeaders partHeaders = new HttpHeaders();
        partHeaders.setContentType(MediaType.parseMediaType("application/zip"));
        body.add("file", new HttpEntity<>(zipResource, partHeaders));

        TrainingResultJson result = addTrainingAuth(
                trainingRestClient.post()
                        .uri("/segtrain")
                        .contentType(MediaType.MULTIPART_FORM_DATA))
                .body(body)
                .retrieve()
                .body(TrainingResultJson.class);

        if (result == null) return new OcrClient.TrainingResult(null, null, null, null);
        return new OcrClient.TrainingResult(result.loss(), result.accuracy(), result.cer(), result.epochs());
    }

    @Override
    public OcrClient.TrainingResult trainSenderModel(byte[] trainingDataZip, String outputModelPath) {
        ByteArrayResource zipResource = new ByteArrayResource(trainingDataZip) {
            @Override
            public String getFilename() { return "sender-training-data.zip"; }
        };

        MultiValueMap<String, Object> body = new LinkedMultiValueMap<>();
        HttpHeaders partHeaders = new HttpHeaders();
        partHeaders.setContentType(MediaType.parseMediaType("application/zip"));
        body.add("file", new HttpEntity<>(zipResource, partHeaders));
        body.add("output_model_path", outputModelPath);

        TrainingResultJson result = addTrainingAuth(
                trainingRestClient.post()
                        .uri("/train-sender")
                        .contentType(MediaType.MULTIPART_FORM_DATA))
                .body(body)
                .retrieve()
                .body(TrainingResultJson.class);

        if (result == null) return new OcrClient.TrainingResult(null, null, null, null);
        return new OcrClient.TrainingResult(result.loss(), result.accuracy(), result.cer(), result.epochs());
    }

    @Override
    public boolean isHealthy() {
        try {
            restClient.get()
                    .uri("/health")
                    .retrieve()
                    .toBodilessEntity();
            return true;
        } catch (Exception e) {
            log.warn("OCR service health check failed: {}", e.getMessage());
            return false;
        }
    }

    @Override
    public void streamBlocks(String pdfUrl, ScriptType scriptType,
                              List<OcrRegion> regions, @Nullable String senderModelPath,
                              Consumer<OcrStreamEvent> handler) {
        String body;
        try {
            var requestMap = new java.util.LinkedHashMap<String, Object>();
            requestMap.put("pdfUrl", pdfUrl);
            requestMap.put("scriptType", scriptType.name());
            requestMap.put("language", "de");
            if (regions != null && !regions.isEmpty()) {
                requestMap.put("regions", regions);
            }
            if (senderModelPath != null) {
                requestMap.put("senderModelPath", senderModelPath);
            }
            body = NDJSON_MAPPER.writeValueAsString(requestMap);
        } catch (IOException e) {
            throw new RuntimeException("Failed to serialize OCR request", e);
        }

        HttpRequest request = HttpRequest.newBuilder()
                .uri(URI.create(baseUrl + "/ocr/stream"))
                .header("Content-Type", "application/json")
                .POST(HttpRequest.BodyPublishers.ofString(body))
                .timeout(Duration.ofMinutes(5))
                .build();

        try {
            HttpResponse<InputStream> response = streamingHttpClient.send(
                    request, HttpResponse.BodyHandlers.ofInputStream());

            if (response.statusCode() == 404) {
                log.info("OCR service does not support /ocr/stream (404), falling back to /ocr");
                List<OcrBlockResult> allBlocks = extractBlocks(pdfUrl, scriptType);
                handler.accept(new OcrStreamEvent.Start(0));
                for (OcrBlockResult block : allBlocks) {
                    handler.accept(new OcrStreamEvent.Page(block.pageNumber(), List.of(block)));
                }
                handler.accept(new OcrStreamEvent.Done(allBlocks.size(), 0));
                return;
            }

            try (InputStream inputStream = response.body()) {
                parseNdjsonStream(inputStream, handler);
            }
        } catch (IOException | InterruptedException e) {
            if (e instanceof InterruptedException) {
                Thread.currentThread().interrupt();
            }
            throw new RuntimeException("NDJSON stream failed: " + e.getMessage(), e);
        }
    }

    static void parseNdjsonStream(InputStream inputStream, Consumer<OcrStreamEvent> handler) {
        try (BufferedReader reader = new BufferedReader(
                new InputStreamReader(inputStream, StandardCharsets.UTF_8))) {
            String line;
            while ((line = reader.readLine()) != null) {
                if (line.isBlank()) continue;

                JsonNode node = NDJSON_MAPPER.readTree(line);
                String type = node.path("type").asText();

                switch (type) {
                    case "start" -> handler.accept(
                            new OcrStreamEvent.Start(node.path("totalPages").asInt()));
                    case "preprocessing" -> handler.accept(
                            new OcrStreamEvent.Preprocessing(node.path("pageNumber").asInt()));
                    case "page" -> {
                        int pageNumber = node.path("pageNumber").asInt();
                        List<OcrBlockResult> blocks = NDJSON_MAPPER.convertValue(
                                node.path("blocks"),
                                new TypeReference<>() {});
                        handler.accept(new OcrStreamEvent.Page(pageNumber, blocks));
                    }
                    case "error" -> handler.accept(
                            new OcrStreamEvent.Error(
                                    node.path("pageNumber").asInt(),
                                    node.path("message").asText()));
                    case "done" -> handler.accept(
                            new OcrStreamEvent.Done(
                                    node.path("totalBlocks").asInt(),
                                    node.path("skippedPages").asInt()));
                    default -> log.debug("Ignoring unknown NDJSON event type: {}", type);
                }
            }
        } catch (IOException e) {
            throw new RuntimeException("Failed to parse NDJSON stream: " + e.getMessage(), e);
        }
    }

    record TrainingResultJson(Double loss, Double accuracy, Double cer, Integer epochs) {}

    record OcrBlockJson(
            @JsonProperty("pageNumber") int pageNumber,
            double x,
            double y,
            double width,
            double height,
            List<List<Double>> polygon,
            String text,
            String annotationId
    ) {
        OcrBlockResult toResult() {
            return new OcrBlockResult(pageNumber, x, y, width, height, polygon, text, annotationId);
        }
    }
}
