from flask_wtf import FlaskForm
from flask_login import current_user
from flask_wtf.file import FileField, FileRequired
from wtforms import (
    BooleanField,
    DecimalRangeField,
    MultipleFileField,
    SelectField,
    StringField,
    SubmitField,
    ValidationError
)
from wtforms.validators import InputRequired, Length
from app.models import SpaCyNLPPipelineModel, TesseractOCRPipelineModel
from . import SERVICES


class CreateJobBaseForm(FlaskForm):
    description = StringField(
        'Description',
        validators=[InputRequired(), Length(max=255)]
    )
    title = StringField(
        'Title',
        validators=[InputRequired(), Length(max=32)]
    )
    version = SelectField('Version', validators=[InputRequired()])
    submit = SubmitField()


class CreateFileSetupPipelineJobForm(CreateJobBaseForm):
    images = MultipleFileField('File(s)', validators=[InputRequired()])

    def validate_images(form, field):
        valid_mimetypes = ['image/jpeg', 'image/png', 'image/tiff']
        for image in field.data:
            if image.mimetype not in valid_mimetypes:
                raise ValidationError('JPEG, PNG and TIFF files only!')

    def __init__(self, *args, **kwargs):
        if 'prefix' not in kwargs:
            kwargs['prefix'] = 'create-file-setup-pipeline-job-form'
        service_manifest = SERVICES['file-setup-pipeline']
        version = kwargs.pop('version', service_manifest['latest_version'])
        super().__init__(*args, **kwargs)
        self.version.choices = [(x, x) for x in service_manifest['versions']]
        self.version.data = version
        self.version.default = service_manifest['latest_version']


class CreateTesseractOCRPipelineJobForm(CreateJobBaseForm):
    binarization = BooleanField('Binarization')
    pdf = FileField('File', validators=[FileRequired()])
    model = SelectField('Model', validators=[InputRequired()])
    ocropus_nlbin_threshold = DecimalRangeField(
        render_kw={'min': 0, 'max': 1, 'step': 0.1, 'start': [0.5], 'disabled': True}
    )

    def validate_binarization(self, field):
        service_info = SERVICES['tesseract-ocr-pipeline']['versions'][self.version.data]
        if field.data:
            if not('methods' in service_info and 'binarization' in service_info['methods']):
                raise ValidationError('Binarization is not available')
              
    def validate_pdf(self, field):
        if field.data.mimetype != 'application/pdf':
            raise ValidationError('PDF files only!')

    def __init__(self, *args, **kwargs):
        if 'prefix' not in kwargs:
            kwargs['prefix'] = 'create-tesseract-ocr-pipeline-job-form'
        service_manifest = SERVICES['tesseract-ocr-pipeline']
        version = kwargs.pop('version', service_manifest['latest_version'])
        super().__init__(*args, **kwargs)
        service_info = service_manifest['versions'][version]
        if self.binarization.render_kw is None:
            self.binarization.render_kw = {}
        self.binarization.render_kw['disabled'] = True
        if self.ocropus_nlbin_threshold.render_kw is None:
            self.ocropus_nlbin_threshold.render_kw = {}
        self.ocropus_nlbin_threshold.render_kw['disabled'] = True
        if 'methods' in service_info:
            if 'binarization' in service_info['methods']:
                del self.binarization.render_kw['disabled']
                if 'ocropus_nlbin_threshold' in service_info['methods']:
                    del self.ocropus_nlbin_threshold.render_kw['disabled']
        models = [
            x for x in TesseractOCRPipelineModel.query.order_by(TesseractOCRPipelineModel.title).all()
            if version in x.compatible_service_versions and (x.is_public == True or x.user == current_user)
        ]
        self.model.choices = [('', 'Choose your option')]
        self.model.choices += [(x.hashid, f'{x.title} [{x.version}]') for x in models]
        self.model.default = ''
        self.version.choices = [(x, x) for x in service_manifest['versions']]
        self.version.data = version
        self.version.default = service_manifest['latest_version']


