package info.selflearner.ocr.util;

import cn.easyproject.easyocr.EasyOCR;
import info.selflearner.ocr.model.Passport;
import net.sourceforge.tess4j.*;
import net.sourceforge.tess4j.util.LoadLibs;
import org.json.JSONArray;
import org.json.JSONObject;

import javax.imageio.ImageIO;
import java.awt.*;
import java.awt.geom.AffineTransform;
import java.awt.image.AffineTransformOp;
import java.awt.image.BufferedImage;
import java.io.File;
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;

public class OCR {
    private static final ITesseract TESSERACT = new Tesseract();
    private static final EasyOCR EASY_OCR = new EasyOCR();

    static {
        TESSERACT.setDatapath("src/main/resources/tessdata/");
        TESSERACT.setVariable("user_defined_dpi", "300");
        System.setProperty("java.library.path", LoadLibs.extractTessResources("win32-x86-64").getPath());
    }

    public static String perform(File file, String languages) throws TesseractException {
        TESSERACT.setLanguage(languages);
        String ocrResult = TESSERACT.doOCR(file);
        file.delete();
        return ocrResult;
    }

    public static Passport performPassport(File file) throws IOException, InterruptedException {
        for (int imageRotateAngle = 0; imageRotateAngle <= 270; imageRotateAngle += 90) {
            if (imageRotateAngle == 90) {
                System.out.println("Rotating image to 90 degrees...");
                rotateImageVertically(file, true);
            } else if (imageRotateAngle == 180) {
                System.out.println("Rotating image to 180 degrees...");
                rotateImageHorizontally(file);
            } else if (imageRotateAngle == 270) {
                System.out.println("Rotating image to 270 degrees...");
                rotateImageVertically(file, false);
            }

            Runtime.getRuntime().exec(new String[]{"surya_ocr", file.getPath(), "--results_dir", System.getProperty("java.io.tmpdir"), "--langs=en"}).waitFor();
            JSONArray jsonTextLines = new JSONObject(Files.readAllLines(Path.of(System.getProperty("java.io.tmpdir"), "/", file.getName().split("\\.")[0], "/results.json")).getFirst())
                    .getJSONArray(file.getName().split("\\.")[0]).getJSONObject(0).getJSONArray("text_lines");
            int mrzStartingIndex = -1;

            for (int arrayIndex = 0; arrayIndex < jsonTextLines.length() && mrzStartingIndex == -1; arrayIndex++) {
                if (jsonTextLines.getJSONObject(arrayIndex).getString("text").indexOf("P<") == 0) {
                    mrzStartingIndex = arrayIndex;
                }
            }

            if (mrzStartingIndex != -1) {
                System.out.println("MRZ1: " + jsonTextLines.getJSONObject(mrzStartingIndex).getString("text"));
                System.out.println("MRZ2: " + jsonTextLines.getJSONObject(mrzStartingIndex + 1).getString("text"));
                try {
                    return new Passport(new String[]{jsonTextLines.getJSONObject(mrzStartingIndex).getString("text"), jsonTextLines.getJSONObject(mrzStartingIndex + 1).getString("text")},
                            new int[]{(int) (jsonTextLines.getJSONObject(mrzStartingIndex).getDouble("confidence") * 100), (int) (jsonTextLines.getJSONObject(mrzStartingIndex + 1).getDouble("confidence") * 100)});
                } catch (Exception _) {

                }
            }
        }
        throw new IOException("Passport not found");
    }

    private static void rotateImageVertically(File image, boolean counterClockwise) throws IOException {
        BufferedImage bufferedImage = ImageIO.read(image);
        BufferedImage output = new BufferedImage(bufferedImage.getHeight(), bufferedImage.getWidth(), bufferedImage.getType());

        AffineTransform affineTransform = getAffineTransform(counterClockwise, bufferedImage);

        AffineTransformOp affineTransformOp = new AffineTransformOp(affineTransform, AffineTransformOp.TYPE_BILINEAR);
        affineTransformOp.filter(bufferedImage, output);

        ImageIO.write(output, "jpg", image);
    }

    public static void rotateImageHorizontally(File image) throws IOException {
        BufferedImage imageToRotate = ImageIO.read(image);

        int widthOfImage = imageToRotate.getWidth();
        int heightOfImage = imageToRotate.getHeight();
        int typeOfImage = imageToRotate.getType();

        BufferedImage output = new BufferedImage(widthOfImage, heightOfImage, typeOfImage);

        Graphics2D graphics2D = output.createGraphics();

        graphics2D.rotate(Math.toRadians(180), widthOfImage / 2.0, heightOfImage / 2.0);
        graphics2D.drawImage(imageToRotate, null, 0, 0);

        ImageIO.write(output, "jpg", image);
    }

    private static AffineTransform getAffineTransform(boolean counterClockwise, BufferedImage bufferedImage) {
        int imageWidth = bufferedImage.getWidth();
        int imageHeight = bufferedImage.getHeight();

        AffineTransform affineTransform = new AffineTransform();
        affineTransform.rotate((!counterClockwise ? Math.PI : -Math.PI) / 2, imageWidth / 2.0, imageHeight / 2.0);

        double offset = (imageWidth - imageHeight) / 2.0;
        if (!counterClockwise) {
            affineTransform.translate(offset, offset);
        } else {
            affineTransform.translate(-offset, -offset);
        }
        return affineTransform;
    }

//    public static Passport performPassport(File file) throws TesseractException, IOException {
//        TESSERACT.setLanguage("eng");
//
//        for (int i = 0; i < 4; i++) {
//            String ocrResult = null;
//
//            switch (i) {
//                case 0 -> ocrResult = TESSERACT.doOCR(file);
//                case 1 -> ocrResult = EASY_OCR.discern(file);
//                case 2 -> ocrResult = EASY_OCR.discernAndAutoCleanImage(file, ImageType.CAPTCHA_INTERFERENCE_LINE);
//            }
//
//            for (double imageWidthRatio = 0.8; imageWidthRatio <= 1.6; imageWidthRatio += 0.2) {
//                for (double imageHeightRatio = 0.8; imageHeightRatio <= 1.6; imageHeightRatio += 0.2) {
//                    if (i == 3) {
//                        ocrResult = EASY_OCR.discernAndAutoCleanImage(file, ImageType.CAPTCHA_NORMAL, imageWidthRatio, imageHeightRatio);
//                        System.out.println("imageWidthRatio: " + imageWidthRatio + ", imageHeightRatio: " + imageHeightRatio);
//                    }
//
//                    int mrzStartingIndex = ocrResult.indexOf("P<");
//
//                    if (mrzStartingIndex != -1 && mrzStartingIndex + 89 <= ocrResult.length()) {
//                        String[] mrz = ocrResult.substring(mrzStartingIndex, mrzStartingIndex + 89).split("\n");
//
//                        if (mrz[0].length() >= 44 && mrz[1].length() >= 44) {
//                            System.out.println("mrz: " + Arrays.toString(mrz));
//                            Passport passport;
//
//                            try {
//                                passport = new Passport(mrz);
//                            } catch (StringIndexOutOfBoundsException | NumberFormatException exception) {
//                                continue;
//                            }
//
//                            return passport;
//                        }
//                    }
//
//                    if (i != 3) break;
//                }
//                if (i != 3) break;
//            }
//        }
//
//        throw new IOException("Passport not found");
//    }

}