# Implementace jobů pracujících se submity

from dataclasses import dataclass
import os
import re
import shutil
from sqlalchemy import and_
from sqlalchemy.orm import joinedload
from tempfile import NamedTemporaryFile
from typing import List, Optional
import unicodedata
import werkzeug.utils
import zipfile

import mo.db as db
from mo.jobs import TheJob, job_handler, job_file_path
from mo.submit import Submitter, SubmitException
from mo.util import logger, data_dir
from mo.util_format import inflect_number, inflect_by_number, data_size


def schedule_download_submits(paper_ids: List[int], description: str, for_user: db.User, want_subdirs: bool):
    the_job = TheJob()
    job = the_job.create(db.JobType.download_submits, for_user)
    job.description = description
    job.in_json = {'papers': paper_ids, 'want_subdirs': want_subdirs}
    the_job.submit()


@job_handler(db.JobType.download_submits)
def handle_download_submits(the_job: TheJob):
    """Zazipuje zadané papíry.

    Vstupní JSON:
        { 'papers': [ seznam paper_id ke stažení ],
          'want_feedback': true/false,
        }

    Výstupní JSON:
        null
    """

    job = the_job.job
    assert job.in_json is not None
    ids: List[int] = job.in_json['papers']  # type: ignore
    want_subdirs: bool = job.in_json['want_subdirs']  # type: ignore

    sess = db.get_session()
    papers = (sess.query(db.Paper, db.User, db.Task.code, db.Place)
              .select_from(db.Paper)
              .filter(db.Paper.paper_id.in_(ids))
              .join(db.User, db.User.user_id == db.Paper.for_user)
              .join(db.Task, db.Task.task_id == db.Paper.for_task)
              .join(db.Participation, db.Participation.user_id == db.Paper.for_user)
              .join(db.Contest, and_(db.Contest.master_contest_id == db.Participation.contest_id, db.Contest.round_id == db.Task.round_id))
              .join(db.Place, db.Place.place_id == db.Contest.place_id)
              .all())
    papers.sort(key=lambda p: (p[1].sort_key(), p[2]))

    temp_file = NamedTemporaryFile(suffix='.zip', dir=data_dir('tmp'), mode='w+b')
    logger.debug('Job: Vytvářím archiv %s', temp_file.name)

    cnt = 0
    with zipfile.ZipFile(temp_file, mode='w') as zip:
        for p, u, task_code, place in papers:
            cnt += 1
            full_name = u.full_name()
            ascii_name = (unicodedata.normalize('NFD', full_name)
                          .encode('ascii', 'ignore')
                          .decode('utf-8'))
            fn = f'{task_code}_{cnt:04d}_{u.user_id}_{p.paper_id}_{ascii_name}.pdf'
            fn = werkzeug.utils.secure_filename(fn)
            if want_subdirs:
                fn = f'{place.get_code()}/{fn}'
            logger.debug('Job: Přidávám %s', fn)
            zip.write(filename=os.path.join(data_dir('submits'), p.file_name or p.orig_file_name),
                      arcname=fn)

    job.out_file = the_job.attach_file(temp_file.name, '.zip')
    out_size = temp_file.tell()
    job.result = 'Celkem ' + inflect_number(cnt, 'soubor', 'soubory', 'souborů') + ', ' + data_size(out_size)
    temp_file.close()


def schedule_upload_feedback(round: db.Round, tmp_file: str, description: str, for_user: db.User,
                             only_contest: Optional[db.Contest],
                             only_site: Optional[db.Place],
                             only_region: Optional[db.Place],
                             only_task: Optional[db.Task]):
    the_job = TheJob()
    job = the_job.create(db.JobType.upload_feedback, for_user)
    job.description = description
    job.in_json = {
        'round_id': round.round_id,
        'only_contest_id': only_contest.contest_id if only_contest is not None else None,
        'only_site_id': only_site.place_id if only_site is not None else None,
        'only_region_id': only_region.place_id if only_region is not None else None,
        'only_task_id': only_task.task_id if only_task is not None else None,
    }
    job.in_file = the_job.attach_file(tmp_file, '.zip')
    the_job.submit()


