
from ..base import *
from .schema.schema import *

import cv2
from tempfile import NamedTemporaryFile
import numpy as np
from io import BytesIO
from zipfile import ZipFile
from PIL import Image as PILImage
from deep_translator import GoogleTranslator
from pytesseract import Output, image_to_string
from fastapi import (
    UploadFile, 
    File, 
    HTTPException, 
    status,
    Depends
)

async def stream_image_bytes(image):
    try:
        img = PILImage.open(BytesIO(image))
        return img
    except Exception as err:
        return err


async def upload_file(
    file: UploadFile = File(..., description="Upload Single Image File")
    ) -> UploadFile:
    return file

async def read_upload_file(
    file = Depends(upload_file)
    ) -> bytes:
    im = await stream_image_bytes(await file.read())
    return im

async def bytes_to_tesseract(
    translate_from: TranslatorLangCodes, 
    translate_to: TranslatorLangCodes,
    image_bytes: bytes = Depends(read_upload_file)
    ):
    config = r'-l {0} --psm 4'.format(TessLangCodes[translate_from])
    text = image_to_string(
        image_bytes, 
        output_type=Output.STRING, 
        config=config
    )
    return (translate_from, translate_to, text)


async def translate_text_string( 
    data: Tuple[str, str, str] = Depends(bytes_to_tesseract)
    ) -> str:
    '''https://medium.com/analytics-vidhya/how-to-translate-text-with-python-9d203139dcf5'''

    translate_from, translate_to, text = data
    try:
        text = GoogleTranslator(
            source=translate_from, 
            target=translate_to
            ).translate(text)
        return text
    except Exception as err:
        logger.warn(f'Error Translate Text: {err}')
        return err

async def ocr_from_file(image, lang='deu'):
    im = await stream_image_bytes(image)
    config = r'-l {0} --psm 4'.format(lang)
    text = image_to_string(
        im, 
        output_type=Output.STRING, 
        config=config
    )
    return text


class BaseOCRTranslate(GoogleTranslator):

    async def available_languages(self) -> List[str]:
        try:
            return [code.value for code in TranslatorLangCodes]
        except Exception as err:
            logger.warn(f'Error Available Langs: {err}')
            return err

    async def goog_langs(self, **kwargs):
        """output: [arabic, french, english etc...]
        Returns:
            (List | Dict): list or dict or supported translation laguages 
        """
        langs = GoogleTranslator().get_supported_languages(kwargs)
        return langs

    @property
    def get_langs(self):
        return self.get_supported_languages()

class ExtractText:

    async def extract_image_text(
        self,
        image: UploadFile = File(...),
        lang: TranslatorLangCodes = TranslatorLangCodes.english
        ) -> str:
        '''Returns text from image for a given language
        - if no language is specified textract will auto detect.
            - assuming tessract language model exists.
        '''
        try:
            data = np.fromfile(image.file, dtype=np.uint8)
            image = cv2.imdecode(data, cv2.IMREAD_UNCHANGED)
            f_path = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

            lang = TessLangCodes[lang.value].value
            logger.info(f'LANG_CODE:   {lang}')

            text = image_to_string(
                f_path, 
                lang=lang,
                output_type=Output.STRING
            )
            logger.info(text)
            return text
        except Exception as e:
            err_msg = str(e)
            logger.error(err_msg)
            raise HTTPException(
                status_code=status.HTTP_404_NOT_FOUND, 
                detail=err_msg
            )


class TranslateFiles(BaseOCRTranslate):
    '''
    Tessdata lang codes: https://tesseract-ocr.github.io/tessdoc/Data-Files
    google_trans_new: https://github.com/lushan88a/google_trans_new
    https://towardsdatascience.com/language-translation-using-python-bd8020772ccc
    '''
    async def translate_image_file(
        self, 
        text = Depends(translate_text_string),
        ) -> str:
        try:
            return text
        except Exception as err:
            logger.warn(f'Error Translate Image File: {err}')
            return err


    async def translate_compressed_file(
        self,
        from_lang: TranslatorLangCodes, 
        to_lang: TranslatorLangCodes,
        file = Depends(upload_file)
        ) -> List[str]:
        lst = []
        with NamedTemporaryFile(suffix='.zip') as temp:
            temp.write(await file.read())
            with ZipFile(temp.name) as zp:
                file_list = zp.namelist()
                for file_name in file_list:
                    with zp.open(file_name) as z_file:
                        text = await ocr_from_file(
                            z_file.read(), 
                            lang=TessLangCodes[from_lang]
                        )
                        lst.append(text)
                        z_file.close()
                zp.close()
            temp.close()

        tran_lst = GoogleTranslator(from_lang, to_lang).translate_batch(lst)
        return tran_lst

    async def translate_image_file_list(
        self,
        files: List[UploadFile] = File(..., description="Upload Multiple Files")
        ):
        return {"filenames": [file.filename for file in files]}

