# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
""" Helper classes for multiprocess captcha image generation
This module also provides script for saving captcha images to file using CLI.
"""

from __future__ import print_function
import random

import numpy as np
from captcha.image import ImageCaptcha
import cv2
from multiproc_data import MPData


class CaptchaGen(object):
    """Generates a captcha image
    """
    def __init__(self, h, w, font_paths):
        """
        Parameters
        ----------
        h: int
            Height of the generated images
        w: int
            Width of the generated images
        font_paths: list of str
            List of all fonts in ttf format
        """
        self.captcha = ImageCaptcha(fonts=font_paths)
        self.h = h
        self.w = w

    def image(self, captcha_str):
        """Generate a greyscale captcha image representing number string

        Parameters
        ----------
        captcha_str: str
            string a characters for captcha image

        Returns
        -------
        numpy.ndarray
            Generated greyscale image in np.ndarray float type with values normalized to [0, 1]
        """
        img = self.captcha.generate(captcha_str)
        img = np.fromstring(img.getvalue(), dtype='uint8')
        img = cv2.imdecode(img, cv2.IMREAD_GRAYSCALE)
        img = cv2.resize(img, (self.h, self.w))
        img = img.transpose(1, 0)
        img = np.multiply(img, 1 / 255.0)
        return img


class DigitCaptcha(object):
    """Provides shape() and get() interface for digit-captcha image generation
    """
    def __init__(self, font_paths, h, w, num_digit_min, num_digit_max):
        """
        Parameters
        ----------
        font_paths: list of str
            List of path to ttf font files
        h: int
            height of the generated image
        w: int
            width of the generated image
        num_digit_min: int
            minimum number of digits in generated captcha image
        num_digit_max: int
            maximum number of digits in generated captcha image
        """
        self.num_digit_min = num_digit_min
        self.num_digit_max = num_digit_max
        self.captcha = CaptchaGen(h=h, w=w, font_paths=font_paths)

    @property
    def shape(self):
        """Returns shape of the image data generated

        Returns
        -------
        tuple(int, int)
        """
        return self.captcha.h, self.captcha.w

    def get(self):
        """Get an image from the queue

        Returns
        -------
        np.ndarray
            A captcha image, normalized to [0, 1]
        """
        return self._gen_sample()

    @staticmethod
    def get_rand(num_digit_min, num_digit_max):
        """Generates a character string of digits. Number of digits are
        between self.num_digit_min and self.num_digit_max
        Returns
        -------
        str
        """
        buf = ""
        max_len = random.randint(num_digit_min, num_digit_max)
        for i in range(max_len):
            buf += str(random.randint(0, 9))
        return buf

    def _gen_sample(self):
        """Generate a random captcha image sample
        Returns
        -------
        (numpy.ndarray, str)
            Tuple of image (numpy ndarray) and character string of digits used to generate the image
        """
        num_str = self.get_rand(self.num_digit_min, self.num_digit_max)
        return self.captcha.image(num_str), num_str


class MPDigitCaptcha(DigitCaptcha):
    """Handles multi-process captcha image generation
    """
    def __init__(self, font_paths, h, w, num_digit_min, num_digit_max, num_processes, max_queue_size):
        """Parameters
        ----------
        font_paths: list of str
            List of path to ttf font files
        h: int
            height of the generated image
        w: int
            width of the generated image
        num_digit_min: int
            minimum number of digits in generated captcha image
        num_digit_max: int
            maximum number of digits in generated captcha image
        num_processes: int
            Number of processes to spawn
        max_queue_size: int
            Maximum images in queue before processes wait
        """
        super(MPDigitCaptcha, self).__init__(font_paths, h, w, num_digit_min, num_digit_max)
        self.mp_data = MPData(num_processes, max_queue_size, self._gen_sample)

    def start(self):
        """Starts the processes"""
        self.mp_data.start()

    def get(self):
        """Get an image from the queue

        Returns
        -------
        np.ndarray
            A captcha image, normalized to [0, 1]
        """
        return self.mp_data.get()

    def reset(self):
        """Resets the generator by stopping all processes"""
        self.mp_data.reset()


if __name__ == '__main__':
    import argparse

    def main():
        """Program entry point"""
        parser = argparse.ArgumentParser()
        parser.add_argument("font_path", help="Path to ttf font file")
        parser.add_argument("output", help="Output filename including extension (e.g. 'sample.jpg')")
        parser.add_argument("--num", help="Up to 4 digit number [Default: random]")
        args = parser.parse_args()

        captcha = ImageCaptcha(fonts=[args.font_path])
        captcha_str = args.num if args.num else DigitCaptcha.get_rand(3, 4)
        img = captcha.generate(captcha_str)
        img = np.fromstring(img.getvalue(), dtype='uint8')
        img = cv2.imdecode(img, cv2.IMREAD_GRAYSCALE)
        cv2.imwrite(args.output, img)
        print("Captcha image with digits {} written to {}".format([int(c) for c in captcha_str], args.output))

    main()