@dataclass
class UploadFeedback:
    file_name: str
    task_code: str
    user_id: int
    task: Optional[db.Task] = None
    user: Optional[db.User] = None
    tmp_name: Optional[str] = None


def parse_feedback_name(name: str) -> Optional[UploadFeedback]:
    name = os.path.basename(name)

    # Formát jmen generovaný dávkovým stahováním
    m = re.match(r'(?P<task>[^_]+)_(?P<order>\d+)_(?P<user_id>\d+)_', name)
    if m:
        return UploadFeedback(
            file_name=name,
            task_code=m['task'],
            user_id=int(m['user_id']),
        )

    # Formát jmen, pod kterými se ukládají jednotlivě stahovaná řešení
    m = re.match(r'(?P<task>[^_]+)_(reseni|opravene)_(?P<user_id>\d+)_', name)
    if m:
        return UploadFeedback(
            file_name=name,
            task_code=m['task'],
            user_id=int(m['user_id']),
        )

    # Formát jmen, pod kterými se dříve ukládala jednotlivě stahovaná řešení
    m = re.match(r'(?P<task>.+)-(reseni|opravene)-(?P<paper_id>\d+)\.', name)
    if m:
        paper = db.get_session().query(db.Paper).get(int(m['paper_id']))
        if paper:
            return UploadFeedback(
                file_name=name,
                task_code=m['task'],
                user_id=paper.for_user,
            )

    return None


