# Flyto2 Core - Source Available License
# Copyright (c) 2025 Flyto2. All Rights Reserved.
#
# This source code is licensed under the Flyto2 Source Available License v1.0.
# Commercial use requires a license. See LICENSE for details.

"""
Image OCR Module
Extract text from images using Tesseract OCR.
"""
import asyncio
import logging
import os
from typing import Any, Dict

from ...registry import register_module
from ...schema import compose
from ...schema.builders import field
from ...schema.constants import FieldGroup
from ...errors import ValidationError, ModuleError


logger = logging.getLogger(__name__)


@register_module(
    module_id='image.ocr',
    version='1.0.0',
    category='image',
    subcategory='analysis',
    tags=['image', 'ocr', 'text', 'extract', 'tesseract'],
    label='OCR - Extract Text',
    label_key='modules.image.ocr.label',
    description='Extract text from images using Tesseract OCR',
    description_key='modules.image.ocr.description',
    icon='FileText',
    color='#3B82F6',
    input_types=['file'],
    output_types=['string'],

    can_receive_from=['file.*', 'image.*', 'browser.*', 'http.*', 'flow.*', 'start'],
    can_connect_to=['text.*', 'string.*', 'data.*', 'flow.*'],

    retryable=True,
    max_retries=2,
    concurrent_safe=True,
    timeout_ms=120000,

    requires_credentials=False,
    handles_sensitive_data=False,
    required_permissions=[],

    params_schema=compose(
        field(
            'image_path',
            type='string',
            label='Image Path',
            label_key='modules.image.ocr.params.image_path.label',
            description='Path to the image file',
            description_key='modules.image.ocr.params.image_path.description',
            required=True,
            placeholder='/path/to/image.png',
            group=FieldGroup.BASIC,
        ),
        field(
            'language',
            type='string',
            label='Language',
            label_key='modules.image.ocr.params.language.label',
            description='OCR language code (e.g. eng, chi_tra, jpn)',
            description_key='modules.image.ocr.params.language.description',
            default='eng',
            placeholder='eng',
            group=FieldGroup.OPTIONS,
        ),
        field(
            'psm',
            type='number',
            label='Page Segmentation Mode',
            label_key='modules.image.ocr.params.psm.label',
            description='Tesseract page segmentation mode (0-13)',
            description_key='modules.image.ocr.params.psm.description',
            default=3,
            group=FieldGroup.ADVANCED,
        ),
        field(
            'output_type',
            type='string',
            label='Output Type',
            label_key='modules.image.ocr.params.output_type.label',
            description='Type of OCR output',
            description_key='modules.image.ocr.params.output_type.description',
            default='text',
            options=[
                {'value': 'text', 'label': 'Plain Text'},
                {'value': 'data', 'label': 'Structured Data'},
                {'value': 'boxes', 'label': 'Bounding Boxes'},
            ],
            group=FieldGroup.OPTIONS,
        ),
    ),
    output_schema={
        'text': {
            'type': 'string',
            'description': 'Extracted text from the image',
            'description_key': 'modules.image.ocr.output.text.description',
        },
        'confidence': {
            'type': 'number',
            'description': 'Average OCR confidence score (0-100)',
            'description_key': 'modules.image.ocr.output.confidence.description',
        },
        'language': {
            'type': 'string',
            'description': 'Language used for OCR',
            'description_key': 'modules.image.ocr.output.language.description',
        },
    },
    examples=[
        {
            'title': 'Extract text from image',
            'title_key': 'modules.image.ocr.examples.basic.title',
            'params': {
                'image_path': '/path/to/document.png',
                'language': 'eng',
            },
        },
    ],
    author='Flyto Team',
    license='MIT',
)
async def image_ocr(context: Dict[str, Any]) -> Dict[str, Any]:
    """Extract text from images using OCR."""
    params = context['params']
    image_path = params.get('image_path')
    language = params.get('language', 'eng')
    psm = params.get('psm', 3)
    output_type = params.get('output_type', 'text')

    if not image_path:
        raise ValidationError("Missing required parameter: image_path", field="image_path")

    if not os.path.exists(image_path):
        raise ModuleError(f"Image file not found: {image_path}")

    def _ocr():
        try:
            import pytesseract
            from PIL import Image
        except ImportError:
            raise ModuleError(
                "pytesseract and Pillow are required for image.ocr. "
                "Install with: pip install pytesseract Pillow"
            )

        img = Image.open(image_path)
        custom_config = f'--psm {psm}'

        if output_type == 'text':
            text = pytesseract.image_to_string(img, lang=language, config=custom_config)
            # Get confidence from detailed data
            try:
                data = pytesseract.image_to_data(img, lang=language, config=custom_config, output_type=pytesseract.Output.DICT)
                confidences = [int(c) for c in data['conf'] if int(c) > 0]
                avg_confidence = sum(confidences) / len(confidences) if confidences else 0.0
            except Exception:
                avg_confidence = 0.0

            return {
                'text': text.strip(),
                'confidence': round(avg_confidence, 2),
                'language': language,
            }

        elif output_type == 'data':
            data = pytesseract.image_to_data(img, lang=language, config=custom_config, output_type=pytesseract.Output.DICT)
            confidences = [int(c) for c in data['conf'] if int(c) > 0]
            avg_confidence = sum(confidences) / len(confidences) if confidences else 0.0
            words = []
            for i, word in enumerate(data['text']):
                if word.strip():
                    words.append({
                        'text': word,
                        'confidence': int(data['conf'][i]),
                        'x': data['left'][i],
                        'y': data['top'][i],
                        'width': data['width'][i],
                        'height': data['height'][i],
                    })

            full_text = ' '.join(w['text'] for w in words)
            return {
                'text': full_text,
                'confidence': round(avg_confidence, 2),
                'language': language,
                'words': words,
            }

        elif output_type == 'boxes':
            boxes = pytesseract.image_to_boxes(img, lang=language, config=custom_config)
            text = pytesseract.image_to_string(img, lang=language, config=custom_config)
            return {
                'text': text.strip(),
                'confidence': 0.0,
                'language': language,
                'boxes': boxes,
            }

        else:
            raise ValidationError(f"Unsupported output_type: {output_type}", field="output_type")

    result = await asyncio.to_thread(_ocr)
    logger.info(f"OCR extracted {len(result['text'])} characters from {image_path}")

    return {
        'ok': True,
        'data': result,
    }
