import sys
import json
import io
import logging
import time
import os
from paddleocr import PaddleOCR

# Configure stdout and stderr
sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8')
sys.stderr = io.TextIOWrapper(sys.stderr.buffer, encoding='utf-8')

# Suppress PaddleOCR logs
logging.getLogger('ppocr').setLevel(logging.ERROR)
logging.getLogger('paddleocr').setLevel(logging.ERROR)

reader = None

def init_reader(languages):
    start_time = time.time()
    global reader
    try:
        lang = languages.split(',')[0].strip() if languages else 'en'
        reader = PaddleOCR(
            use_textline_orientation=True,
            lang=lang,
            use_doc_unwarping=False,
            text_det_thresh=0.3,
            text_det_box_thresh=0.5,
            text_det_limit_side_len=960,
            text_det_unclip_ratio=1.2,
            text_recognition_batch_size=1,
            enable_mkldnn=True,
            cpu_threads=4,
        )
        end_time = time.time()
        print(f"[TIME] Initialization: {end_time - start_time:.3f}s", file=sys.stderr)
        return json.dumps({"status": "success", "message": "Reader initialized"})
    except Exception as e:
        print(f"[TIME] Initialization failed: {time.time() - start_time:.3f}s", file=sys.stderr)
        return json.dumps({"status": "error", "message": f"Failed to initialize: {str(e)}"})


def read_text(image_path):
    if reader is None:
        return json.dumps({"status": "error", "message": "Reader not initialized"})
    try:
        start_time = time.time()
        cleaned_image_path = image_path.strip()
        
        if not os.path.exists(cleaned_image_path):
            return json.dumps({"status": "error", "message": f"Image not found: {cleaned_image_path}"})

        ocr_start = time.time()
        result = reader.predict(cleaned_image_path)
        ocr_end = time.time()
        print(f"[TIME] OCR Processing: {ocr_end - ocr_start:.3f}s", file=sys.stderr)
        
        if not result:
            end_time = time.time()
            print(f"[TIME] Total read_text: {end_time - start_time:.3f}s", file=sys.stderr)
            return json.dumps({"status": "success", "data": []})

        formatted_result = []
        
        # Handle the PaddleOCR result structure
        for page_data in result:
            # Check if this is the detailed result format (dict with keys)
            if isinstance(page_data, dict) and 'rec_polys' in page_data and 'rec_texts' in page_data and 'rec_scores' in page_data:
                # Extract data from the detailed format
                polys = page_data['rec_polys']
                texts = page_data['rec_texts']
                scores = page_data['rec_scores']
                
                # Process each detected text region
                for i in range(len(texts)):
                    try:
                        # Convert numpy array to list of coordinate pairs
                        if i < len(polys):
                            bbox = [[int(coord[0]), int(coord[1])] for coord in polys[i]]
                        else:
                            bbox = []
                        
                        text = texts[i] if i < len(texts) else ""
                        confidence = float(scores[i]) if i < len(scores) else 0.0
                        
                        formatted_result.append({
                            'bbox': bbox,
                            'text': text,
                            'confidence': confidence
                        })
                    except Exception as e:
                        print(f"[ERROR] Processing item {i}: {str(e)}", file=sys.stderr)
                        continue
                        
            # Handle the simple list format: [[[bbox], [text, confidence]], ...]
            elif isinstance(page_data, list):
                for line in page_data:
                    try:
                        if len(line) >= 2:  # Ensure we have both bbox and text/confidence
                            bbox = [[int(coord[0]), int(coord[1])] for coord in line[0]]
                            
                            # Handle text and confidence extraction
                            if isinstance(line[1], (list, tuple)) and len(line[1]) >= 2:
                                text = str(line[1][0])
                                confidence = float(line[1][1])
                            elif isinstance(line[1], str):
                                text = line[1]
                                confidence = 1.0  # Default confidence if not provided
                            else:
                                text = str(line[1])
                                confidence = 0.0
                            
                            formatted_result.append({
                                'bbox': bbox,
                                'text': text,
                                'confidence': confidence
                            })
                    except Exception as e:
                        print(f"[ERROR] Processing line: {str(e)}", file=sys.stderr)
                        continue

        end_time = time.time()
        print(f"[TIME] Total read_text: {end_time - start_time:.3f}s", file=sys.stderr)
        print(f"[INFO] Processed {len(formatted_result)} text regions", file=sys.stderr)
        return json.dumps({"status": "success", "data": formatted_result})
        
    except Exception as e:
        end_time = time.time()
        print(f"[TIME] read_text failed: {end_time - start_time:.3f}s", file=sys.stderr)
        print(f"[ERROR] Exception: {str(e)}", file=sys.stderr)
        return json.dumps({"status": "error", "message": f"Error reading text: {str(e)}"})


def process_command(command, args):
    start_time = time.time()
    if command == 'init':
        result = init_reader(args)
    elif command == 'read_text':
        result = read_text(args)
    elif command == 'close':
        result = json.dumps({"status": "success", "message": "Closing process"})
    else:
        result = json.dumps({"status": "error", "message": "Invalid command"})
    end_time = time.time()
    print(f"[TIME] Command '{command}': {end_time - start_time:.3f}s", file=sys.stderr)
    return result


if __name__ == '__main__':
    while True:
        try:
            line = sys.stdin.readline().strip()
            if not line:
                break

            parts = line.split(maxsplit=1)
            command = parts[0]
            args = parts[1] if len(parts) > 1 else ""

            result = process_command(command, args)
            sys.stdout.write(result + '\n')
            sys.stdout.flush()

            if command == 'close':
                break
        except Exception as e:
            sys.stdout.write(json.dumps({"status": "error", "message": str(e)}) + '\n')
            sys.stdout.flush