@job_handler(db.JobType.upload_feedback)
def handle_upload_feedback(the_job: TheJob):
    """Uploaduje opravená řešení.

    Vstupní JSON:
        { 'round_id': <id>,
          'only_task_id': <id_or_null>,
          'only_contest_id': <id_or_null>,
          'only_site_id': <id_or_null>,
        }

    Výstupní JSON:
        null
    """

    job = the_job.job
    assert job.in_file is not None
    in_json = job.in_json
    assert in_json is not None
    round_id: int = in_json['round_id']  # type: ignore
    only_contest_id: Optional[int] = in_json['only_contest_id']  # type: ignore
    only_site_id: Optional[int] = in_json['only_site_id']  # type: ignore
    only_region_id: Optional[int] = in_json['only_region_id']  # type: ignore
    only_task_id: Optional[int] = in_json['only_task_id']  # type: ignore

    sess = db.get_session()
    round = sess.query(db.Round).get(round_id)
    assert round is not None

    if only_region_id is not None:
        only_region = sess.query(db.Place).get(only_region_id)
        assert only_region is not None
    else:
        only_region = None

    files: List[UploadFeedback] = []

    def parse_zip(in_path: str):
        try:
            with zipfile.ZipFile(in_path, mode='r') as zip:
                contents = zip.infolist()
                for item in contents:
                    if not item.is_dir():
                        fb = parse_feedback_name(item.filename)
                        if fb:
                            tmp_file = NamedTemporaryFile(dir=data_dir('tmp'), mode='w+b', delete=False)
                            logger.debug(f'Job: Extrahuji {item.filename} do {tmp_file.name}')
                            with zip.open(item) as item_file:
                                shutil.copyfileobj(item_file, tmp_file)
                            tmp_file.close()
                            fb.tmp_name = tmp_file.name
                            files.append(fb)
                        else:
                            the_job.error(f'Nerozpoznáno jméno souboru {item.filename}')
        except zipfile.BadZipFile:
            the_job.error('Chybný formát souboru. Je to opravdu ZIP?')

    def resolve_tasks(files):
        task_dict = {f.task_code: None for f in files}
        tasks = sess.query(db.Task).filter_by(round=round).filter(db.Task.code.in_(task_dict.keys())).all()
        for task in tasks:
            task_dict[task.code] = task
        for code, task in task_dict.items():
            if task is None:
                the_job.error(f'Neznámá úloha {code}')
        for f in files:
            f.task = task_dict[f.task_code]
            if f.task is None:
                the_job.error(f'{f.file_name}: Neznámá úloha {code}')
            elif only_task_id is not None and f.task.task_id != only_task_id:
                the_job.error(f'{f.file_name}: Tato úloha nebyla vybraná k nahrávání')

    def resolve_users(files):
        user_dict = {f.user_id: None for f in files}
        contest_dict = {f.user_id: None for f in files}
        site_id_dict = {f.user_id: None for f in files}

        rows = (sess.query(db.User, db.Participation, db.Contest)
                .select_from(db.Participation)
                .join(db.User, db.User.user_id == db.Participation.user_id)
                .join(db.Contest, db.Contest.master_contest_id == db.Participation.contest_id)
                .filter(db.Contest.round == round)
                .filter(db.Participation.user_id.in_(user_dict.keys()))
                .options(joinedload(db.Contest.place))
                .all())

        user_rights = {}
        for user, pion, contest in rows:
            user_dict[user.user_id] = user
            contest_dict[user.user_id] = contest
            site_id_dict[user.user_id] = pion.place_id
            rr = the_job.gatekeeper.rights_for_contest(contest)
            user_rights[user.user_id] = rr.can_upload_feedback()

        for f in files:
            f.user = user_dict[f.user_id]
            if not f.user:
                the_job.error(f'{f.file_name}: Neznámý účastník #{f.user_id}')
            elif only_contest_id is not None and contest_dict[f.user_id].contest_id != only_contest_id:
                the_job.error(f'{f.file_name}: Účastník leží mimo vybranou soutěž')
            elif only_site_id is not None and site_id_dict[f.user_id] != only_site_id:
                the_job.error(f'{f.file_name}: Účastník leží mimo vybrané soutěžní místo')
            elif only_region is not None and not the_job.gatekeeper.is_ancestor_of(only_region, contest_dict[f.user_id].place):
                the_job.error(f'{f.file_name}: Účastník leží mimo vybraný region')
            elif not user_rights[f.user_id]:
                the_job.error(f'{f.file_name}: K tomuto účastníkovi nemáte dostatečná oprávnění')

    def process_file(fb: UploadFeedback) -> bool:
        assert fb.user and fb.task
        paper = db.Paper(
            for_user_obj=fb.user,
            task=fb.task,
            type=db.PaperType.feedback,
            uploaded_by_obj=job.user,
        )
        try:
            assert fb.tmp_name
            smtr = Submitter()
            smtr.submit_paper(paper, fb.tmp_name)
            sess.add(paper)

            sol = (sess.query(db.Solution)
                   .filter_by(task=fb.task, user=fb.user)
                   .with_for_update()
                   .one())
            sol.final_feedback_obj = paper

            sess.commit()
            return True
        except SubmitException as e:
            the_job.error(f'{fb.file_name}: {e}')
            return False

    cnt_good = 0
    parse_zip(job_file_path(job.in_file))

    if not the_job.errors:
        resolve_tasks(files)
        resolve_users(files)

    if not the_job.errors:
        for f in files:
            if process_file(f):
                cnt_good += 1

    for f in files:
        if f.tmp_name is not None:
            os.unlink(f.tmp_name)

    job.result = (inflect_by_number(cnt_good, 'Nahrán', 'Nahrány', 'Nahráno')
                  + f' {cnt_good} z '
                  + inflect_number(len(files), 'souboru', 'souborů', 'souborů'))
    if the_job.errors:
        job.result += (', '
                       + inflect_by_number(len(the_job.errors), 'nastala', 'nastaly', 'nastalo')
                       + ' '
                       + inflect_number(len(the_job.errors), 'chyba', 'chyby', 'chyb'))