#!/usr/bin/python3
# coding=utf-8

import db
from .utils import filenamefy

conn = db.get_connection()
try:
    cur = conn.cursor()
    try:
        new_lidentifier = None
        # Get or add language.
        cur.execute('select id, name, identifier from language order by id asc')
        languages = cur.fetchall()
        print('Languages:')
        for lid, lname, lidentifier in languages:
            print('  {}: {}'.format(lid, lname))

        new_lid = input("Enter language ID or 'n' for new): ")
        if new_lid == 'n':
            new_lname = input('Enter name of the new language: ')
            new_lidentifier = filenamefy(new_lname)
            cur.execute('insert into language (name, identifier) values (%s, %s) returning id',
                        (new_lname, new_lidentifier))
            new_lid = cur.fetchone()[0]
            print('Added new language "{}" with ID {} and identifier {}'.format(
                new_lname, new_lid, new_lidentifier))
        else:
            new_lid = int(new_lid)
            for lid, lname, lidentifier in languages:
                print(lid, lname, lidentifier)
                if lid == new_lid:
                    new_lidentifier = lidentifier
                    break
        if new_lidentifier is None:
            raise Exception('Language with ID {} does not exist'.format(new_lid))
        print('Selected langauge {}'.format(new_lid))
        print()

        # Get or add problem group.
        new_gidentifier = None
        cur.execute('select id, name, identifier from problem_group order by id asc')
        groups = cur.fetchall()
        print('Problem groups:')
        for gid, gname, gidentifier in groups:
            print('  {}: {}'.format(gid, gname))
        new_gid = input("Enter problem group ID or 'n' for new): ")
        if new_gid == 'n':
            new_gname = input('Enter name of the new problem group: ')
            new_gidentifier = filenamefy(new_gname)
            cur.execute('insert into problem_group (name, identifier) values (%s, %s) returning id',
                        (new_gname, new_gidentifier))
            new_gid = cur.fetchone()[0]
            print('Added new problem group "{}" with ID {} and identifier {}'.format(
                new_gname, new_gid, new_gidentifier))
        else:
            new_gid = int(new_gid)
            for gid, gname, gidentifier in groups:
                if gid == new_gid:
                    new_gidentifier = gidentifier
                    break
        if new_gidentifier is None:
            raise Exception('Group with ID {} does not exist'.format(new_gid))
        print('Selected problem group {}'.format(new_gid))
        print()

        # Add problem.
        new_pname = input('Enter name of the new problem: ')
        new_pidentifier = filenamefy(new_pname)
        cur.execute('insert into problem (language_id, problem_group_id, name, identifier, is_visible) values (%s, %s, %s, %s, %s) returning id',
                    (new_lid, new_gid, new_pname, new_pidentifier, True))
        new_pid = cur.fetchone()[0]
        print('Added new problem "{}" with ID {} and identifier {}'.format(
            new_pname, new_pid, new_pidentifier))
        print('Data files should be placed in "{}/problems/{}/{}"'.format(
            new_lidentifier, new_gidentifier, new_pidentifier))

    finally:
        cur.close()
    conn.commit()
except:
    conn.rollback()
    raise
finally:
    db.return_connection(conn)