package net.micode.myapplication.service;

import android.content.Context;
import android.graphics.Bitmap;
import android.graphics.BitmapFactory;
import android.graphics.Matrix;
import android.media.ExifInterface;
import android.net.Uri;
import android.util.Log;

import com.google.mlkit.vision.common.InputImage;
import com.google.mlkit.vision.text.Text;
import com.google.mlkit.vision.text.TextRecognition;
import com.google.mlkit.vision.text.TextRecognizer;
import com.google.mlkit.vision.text.chinese.ChineseTextRecognizerOptions;

import com.tom_roush.pdfbox.pdmodel.PDDocument;
import com.tom_roush.pdfbox.rendering.PDFRenderer;

import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicReference;

public class OcrService {
    private static final String TAG = "OcrService";
    private static final int MAX_IMAGE_SIZE = 2048; // 最大图片尺寸
    private static final int PDF_DPI = 150; // PDF渲染DPI
    
    private Context context;
    private TextRecognizer textRecognizer;
    
    public OcrService(Context context) {
        this.context = context;
        // 初始化中文文本识别器
        this.textRecognizer = TextRecognition.getClient(new ChineseTextRecognizerOptions.Builder().build());
    }
    
    /**
     * 从文件提取文本（支持图片和PDF）
     */
    public String extractTextFromFile(String filePath) throws Exception {
        if (filePath == null || filePath.isEmpty()) {
            throw new IllegalArgumentException("文件路径不能为空");
        }
        
        File file = new File(filePath);
        if (!file.exists()) {
            throw new IOException("文件不存在: " + filePath);
        }
        
        String fileName = file.getName().toLowerCase();
        
        try {
            if (fileName.endsWith(".pdf")) {
                return extractTextFromPdf(filePath);
            } else if (isImageFile(fileName)) {
                return extractTextFromImage(filePath);
            } else {
                throw new UnsupportedOperationException("不支持的文件类型: " + fileName);
            }
        } catch (Exception e) {
            Log.e(TAG, "OCR提取文本失败: " + filePath, e);
            throw e;
        }
    }
    
    /**
     * 从图片文件提取文本
     */
    public String extractTextFromImage(String imagePath) throws Exception {
        Bitmap bitmap = loadAndProcessImage(imagePath);
        if (bitmap == null) {
            throw new IOException("无法加载图片: " + imagePath);
        }
        
        try {
            InputImage image = InputImage.fromBitmap(bitmap, 0);
            return performTextRecognition(image);
        } finally {
            if (bitmap != null && !bitmap.isRecycled()) {
                bitmap.recycle();
            }
        }
    }
    
    /**
     * 从Bitmap提取文本
     */
    public String extractTextFromBitmap(Bitmap bitmap) throws Exception {
        if (bitmap == null) {
            throw new IllegalArgumentException("Bitmap不能为空");
        }
        
        try {
            InputImage image = InputImage.fromBitmap(bitmap, 0);
            return performTextRecognition(image);
        } catch (Exception e) {
            Log.e(TAG, "从Bitmap提取文本失败", e);
            throw e;
        }
    }
    
    /**
     * 从PDF文件提取文本
     */
    public String extractTextFromPdf(String pdfPath) throws Exception {
        PDDocument document = null;
        StringBuilder result = new StringBuilder();
        
        try {
            document = PDDocument.load(new File(pdfPath));
            PDFRenderer renderer = new PDFRenderer(document);
            
            int pageCount = document.getNumberOfPages();
            Log.i(TAG, "PDF总页数: " + pageCount);
            
            for (int pageIndex = 0; pageIndex < pageCount; pageIndex++) {
                Log.i(TAG, "处理PDF第 " + (pageIndex + 1) + " 页");
                
                // 将PDF页面渲染为Bitmap
                Bitmap bitmap = renderer.renderImageWithDPI(pageIndex, PDF_DPI);
                if (bitmap != null) {
                    try {
                        // 从Bitmap提取文本
                        String pageText = extractTextFromBitmap(bitmap);
                        if (pageText != null && !pageText.trim().isEmpty()) {
                            result.append(pageText);
                            if (pageIndex < pageCount - 1) {
                                result.append("\n\n"); // 页间分隔
                            }
                        }
                    } finally {
                        if (!bitmap.isRecycled()) {
                            bitmap.recycle();
                        }
                    }
                }
            }
            
            return result.toString().trim();
            
        } catch (Exception e) {
            Log.e(TAG, "PDF处理失败: " + pdfPath, e);
            throw e;
        } finally {
            if (document != null) {
                try {
                    document.close();
                } catch (IOException e) {
                    Log.e(TAG, "关闭PDF文档失败", e);
                }
            }
        }
    }
    
