"""Clear out matching diffs.

Revision ID: 1cf84ddb034c
Revises: 4d7019b218d4
Create Date: 2013-10-28 08:46:17.516023

"""

# revision identifiers, used by Alembic.
revision = '1cf84ddb034c'
down_revision = '4d7019b218d4'

from alembic import context, op
from collections import defaultdict
from sqlalchemy.sql import and_, table, column
import os
import pickle
import sqlalchemy as sa

file_ = table('file',
              column('id', sa.Integer),
              column('sha1', sa.String),
              column('size', sa.Integer))

testcase = table('testcase',
                 column('id', sa.Integer),
                 column('output_type', sa.Enum('diff', 'image', 'text',
                                               name='output_type')))

testcaseresult = table('testcaseresult',
                       column('diff_id', sa.Integer),
                       column('submission_id', sa.Integer),
                       column('test_case_id', sa.Integer))


def no_diff(sha1, directory):
    path = os.path.join(directory, sha1[:2], sha1[2:4], sha1[4:])
    diff = pickle.load(open(path))
    return diff._diff is None


def upgrade():
    ### commands auto generated by Alembic - please adjust! ###
    directory = context.config.file_config.get('app:main', 'file_directory')

    conn = op.get_bind()

    # test cases whose results store diffs
    diff_tcs = [tc_id for (tc_id, output) in conn.execute(testcase.select())
                if output == 'diff']

    diff_to_tcrs = defaultdict(list)
    for (diff_id, s_id, tc_id) in conn.execute(
        testcaseresult.select().where(testcaseresult.c.test_case_id
                                      .in_(diff_tcs))):
        if diff_id:
            diff_to_tcrs[diff_id].append({'s_id': s_id, 'tc_id': tc_id})

    # Prune diffs that are not significantly unique
    diff_to_tcrs = {x: y for (x, y) in diff_to_tcrs.items() if len(y) > 64}

    # Find test case results with no diff
    empty_tcrs = []
    for (f_id, sha1, size) in conn.execute(file_.select()
                                           .where(file_.c.id
                                                  .in_(diff_to_tcrs.keys()))):
        if size == 0:
            print('Found unexpected zero byte diff on: {}'
                  .format(diff_to_tcrs[f_id]))
            empty_tcrs.extend(diff_to_tcrs[f_id])
        elif no_diff(sha1, directory):
            empty_tcrs.extend(diff_to_tcrs[f_id])

    if not empty_tcrs:  # Exit early
        return

    # Update diff to None for appropriate TestCaseResults
    conn.execute(testcaseresult.update().where(and_(
                testcaseresult.c.submission_id == sa.bindparam('s_id'),
                testcaseresult.c.test_case_id == sa.bindparam('tc_id')))
               .values(diff_id=None), empty_tcrs)
    ### end Alembic commands ###


def downgrade():
    ### commands auto generated by Alembic - please adjust! ###
    pass
    ### end Alembic commands ###
