from flask import abort, current_app, flash, redirect, render_template, url_for
from flask_login import current_user, login_required
from werkzeug.utils import secure_filename
from . import services
from .forms import NewNLPJobForm, NewOCRJobForm
from .. import db
from ..models import Job, JobInput
import json
import os


SERVICES = {'nlp': {'name': 'Natural Language Processing',
                    'resources': {'mem_mb': 4096, 'n_cores': 2},
                    'new_job_form': NewNLPJobForm},
            'ocr': {'name': 'Optical Character Recognition',
                    'resources': {'mem_mb': 8192, 'n_cores': 4},
                    'new_job_form': NewOCRJobForm}}


@services.route('/<service_handle>', methods=['GET', 'POST'])
@login_required
def service(service_handle):
    if service_handle not in SERVICES:
        abort(404)
    new_job_form = SERVICES[service_handle]['new_job_form']()
    if new_job_form.validate_on_submit():
        _service_args = ['-l {}'.format(new_job_form.language.data)]
        if service_handle == 'ocr':
            if not new_job_form.binarization.data:
                _service_args.append('--skip-binarisation')
        job = Job(
            creator=current_user,
            description=new_job_form.description.data,
            mem_mb=SERVICES[service_handle]['resources']['mem_mb'],
            n_cores=SERVICES[service_handle]['resources']['n_cores'],
            service=service_handle,
            service_args=json.dumps(_service_args),
            service_version=new_job_form.version.data,
            status='preparing',
            title=new_job_form.title.data
        )
        db.session.add(job)
        db.session.commit()
        dir = os.path.join(current_app.config['OPAQUE_STORAGE_DIRECTORY'],
                           str(job.user_id),
                           'jobs',
                           str(job.id))
        try:
            os.makedirs(dir)
        except OSError:
            flash('OSError!')
            db.session.remove(job)
            db.session.commit()
        else:
            for file in new_job_form.files.data:
                filename = secure_filename(file.filename)
                file.save(os.path.join(dir, filename))
                job_input = JobInput(filename=filename, job=job)
                db.session.add(job_input)
            job.status = 'submitted'
            db.session.commit()
            flash('Job created!')
        return redirect(
            url_for('services.service', service_handle=service_handle)
        )
    return render_template('services/{}.html.j2'.format(service_handle),
                           title=SERVICES[service_handle]['name'],
                           new_job_form=new_job_form)