    /**
     * 执行文本识别
     */
    private String performTextRecognition(InputImage image) throws Exception {
        final CountDownLatch latch = new CountDownLatch(1);
        final AtomicReference<String> result = new AtomicReference<>("");
        final AtomicReference<Exception> error = new AtomicReference<>(null);
        
        textRecognizer.process(image)
            .addOnSuccessListener(visionText -> {
                try {
                    String extractedText = processVisionText(visionText);
                    result.set(extractedText);
                } catch (Exception e) {
                    error.set(e);
                } finally {
                    latch.countDown();
                }
            })
            .addOnFailureListener(e -> {
                error.set(e);
                latch.countDown();
            });
        
        // 等待识别完成（最多30秒）
        if (!latch.await(30, TimeUnit.SECONDS)) {
            throw new Exception("文本识别超时");
        }
        
        if (error.get() != null) {
            throw error.get();
        }
        
        return result.get();
    }
    
    /**
     * 处理Vision Text结果
     */
    private String processVisionText(Text visionText) {
        StringBuilder result = new StringBuilder();
        
        for (Text.TextBlock block : visionText.getTextBlocks()) {
            for (Text.Line line : block.getLines()) {
                result.append(line.getText()).append("\n");
            }
            result.append("\n"); // 块间分隔
        }
        
        return result.toString().trim();
    }
    
    /**
     * 加载并处理图片
     */
    private Bitmap loadAndProcessImage(String imagePath) {
        BitmapFactory.Options options = new BitmapFactory.Options();
        options.inJustDecodeBounds = true;
        BitmapFactory.decodeFile(imagePath, options);
        
        // 计算缩放比例
        int scaleFactor = calculateInSampleSize(options, MAX_IMAGE_SIZE, MAX_IMAGE_SIZE);
        
        options.inJustDecodeBounds = false;
        options.inSampleSize = scaleFactor;
        
        Bitmap bitmap = BitmapFactory.decodeFile(imagePath, options);
        if (bitmap == null) {
            return null;
        }
        
        // 处理图片方向
        bitmap = handleImageOrientation(bitmap, imagePath);
        
        return bitmap;
    }
    
    /**
     * 计算图片缩放比例
     */
    private int calculateInSampleSize(BitmapFactory.Options options, int reqWidth, int reqHeight) {
        final int height = options.outHeight;
        final int width = options.outWidth;
        int inSampleSize = 1;
        
        if (height > reqHeight || width > reqWidth) {
            final int halfHeight = height / 2;
            final int halfWidth = width / 2;
            
            while ((halfHeight / inSampleSize) >= reqHeight && (halfWidth / inSampleSize) >= reqWidth) {
                inSampleSize *= 2;
            }
        }
        
        return inSampleSize;
    }
    
    /**
     * 处理图片方向
     */
    private Bitmap handleImageOrientation(Bitmap bitmap, String imagePath) {
        try {
            ExifInterface exif = new ExifInterface(imagePath);
            int orientation = exif.getAttributeInt(ExifInterface.TAG_ORIENTATION, 
                ExifInterface.ORIENTATION_NORMAL);
            
            Matrix matrix = new Matrix();
            switch (orientation) {
                case ExifInterface.ORIENTATION_ROTATE_90:
                    matrix.postRotate(90);
                    break;
                case ExifInterface.ORIENTATION_ROTATE_180:
                    matrix.postRotate(180);
                    break;
                case ExifInterface.ORIENTATION_ROTATE_270:
                    matrix.postRotate(270);
                    break;
                default:
                    return bitmap;
            }
            
            Bitmap rotatedBitmap = Bitmap.createBitmap(bitmap, 0, 0, bitmap.getWidth(), 
                bitmap.getHeight(), matrix, true);
            
            // 回收原始bitmap
            if (rotatedBitmap != bitmap && !bitmap.isRecycled()) {
                bitmap.recycle();
            }
            
            return rotatedBitmap;
            
        } catch (IOException e) {
            Log.e(TAG, "处理图片方向失败", e);
            return bitmap;
        }
    }
    
    /**
     * 检查是否为图片文件
     */
    private boolean isImageFile(String fileName) {
        String[] imageExtensions = {".jpg", ".jpeg", ".png", ".bmp", ".webp", ".gif"};
        String lowerFileName = fileName.toLowerCase();
        
        for (String ext : imageExtensions) {
            if (lowerFileName.endsWith(ext)) {
                return true;
            }
        }
        return false;
    }
    
    /**
     * 关闭OCR服务
     */
    public void close() {
        if (textRecognizer != null) {
            textRecognizer.close();
        }
    }
}