class CreateTranskribusHTRPipelineJobForm(CreateJobBaseForm):
    binarization = BooleanField('Binarization')
    pdf = FileField('File', validators=[FileRequired()])
    model = SelectField('Model', validators=[InputRequired()])

    def validate_binarization(self, field):
        service_info = SERVICES['transkribus-htr-pipeline']['versions'][self.version.data]
        if field.data:
            if(
                'methods' not in service_info
                or 'binarization' not in service_info['methods']
            ):
                raise ValidationError('Binarization is not available')

    def validate_pdf(self, field):
        if field.data.mimetype != 'application/pdf':
            raise ValidationError('PDF files only!')

    def __init__(self, *args, **kwargs):
        if 'prefix' not in kwargs:
            kwargs['prefix'] = 'create-transkribus-htr-pipeline-job-form'
        transkribus_htr_pipeline_models = kwargs.pop('transkribus_htr_pipeline_models', [])
        service_manifest = SERVICES['transkribus-htr-pipeline']
        version = kwargs.pop('version', service_manifest['latest_version'])
        super().__init__(*args, **kwargs)
        service_info = service_manifest['versions'][version]
        if self.binarization.render_kw is None:
            self.binarization.render_kw = {}
        self.binarization.render_kw['disabled'] = True
        if 'methods' in service_info:
            if 'binarization' in service_info['methods']:
                del self.binarization.render_kw['disabled']
        self.model.choices = [('', 'Choose your option')]
        self.model.choices += [(x['modelId'], x['name']) for x in transkribus_htr_pipeline_models]
        self.model.default = ''
        self.version.choices = [(x, x) for x in service_manifest['versions']]
        self.version.data = version
        self.version.default = service_manifest['latest_version']


class CreateSpacyNLPPipelineJobForm(CreateJobBaseForm):
    encoding_detection = BooleanField('Encoding detection', render_kw={'disabled': True})
    txt = FileField('File', validators=[FileRequired()])
    model = SelectField('Model', validators=[InputRequired()])
    
    def validate_encoding_detection(self, field):
        service_info = SERVICES['spacy-nlp-pipeline']['versions'][self.version.data]
        if field.data:
            if(
                'methods' not in service_info
                or 'encoding_detection' not in service_info['methods']
            ):
                raise ValidationError('Encoding detection is not available')

    def validate_txt(form, field):
        if field.data.mimetype != 'text/plain':
            raise ValidationError('Plain text files only!')

    def __init__(self, *args, **kwargs):
        if 'prefix' not in kwargs:
            kwargs['prefix'] = 'create-spacy-nlp-pipeline-job-form'
        service_manifest = SERVICES['spacy-nlp-pipeline']
        version = kwargs.pop('version', service_manifest['latest_version'])
        super().__init__(*args, **kwargs)
        service_info = service_manifest['versions'][version]
        print(service_info)
        if self.encoding_detection.render_kw is None:
            self.encoding_detection.render_kw = {}
        self.encoding_detection.render_kw['disabled'] = True
        if 'methods' in service_info:
            if 'encoding_detection' in service_info['methods']:
                del self.encoding_detection.render_kw['disabled']
        user_models = [
            x for x in current_user.spacy_nlp_pipeline_models.order_by(SpaCyNLPPipelineModel.title).all()
        ]
        models = [
            x for x in SpaCyNLPPipelineModel.query.filter(SpaCyNLPPipelineModel.user != current_user, SpaCyNLPPipelineModel.is_public == True).order_by(SpaCyNLPPipelineModel.title).all()
            if version in x.compatible_service_versions
        ]
        self.model.choices = {
            '': [('', 'Choose your option')],
            'Your models': [(x.hashid, f'{x.title} [{x.version}]') for x in user_models],
            'Public models': [(x.hashid, f'{x.title} [{x.version}]') for x in models]
        }
        self.model.default = ''
        self.version.choices = [(x, x) for x in service_manifest['versions']]
        self.version.data = version
        self.version.default = version