import tacticenv

from pyasm.security import *
from pyasm.unittest import UnittestEnvironment, Sample3dEnvironment
from pyasm.common.spt_date import SPTDate

from .sql import *
from .search import *
from .transaction import *
from .database_impl import *
from pyasm.unittest import *
from pyasm.biz import Project
import unittest
from datetime import datetime
from dateutil import parser
import pytz

class SearchTest(unittest.TestCase):

    def _commit(self, impl):
        if impl.commit_on_schema_change():
            DbContainer.commit_thread_sql()

    def test_all(self):
        # start batch environment
        Batch()
        from pyasm.web.web_init import WebInit
        WebInit().execute()

        sample3d_test_env = Sample3dEnvironment()
        sample3d_test_env.create()

        test_env = UnittestEnvironment()
        test_env.create()

        db_res = DbResource.get_default('unittest')
        sql = DbContainer.get(db_res)
        impl = sql.get_database_impl()
        db_type = impl.get_database_type()
        if db_type == "PostgreSQL":
            self.prefix = '''"unittest"."public".'''
            self.sthpw_prefix = '''"sthpw"."public".'''
        elif db_type == "Sqlite":
            self.prefix = ""
            self.sthpw_prefix = ""
        elif db_type == "SQLServer":
            self.prefix = '''"unittest"."dbo".'''
            self.sthpw_prefix = '''"sthpw"."dbo".'''

        else:
            self.prefix = '''"unittest".'''
            self.sthpw_prefix = '''"sthpw".'''

        self.transaction = Transaction.get(create=True)
        try:

            self.person = Person.create( "5", "a",
                    "ComputerWorld", "1")
            self.person = Person.create( "4", "b",
                    "ComputerWorld", "3")
            self.person = Person.create( "3", "c",
                    "ComputerWorld", "3")
            self.person = Person.create( "2", "d",
                    "ComputerWorld", "4")
            self.person = Person.create( "1", "e",
                    "ComputerWorld", "5")

            self._test_no_id()

            self._test_order_by()
            self._test_search_key()
            self._test_search()
            self._test_multi_db_subselect()

            # FIXME: this requires sample3d project
            self._test_search_other_project()
            self._test_search_type()
            self._test_metadata()
            self._test_search_type_existence()
            self._test_project()
            self._test_search_filter()
            self._test_dates_search()
            self._test_child_search()
            self._test_parent_search()
            self._test_add_column_search()
            self._test_commit()
            self._test_set_value()
            self._test_search_set_value()
            self._test_get_by_statement()
            self._test_views()
            self._test_timezone()

        finally:
            self.transaction.rollback()
            Project.set_project('unittest')

            self._commit(impl)

            test_env.delete()
            sample3d_test_env.delete()

    def _test_no_id(self):

        sobject = SearchType.create("unittest/no_id")
        sobject.set_value("name", "Cindy")
        sobject.commit()
        first_id = sobject.get_id()
        self.assertEqual(first_id, sobject.get_value("test_id"))

        sobject = SearchType.create("unittest/no_id")
        sobject.set_value("name", "Mike")
        sobject.commit()
        second_id = sobject.get_id()
        self.assertEqual(second_id, sobject.get_value("test_id"))

        # test update
        sobject.set_value("name", "Michael")
        sobject.commit()

        # test search
        search = Search("unittest/no_id")
        sobjects = search.get_sobjects()
        self.assertEqual( 2, len(sobjects) )
        for sobject in sobjects:
            self.assertEqual( True, sobject.get_id() in [first_id, second_id])

    def _test_order_by(self):

        sobjects = Search.eval("@SOBJECT(unittest/person['@ORDER_BY','description desc, name_first desc'])")
        sobjects1 = Search.eval("@SOBJECT(unittest/person['@ORDER_BY','description, name_first'])")
        sobjects2 = Search.eval("@SOBJECT(unittest/person['@ORDER_BY','description desc, name_first'])")
        name_list = []
        name_list1 = []
        name_list2 = []
        for x in sobjects:
             name = x.get_value('name_last')
             name_list.append(name)
        for x in sobjects1:
             name = x.get_value('name_last')
             name_list1.append(name)
        for x in sobjects2:
             name = x.get_value('name_last')
             name_list2.append(name)
        self.assertEqual(name_list,['e','d','b','c','a'])
        self.assertEqual(name_list1,['a','c','b','d','e'])
        self.assertEqual(name_list2,['e','d','c','b','a'])

        search = Search('unittest/city')
        search.add_order_by('name')
        statement = search.get_statement()
        self.assertEqual(statement, '''SELECT %s"city".* FROM %s"city" ORDER BY "city"."name"''' % (self.prefix, self.prefix))

        search = Search('unittest/person')
        search.add_order_by('unittest/city.unittest/country.code')
        statement = search.get_statement()
        self.assertEqual(statement, '''SELECT %s"person".* FROM %s"person" LEFT OUTER JOIN %s"city" ON "person"."city_code" = "city"."code" LEFT OUTER JOIN %s"country" ON "city"."country_code" = "country"."code" ORDER BY "country"."code"''' % (self.prefix, self.prefix, self.prefix, self.prefix))

        search = Search('unittest/person')
        search.add_order_by('unittest/city.id')
        statement = search.get_statement()
        self.assertEqual(statement, '''SELECT %s"person".* FROM %s"person" LEFT OUTER JOIN %s"city" ON "person"."city_code" = "city"."code" ORDER BY "city"."id"''' % (self.prefix, self.prefix, self.prefix))

        search = Search('unittest/person')
        search.add_order_by('unittest/city.id desc')
        statement = search.get_statement()
        self.assertEqual(statement, '''SELECT %s"person".* FROM %s"person" LEFT OUTER JOIN %s"city" ON "person"."city_code" = "city"."code" ORDER BY "city"."id" desc''' % (self.prefix, self.prefix, self.prefix))

        # with the built-in order-by logic, order by code is added
        search.get_sobjects()
        statement = search.get_statement()
        self.assertEqual(statement, '''SELECT %s"person".* FROM %s"person" LEFT OUTER JOIN %s"city" ON "person"."city_code" = "city"."code" ORDER BY "city"."id" desc, "person"."code"''' % (self.prefix, self.prefix, self.prefix))

    def _test_get_by_statement(self):
        types = ['admin','ben','beth']
        joined_statements = []

        db_res = DbResource.get_default('unittest')
        sql = DbContainer.get(db_res)
        impl = sql.get_database_impl()
        db_type = impl.get_database_type()

        for type in types:
            select = Search('sthpw/login')
            select.add_filter('login', type)
            select.set_show_retired(False)

            if db_type != 'Sqlite':
                select.add_order_by("login")
            statement = select.get_statement()
            joined_statements.append(statement)

        if len(joined_statements) > 1:

            if db_type == 'Sqlite':
                joined_statements = ["%s"%x for x in joined_statements]
            else:
                joined_statements = ["(%s)"%x for x in joined_statements]
            statement = ' union all '.join(joined_statements)
        elif len(joined_statements) == 1:
            statement = joined_statements[0]

        logins =  Login.get_by_statement(statement)

        # This assumes these users actually exist in the database, which they
        # often don't
        #self.assertEqual(len(logins), 3)

        if db_type =='Sqlite':
            self.assertEqual(statement, """SELECT {0}"login".* FROM {0}"login" WHERE "login"."login" = 'admin' union all SELECT {0}"login".* FROM {0}"login" WHERE "login"."login" = 'ben' union all SELECT {0}"login".* FROM {0}"login" WHERE "login"."login" = 'beth'""".format(self.sthpw_prefix))
        else:

            self.assertEqual(statement, '''(SELECT {0}"login".* FROM {0}"login" WHERE "login"."login" = 'admin' ORDER BY "login"."login") union all (SELECT {0}"login".* FROM {0}"login" WHERE "login"."login" = 'ben' ORDER BY "login"."login") union all (SELECT {0}"login".* FROM {0}"login" WHERE "login"."login" = 'beth' ORDER BY "login"."login")'''.format(self.sthpw_prefix))

    def _test_add_column_search(self):
        search = Search('sthpw/task')
        search.add_filter('status','Assignment')
        search.add_column('id')
        search.add_column('timestamp')
        sobject = search.get_sobject()
        self.assertEqual(len(sobject.data), 2)
        self.assertEqual(True, sobject.has_value('timestamp'))

        search = Search('sthpw/task')
        search.add_filter('status','Assignment')
        sobject = search.get_sobject()

        #self.assertEqual(len(sobject.data), 27)
        self.assertEqual(True, sobject.has_value('assigned'))
        self.assertEqual(True, sobject.has_value('process'))

    def _test_child_search(self):

        from pyasm.biz import Task
        person = SearchType.create('unittest/person')
        person.set_value('name_first','pete')
        person.commit()
        for xrange in (1, 50):
            Task.create(person, 'process_AA','some task', 'admin', context='process_AA')
        person2 = SearchType.create('unittest/person')
        person2.set_value('name_first','jamie')
        person2.commit()

        person2_tasks = []
        for xrange in (1, 50):
            person2_tasks.append(Task.create(person2, 'process_BB','some task', 'admin', context='process_BB'))
        task_search = Search('sthpw/task')
        task_search.add_filters('process', ['process_AA', 'process_BB'])
        tasks = task_search.get_sobjects()
        search2 = Search('unittest/person')
        search2.add_relationship_filters(tasks)
        persons = search2.get_sobjects()

        search3 = Search('unittest/person')
        search3.add_relationship_search_filter(task_search)
        persons_fast = search3.get_sobjects()

        self.assertEqual(SObject.get_values(persons, 'id'), SObject.get_values(persons_fast, 'id'))
        self.assertEqual(SObject.get_values(persons_fast, 'name_first'), ['pete','jamie'])

        # if I retire all the tasks for person2
        for task in tasks:
            if task.get_value('process') =='process_BB':
                task.retire()

        task_search = Search('sthpw/task')
        task_search.add_filters('process', ['process_AA', 'process_BB'])
        tasks = task_search.get_sobjects()
        search4 = Search('unittest/person')
        search4.add_relationship_search_filter(task_search)
        persons_fast = search4.get_sobjects()

        search2 = Search('unittest/person')
        search2.add_relationship_filters(tasks)
        persons = search2.get_sobjects()

        self.assertEqual(SObject.get_values(persons, 'id'), SObject.get_values(persons_fast, 'id'))
        self.assertEqual(SObject.get_values(persons_fast, 'name_first'), ['pete'])

        # test add_filters() with an empty array
        task_search = Search('sthpw/task')
        task_search.add_filters('process', [])
        tasks = task_search.get_sobjects()
        self.assertEqual(tasks, [])
        expected = '''SELECT %s"task".* FROM %s"task" WHERE "task"."id" is NULL AND ("task"."s_status" != 'retired' or "task"."s_status" is NULL) ORDER BY "task"."search_type", "task"."search_code"'''%(self.sthpw_prefix, self.sthpw_prefix)
        statement = task_search.get_statement()
        self.assertEqual(statement, expected)

    def _test_parent_search(self):
        from pyasm.biz import Task
        person = SearchType.create('unittest/person')
        person.set_value('name_first','burt')
        person.commit()

        person2 = SearchType.create('unittest/person')
        person2.set_value('name_first','sean')
        person2.commit()

        person_search = Search('unittest/person')
        person_search.add_filters('id', [person.get_id(), person2.get_id()])
        for xrange in (1, 50):
            Task.create(person, 'process_CC','some task', 'admin', context='process_CC')
        for xrange in (1, 50):
            Task.create(person2, 'process_DD','some task', 'admin')

        # find parent of tasks
        search2 = Search('sthpw/task')
        search2.add_relationship_filters([person,person2])
        tasks = search2.get_sobjects()

        search3 = Search('sthpw/task')
        search3.add_relationship_search_filter(person_search)
        tasks2 = search3.get_sobjects()

        self.assertEqual(SObject.get_values(tasks, 'id'), SObject.get_values(tasks2, 'id'))

    def _test_search_other_project(self):
        # while in unittest, search for pacman assets
        search = Search('sthpw/search_object')
        search.add_filter('search_type','prod/asset')
        sobject = search.get_sobject()
        # make sure the database field is not hardcoded pacman
        sobject.set_value('database', '{project}')
        sobject.commit()

        asset = SearchType.create('prod/asset?project=sample3d')
        asset.set_value('code','unittest001')
        asset.set_value('pipeline_code','model')
        asset.commit()

        asset = SearchType.create('prod/asset?project=sample3d')
        asset.set_value('code','unittest002')
        asset.set_value('pipeline_code','model')
        asset.commit()

        search2 = Search("prod/asset?project=sample3d")
        search2.add_filter("pipeline_code", "model")
        sobjects = search2.get_sobjects()
        self.assertEqual(True, len(sobjects) > 0 )
        self.assertEqual("model", sobjects[0].get_value('pipeline_code'))
        self.assertEqual("sample3d", sobjects[0].get_database())
        self.assertEqual("sample3d", sobjects[1].get_database())

        search3 = Search("prod/asset?project=sample3d")
        search3.add_regex_filter("pipeline_code", "model|cg_asset")
        search3.add_filter("code", "chr003")
        sobject = search3.get_sobject()
        if sobject:
            sobject.set_value('description','some char')
            sobject.commit()

            updated_sobject = search3.get_sobject(redo=True)
            self.assertEqual("some char", updated_sobject.get_value('description'))

        search4 = Search("prod/asset?project=sample3d")
        search4.add_regex_filter("pipeline_code", "cg_asset's")

        sobject = search4.get_sobject()
        self.assertEqual(sobject, None)

    def _test_search(self):

        # create the person
        person = SObjectFactory.create("unittest/person")
        person.set_value("name_first", "XXX")
        person.set_value("name_last", "Smith")
        person.set_value("nationality", "Canada")
        person.commit()

        search = Search("unittest/person")
        search.add_filter("name_first", "XXX")
        search.add_filter("name_last", "Smith")
        person = search.get_sobject()
        self.assertEqual("XXX", person.get_value("name_first") )

        xxx_id = person.get_id()

        # test_set_value
        person.set_value("name_first", "YYY")
        person.set_value("name_last", "Cowser")
        person.commit()

        search2 = Search("unittest/person")
        search2.add_filter("name_first", "YYY")
        search2.add_filter("name_last", "Cowser")
        person2 = search2.get_sobject()

        yyy_id = person2.get_id()

        self.assertEqual( xxx_id, yyy_id )

        # test attributes
        name_first = person.get_attr_value("name_first")
        self.assertEqual("YYY", name_first )

        name_last = person.get_attr_value("name_last")
        self.assertEqual("Cowser", name_last )

        # test related search
        sto_search = Search('prod/shot?project=sample3d')
        sobjects = sto_search.get_sobjects()
        self.assertEqual(len(sobjects) > 1, True)

        search = Search('sthpw/login')

        # this should not cause errors if there is schema connection between storyboard and login
        login_sobjects = search.add_relationship_filters(sobjects)
        if sobjects:
            login_sobject = search.add_relationship_filter(sobjects[0])

    def _test_search_type(self):
        '''test that search types behave properly for templating'''
        SearchType.set_global_template("project", "unittest")

        #search_type = SearchType.get("prod/asset")
        sobject = SearchType.create("unittest/person")
        key = sobject.get_search_type()
        self.assertEqual(key, "unittest/person?project=unittest")

    def _test_metadata(self):
        # create an arbitrary sobject
        person = SearchType.create("unittest/person")

        # get metadata that does not exist
        cow = person.get_metadata_value("cow")
        self.assertEqual("", cow)

        # json dict doesn't have this method
        #person.get_metadata_xml().clear_xpath_cache()

        # set the value
        person.set_metadata_value("cow", "angus")
        cow = person.get_metadata_value("cow")
        self.assertEqual("angus", cow)

        # add a whole bunch of entries:
        data = {
            'color': 'blue',
            'height': '195cm',
            'weight': '102Kg'
        }
        person.add_metadata(data)

        color = person.get_metadata_value("color")
        height = person.get_metadata_value("height")
        self.assertEqual("blue", color)
        self.assertEqual("195cm", height)

        # replace the entire structure
        person.add_metadata(data, replace=True)
        cow = person.get_metadata_value("cow")
        self.assertEqual("", cow)

        # if someone set it with non-json data, it should still return as is
        person.metadata = None
        person.set_value('metadata','testing==OK')
        meta_dict = person.get_metadata_dict()
        self.assertEqual('testing==OK', meta_dict)

    def _test_search_type_existence(self):
        '''test a bunch of search types to see if they are registered'''
        # DEPRECATED
        return
        #TODO: add more search types here
        st_list = ['prod/shot_instance','prod/sequence_instance','sthpw/widget_extend']
        for st in st_list:
            search = Search(SearchType.SEARCH_TYPE)
            search.add_filter('search_type', st)
            sobjects = search.get_sobjects()
            self.assertEqual(len(sobjects), 1)

    def _test_search_key(self):

        # test a correct search_key
        search_key = "prod/asset?project=sample3d&code=chr001"

        search_type = SearchKey.extract_base_search_type(search_key)
        self.assertEqual("prod/asset", search_type)

        project = SearchKey.extract_project(search_key)
        self.assertEqual("sample3d", project)

        code = SearchKey.extract_code(search_key)
        self.assertEqual("chr001", code)

        # test a long one
        search_key = "prod/asset?project=sample3d&db=postgres&code=chr001"

        search_type = SearchKey.extract_base_search_type(search_key)
        self.assertEqual("prod/asset", search_type)

        database = SearchKey.extract_database(search_key)
        self.assertEqual("postgres", database)

        project = SearchKey.extract_project(search_key)
        self.assertEqual("sample3d", project)

        code = SearchKey.extract_code(search_key)
        self.assertEqual("chr001", code)

        #eval(@GET(prod/asset.code))
        #eval(@GET(prod/asset.code?project=sample3d)) ???

        search_key = "sthpw/project"
        code = "unittest"

        search = Search(search_key)
        search.add_filter("code", code)
        sobject = search.get_sobject()

        search_key = SearchKey.get_by_sobject(sobject)
        self.assertEqual("sthpw/project?code=unittest", search_key)

    def _test_project(self):

        from pyasm.biz import Project
        sql = Project.get_database_impl()

        # Don't bother running if you don't have sample3d
        if not sql.database_exists('sample3d'):
            return

        Project.set_project("sthpw")

        sobject = SearchType.create('prod/shot?project=sample3d', columns=['code', 'sequence_code', 'pipeline_code'], result=['S001','HT001','shot'])
        self.assertEqual("prod/shot?project=sample3d", sobject.get_search_type())

        if sql.database_exists('sample3d'):
            db_resource = Project.get_db_resource_by_search_type('prod/bin?project=sample3d')
            exists= sql.table_exists(  db_resource ,'bin')
            if exists:
                search = Search('prod/bin', project_code='sample3d')
                self.assertEqual("prod/bin?project=sample3d", search.get_search_type())
        # check that a search type is properly created
        search_type = SearchType.get("prod/shot?project=sample3d")
        base_key = search_type.get_base_key()
        self.assertEqual("prod/shot", base_key)

        # NOTE: search_type get_full_key() method is deprecated.

        # test that the sobject maintains the search type
        sobject = SearchType.create("prod/shot?project=sample3d")
        search_type = sobject.get_search_type()
        self.assertEqual("prod/shot?project=sample3d", search_type)

        # set it back to unittest
        Project.set_project("unittest")

        # test current project is added when there is not project set
        sobject = SearchType.create("prod/shot")
        search_type = sobject.get_search_type()
        self.assertEqual("prod/shot?project=unittest", search_type)

        # test current project is added when there is not project set, even
        # when the project has changed
        sobject = SearchType.create("prod/shot")

        search_type = sobject.get_search_type()
        self.assertEqual("prod/shot?project=unittest", search_type)

        if sql.database_exists('sample3d'):
            Project.set_project("sample3d")

            project_code = Project.get_project_code()
            self.assertEqual("sample3d", project_code)

        # set it back to unittest project
        Project.set_project("unittest")

        # test the search
        if sql.database_exists('sample3d'):
            search_type = "prod/shot?project=sample3d"
            search = Search(search_type)
            project_code = search.get_project_code()
            self.assertEqual("sample3d", project_code)

            # test the search project code even though the project has hanved
            search_type = "prod/shot?project=sample3d"
            search = Search(search_type)
            project_code = search.get_project_code()
            self.assertEqual("sample3d", project_code)

            Project.set_project("admin")
            project_code = search.get_project_code()
            self.assertEqual("sample3d", project_code)

            project_code = Project.get_project_code()
            self.assertEqual("admin", project_code)

        # set it back to unittest project
        Project.set_project("unittest")

    def _test_search_filter(self):
        search = Search("unittest/person")
        search.add_filter("name_first", "YYY")
        search.add_filter("name_last", "Cowser")

        statement = search.get_statement()
        expected = """SELECT %s"person".* FROM %s"person" WHERE "person"."name_first" = 'YYY' AND "person"."name_last" = 'Cowser'""" % (self.prefix, self.prefix)

        self.assertEqual(expected, statement)

        search2 = Search("unittest/city")
        search2.add_search_filter("id", search)

        statement = search2.get_statement()
        expected = """SELECT {0}"city".* FROM {0}"city" WHERE "city"."id" in ( SELECT {0}"person".* FROM {0}"person" WHERE "person"."name_first" = 'YYY' AND "person"."name_last" = 'Cowser' )""".format(self.prefix)

        self.assertEqual(expected, statement)

        search = Search("unittest/person")
        search.add_op('begin')
        search.add_filter("name_first", "YYY")
        search.add_filter("name_last", "Cowser")
        search.add_op('or')
        search.add_filter("city_code", "YYZ")
        statement = search.get_statement()
        expected = """SELECT {0}"person".* FROM {0}"person" WHERE ( "person"."name_first" = 'YYY' OR "person"."name_last" = 'Cowser' ) AND "person"."city_code" = 'YYZ'""".format(self.prefix)

        self.assertEqual(expected, statement)

        # mix "is not" with "is" logic
        search = Search("unittest/person")

        search.add_op('begin')
        search.add_filter('name_first', 'T', op='~')
        search.add_filter('name_first', 'U', op='~')
        search.add_op('or')

        search.add_op('begin')
        search.add_filter('nationality','canadian', op='!=')
        search.add_filter('nationality', None)
        search.add_op('or')

        search.add_op('begin')
        search.add_filter("name_first", "YYY")
        search.add_filter("name_last", "Cowser")
        search.add_filter("city_code", "YYZ")
        search.add_op('or')

        # this is optional
        #search.add_op('and')
        expected = """SELECT {0}"person".* FROM {0}"person" WHERE ( "person"."name_first" ~ 'T' OR "person"."name_first" ~ 'U' ) AND ( "person"."nationality" != 'canadian' OR "person"."nationality" is NULL ) AND ( "person"."name_first" = 'YYY' OR "person"."name_last" = 'Cowser' OR "person"."city_code" = 'YYZ' )""".format(self.prefix)

        statement = search.get_statement()
        self.assertEqual(expected, statement)

        # old style without begin
        search = Search('sthpw/login')
        search.add_where("\"login\".\"license_type\" = 'user'")
        search.add_where("\"login\".\"license_type\" is NULL")
        search.add_op('or')

        expected = """SELECT {0}"login".* FROM {0}"login" WHERE "login"."license_type" = \'user\' OR "login"."license_type" is NULL""".format(self.sthpw_prefix)
        statement = search.get_statement()
        self.assertEqual(expected, statement)

        # extra dangling begin, default back to AND
        search = Search('sthpw/login')
        #search.add_op('begin')
        search.add_op('begin')
        search.add_where("\"login\".\"license_type\" = 'user'")
        search.add_where("\"login\".\"license_type\" is NULL")
        #search.add_op('or')

        expected = """SELECT {0}"login".* FROM {0}"login" WHERE "login"."license_type" = \'user\' AND "login"."license_type" is NULL""".format(self.sthpw_prefix)
        statement = search.get_statement()
        self.assertEqual(expected, statement)

        # extra dangling begin, with 2 ORs applied already
        search = Search('sthpw/login')
        search.add_op('begin')
        search.add_op('begin')
        search.add_where("\"login\".\"license_type\" = 'user'")
        search.add_where("\"login\".\"license_type\" is NULL")
        search.add_op('or')

        search.add_op('begin')
        search.add_where("\"login\".\"license_type\" = 'float'")
        search.add_op('or')

        expected = """SELECT {0}"login".* FROM {0}"login" WHERE ( "login"."license_type" = \'user\' OR "login"."license_type" is NULL ) AND "login"."license_type" = \'float\'""".format(self.sthpw_prefix)
        statement = search.get_statement()
        self.assertEqual(expected, statement)

        # extra dangling begin, with 2 ORs applied already
        search = Search('sthpw/login')
        search.add_op('begin')
        search.add_filter('namespace','sample3d')
        search.add_op('begin')
        search.add_where("\"login\".\"license_type\" = 'user'")
        search.add_where("\"login\".\"license_type\" is NULL")
        search.add_op('or')

        search.add_op('begin')
        search.add_where("\"login\".\"license_type\" = 'float'")
        search.add_op('or')
        search.add_op('or')

        expected = """SELECT {0}"login".* FROM {0}"login" WHERE "login"."namespace" = \'sample3d\' OR ( "login"."license_type" = \'user\' OR "login"."license_type" is NULL ) OR "login"."license_type" = \'float\'""".format(self.sthpw_prefix)
        statement = search.get_statement()
        self.assertEqual(expected, statement)

        # triple ands on level 0, OR on level 1
        search = Search("sthpw/task")

        search.add_op('begin')

        search.add_op('begin')
        search.add_filter('process', 'anim')
        search.add_filter('status', 'Pending')
        search.add_filter('assigned' , 'admin')
        search.add_op('and')

        search.add_op('begin')
        search.add_filter('process', 'layout')
        search.add_filter('status', 'Assignment')
        search.add_filter('assigned' , 'fil')
        search.add_op('and')

        search.add_op('or')

        expected = """SELECT {0}"task".* FROM {0}"task" WHERE ( "task"."process" = \'anim\' AND "task"."status" = \'Pending\' AND "task"."assigned" = \'admin\' ) OR ( "task"."process" = \'layout\' AND "task"."status" = \'Assignment\' AND "task"."assigned" = \'fil\' )""".format(self.sthpw_prefix)

        statement = search.get_statement()
        self.assertEqual(expected, statement)

        # test is NULL and is not NULL
        search = Search("unittest/person")
        search.add_op_filters([("name_first", 'is', 'NULL')])

        statement = search.get_statement()
        expected = """SELECT {0}"person".* FROM {0}"person" WHERE "person"."name_first" is NULL""".format(self.prefix)
        self.assertEqual(expected, statement)

        # lowercase null treated as string
        search = Search("unittest/person")
        search.add_op_filters([("name_first", 'is', 'null')])

        statement = search.get_statement()
        expected = """SELECT {0}"person".* FROM {0}"person" WHERE "person"."name_first" is 'null'""".format(self.prefix)
        self.assertEqual(expected, statement)

        search = Search("unittest/person")
        search.add_op_filters([("name_last", 'is not', 'NULL')])

        statement = search.get_statement()
        expected = """SELECT {0}"person".* FROM {0}"person" WHERE "person"."name_last" is not NULL""".format(self.prefix)
        self.assertEqual(expected, statement)

    def _test_relationship_filter(self):

        search = Search("unittest/person")
        search.add_filter("name_first", "YYY")
        search.add_filter("name_last", "Cowser")

        search2 = Search("unittest/city")
        search2.add_search_filter("id", search)

    def _test_dates_search(self):
        import time
        timezone = time.tzname[time.daylight]

        from pyasm.biz import Project
        database_type = Project.get_by_code("unittest").get_database_type()
        start_range = SPTDate.convert_to_timezone("2010-01-01", timezone)
        end_range = SPTDate.convert_to_timezone("2010-02-02", timezone)

        if database_type == "MySQL":
            start_range = start_range.replace(tzinfo=None)
            end_range = end_range.replace(tzinfo=None)

        search = Search("unittest/person")
        search.add_date_range_filter("birth_date", "2010-01-01", "2010-02-01")
        expected = """SELECT {0}"person".* FROM {0}"person" WHERE "person"."birth_date" >= '{1}' AND "person"."birth_date" < '{2}'""".format(self.prefix, start_range, end_range)
        self.assertEqual(expected, search.get_statement() )

        search = Search("unittest/person")
        search.add_dates_overlap_filter("birth_date", "birth_date", "2010-01-01", "2010-02-01")
        expected = '''SELECT {0}"person".* FROM {0}"person" WHERE "person"."id" in (SELECT {0}"person"."id" FROM {0}"person" '''.format(self.prefix)
        expected += '''WHERE ( "person"."birth_date" <= '{0}' AND "person"."birth_date" >= '{0}' ) '''.format(start_range)
        expected += '''OR ( "person"."birth_date" >= '{0}' AND "person"."birth_date" <= '{0}' ) '''.format(end_range)
        expected += '''OR ( "person"."birth_date" >= '{0}' AND "person"."birth_date" <= '{1}' ))'''.format(start_range, end_range)
        self.assertEqual(expected, search.get_statement() )

    def _test_commit(self):

        from pyasm.biz import Project
        database_type = Project.get_by_code("unittest").get_database_type()
        if database_type == "MySQL":
            print
            print("WARNING: !!!!!!!")
            print("_test_commit is disabled")
            print("WARNING: !!!!!!!")
            print
            return

        person_s = Search('unittest/person')
        person_s.add_filter('name_first','pete')

        person = person_s.get_sobject()
        from pyasm.biz import Note
        Note.create(person, "3 slashes \\\\\\", context="unittest_commit", process="unittest_commit")
        search = Search('sthpw/note')
        search.add_filter('process','unittest_commit')
        search.set_limit(1)
        note = search.get_sobject()
        self.assertEqual(note.get_value('note'), "3 slashes \\\\\\")

    def _test_set_value(self):
        ''' test with different Database Impl'''
        update = Update()
        update.set_database('sthpw')
        update.set_table('task')
        update.set_value('timestamp','2012-12-12')
        self.assertEqual( update.get_statement(), """UPDATE {0}"task" SET "timestamp" = '2012-12-12'""".format(self.sthpw_prefix))

        update = Update()
        update.set_database('sthpw')
        update.set_table('task')

        # Changing database to SQLServer
        sql_impl = DatabaseImpl.get('SQLServer')

        update.impl = sql_impl
        value = '2012-12-12'
        value = update.impl.process_date(value)

        update.set_value('timestamp', value)
        self.assertEqual( update.get_statement(), """UPDATE {0}"task" SET "timestamp" = convert(datetime2, \'2012-12-12\', 0)""".format(self.sthpw_prefix))
        update.set_value('timestamp','NOW')
        self.assertEqual( update.get_statement(), """UPDATE {0}"task" SET "timestamp" = getdate()""".format(self.sthpw_prefix))

        from pyasm.biz import Project
        database_type = Project.get_by_code("unittest").get_database_type()

        # time_dict = {'SQLServer': "convert(datetime2, '2012-12-25', 0)",
        #             'Sqlite':"'2012-12-25'",
        #             'PostgreSQL':"'2012-12-25'",
        #             'MySQL': "'2012-12-25 00:00:00'"}
        #             #'Oracle':"TO_DATE('2012-12-25','YYYY-MM-DD'"}
        # #TODO: test with cx_Oracle installed
        # #TODO: have another test with timezone considered
        # for db_type in ['Sqlite','SQLServer','MySQL','PostgreSQL']:
        #     sql_impl = DatabaseImpl.get(db_type)
        #     update = Update()
        #     update.set_database('sthpw')
        #     update.set_table('task')
        #     update.impl = sql_impl
        #     value = '2012-12-25'
        #     value = update.impl.process_date(value)
        #     update.set_value('timestamp', value)
        #     update.set_value('description','')
        #     if db_type == 'SQLServer':
        #         self.assertEqual( update.get_statement(), """UPDATE %s"task" SET "timestamp" = %s, "description" = N\'\'"""% (self.sthpw_prefix, time_dict.get(db_type)))
        #     else:
        #         self.assertEqual( update.get_statement(), """UPDATE %s"task" SET "timestamp" = %s, "description" = \'\'"""% (self.sthpw_prefix, time_dict.get(db_type)))

        #     update.set_value('description',None)
        #     self.assertEqual( update.get_statement(), """UPDATE %s"task" SET "timestamp" = %s, "description" = NULL"""% (self.sthpw_prefix, time_dict.get(db_type)))

    def _test_timezone(self):

        # timestamp without timezone: set value without timezone conversion

        task = SearchType.create('sthpw/task')
        input_time = datetime.utcnow().replace(microsecond=0)

        task.set_value('timestamp', input_time)
        task.set_value('description', 'search_test')
        output_time = task.get_datetime_value('timestamp')
        self.assertEqual(input_time, output_time)
        task.commit()
        output_time = task.get_datetime_value('timestamp')
        self.assertEqual(input_time, output_time)

        # timestamp without timezone: set value with timezone conversion

        task1 = SearchType.create('sthpw/task')
        input_time1 = input_time
        local_tz = pytz.timezone("Europe/Moscow")
        tz_input_time = input_time1.replace(tzinfo=pytz.utc).astimezone(local_tz)

        task1.set_value('timestamp', tz_input_time)
        output_time = task1.get_datetime_value('timestamp')
        self.assertEqual(SPTDate.convert_to_timezone(tz_input_time, 'UTC'), output_time)
        task1.commit()
        output_time = task1.get_datetime_value('timestamp')
        self.assertEqual(input_time1, output_time)

	    # search for task using both timestamps

        search = Search('sthpw/task')
        search.add_filter('description', 'search_test')
        search.add_filter('timestamp', input_time)
        search_result = search.get_sobject()
        self.assertEqual(task.get_code(), search_result.get_code())

        search = Search('sthpw/task')
        search.add_filter('description', 'search_test')
        search.add_filter('timestamp', tz_input_time)
        search_result = search.get_sobject()
        self.assertEqual(task.get_code(), search_result.get_code())

        # timestamp with timezone: set value without timezone conversion

        ticket = SearchType.create('sthpw/ticket')
        input_time = datetime.utcnow().replace(microsecond=0)

        ticket.set_value('timestamp', input_time)
        ticket.set_value('login', 'admin')
        ticket.set_value('ticket', 'search_test')
        ticket.set_value('expiry', input_time)
        output_time = ticket.get_datetime_value('timestamp')
        self.assertEqual(input_time, output_time)
        ticket.commit()
        output_time = ticket.get_datetime_value('timestamp')
        self.assertEqual(input_time, output_time)

        # timestamp with timezone: set value with timezone conversion

        ticket1 = SearchType.create('sthpw/ticket')
        input_time1 = input_time
        local_tz = pytz.timezone("Europe/Moscow")
        tz_input_time = input_time1.replace(tzinfo=pytz.utc).astimezone(local_tz)

        ticket1.set_value('timestamp', tz_input_time)
        ticket1.set_value('login', 'admin')
        ticket1.set_value('ticket', 'test2')
        ticket1.set_value('expiry', tz_input_time)
        output_time = ticket1.get_datetime_value('timestamp')
        self.assertEqual(SPTDate.convert_to_timezone(tz_input_time, 'UTC'), output_time)
        ticket1.commit()
        output_time = ticket1.get_datetime_value('timestamp')
        self.assertEqual(input_time1, output_time)

	    # search ticket using both timestamps

        search = Search('sthpw/ticket')
        search.add_filter('ticket', 'search_test')
        search.add_filter('timestamp', input_time)
        search_result = search.get_sobject()
        self.assertEqual(ticket.get_code(), search_result.get_code())

        search = Search('sthpw/ticket')
        search.add_filter('ticket', 'search_test')
        search.add_filter('timestamp', tz_input_time)
        search_result = search.get_sobject()
        self.assertEqual(ticket.get_code(), search_result.get_code())

        # test set_value, get_value with string date

        input_time_string = "2019-07-02 18:17:42"
        input_time_date = parser.parse(input_time_string)

        ticket = SearchType.create('sthpw/ticket')
        ticket.set_value('timestamp', input_time_string)
        ticket.set_value('login', 'admin')
        ticket.set_value('ticket', 'set_test')
        ticket.set_value('expiry', input_time_string)
        output_time = ticket.get_datetime_value('timestamp')
        self.assertEqual(input_time_date, output_time)
        ticket.commit()
        output_time = ticket.get_datetime_value('timestamp')
        self.assertEqual(input_time_date, output_time)

        #Test set_value, get_value with timezoned string date

        input_time_string_tz = "2019-07-02 18:17:42+03:00"
        input_time_date_tz = parser.parse(input_time_string_tz)
        input_time_date_utc = SPTDate.convert_to_timezone(input_time_date_tz, 'UTC').replace(tzinfo=None)

        ticket = SearchType.create('sthpw/ticket')
        ticket.set_value('timestamp', input_time_string_tz)
        ticket.set_value('login', 'admin')
        ticket.set_value('ticket', 'set_test_tz')
        ticket.set_value('expiry', input_time_string_tz)
        output_time = ticket.get_datetime_value('timestamp')
        self.assertEqual(input_time_date_tz, output_time)
        ticket.commit()
        output_time = ticket.get_datetime_value('timestamp')
        self.assertEqual(input_time_date_utc, output_time)

    def _test_multi_db_subselect(self):

        from pyasm.biz import Task
        person = SearchType.create('unittest/person')
        person.set_value('name_first','carin')
        person.commit()
        Task.create(person, 'subselect','some task', 'admin', status="a")

        from pyasm.biz import Task
        person = SearchType.create('unittest/person')
        person.set_value('name_first','carin2')
        person.commit()
        Task.create(person, 'subselect','some task', 'admin', status="b")

        from pyasm.biz import Task
        person = SearchType.create('unittest/person')
        person.set_value('name_first','carin3')
        person.commit()
        Task.create(person, 'subselect','some task', 'admin', status="c")

        # find people by carin that have a task "subselect"
        search_person = Search("unittest/person")
        search_task = Search("sthpw/task")
        search_person.add_relationship_search_filter(search_task,use_multidb=False)
        search_task.add_filter("process", "subselect")

        sobjects = search_person.get_sobjects()
        self.assertEqual(len(sobjects), 3)

        # find people by carin that have a task "subselect", using
        # subselect
        search_person = Search("unittest/person")
        search_task = Search("sthpw/task")
        search_person.add_relationship_search_filter(search_task,use_multidb=True)
        search_task.add_filter("process", "subselect")

        sobjects = search_person.get_sobjects()
        self.assertEqual(len(sobjects), 3)

        # find task that carin has
        search_task = Search("sthpw/task")
        search_person = Search("unittest/person")
        search_person.add_filter('name_first','carin')
        search_task.add_relationship_search_filter(search_person,use_multidb=True)

        sobjects = search_task.get_sobjects()

        # test an order by on a task status (cross database join)
        # Note this will not work on Postgres or SQLite which do not
        # support cross database joins
        search_person = Search("unittest/person")
        search_person.add_order_by("sthpw/task.status", direction="desc")
        statement = search_person.get_statement()

        can_join = DatabaseImpl.can_search_types_join( \
                "unittest/person", "sthpw/task")

        if can_join:
            expected = '''SELECT %s"person".* FROM %s"person" LEFT OUTER JOIN %s"task" ON "person"."code" = "task"."search_code" WHERE "task"."search_type" = 'unittest/person?project=unittest' ORDER BY "task"."status" desc''' % (self.prefix, self.prefix, self.sthpw_prefix)
            self.assertEqual(expected, statement)

            sobjects = search_person.get_sobjects()
            names = [x.get_value("name_first") for x in sobjects]
            expected = ['carin3','carint2','carin']
            self.assertEqual(expected, names)

    def _test_search_set_value(self):
        #TODO, test other db impl
        sobject = Search.eval("@SOBJECT(sthpw/task['@LIMIT','1'])", single=True)
        sobject.set_value('timestamp','2012-12-12')
        sobject.set_value('description', '')
        sobject.commit(triggers=False)
        self.assertEqual( sobject.data.get('description'), None)
        self.assertEqual( sobject.get_value('description'), '')

    def _test_views(self):

        car = SearchType.create('unittest/car')
        car.set_value('model','Nissan Sentra')
        car.set_value('class','sedan')
        car.commit()

        car = SearchType.create('unittest/car')
        car.set_value('model','Ferrari 958')
        car.set_value('class','sports_car')
        car.commit()

        car = SearchType.create('unittest/car')
        car.set_value('model','Porsche 911')
        car.set_value('class','sports_car')
        car.commit()

        search = Search("unittest/car")
        cars = search.get_sobjects()

        search = Search("unittest/sports_car")
        sports_cars = search.get_sobjects()

        self.assertEqual(3, len(cars))
        self.assertEqual(2, len(sports_cars))

        #for sports_car in sports_cars:
        #    sports_car.set_value("top_speed", "200mph")
        #    sports_car.commit()

if __name__ == "__main__":
    unittest.main()


from django.db.transaction import atomic
from django.forms import model_to_dict
from django.utils import timezone
from django.utils.translation import ugettext as _

from apps.api import TransferApi
from apps.log_databus.constants import (
    RESTORE_INDEX_SET_PREFIX,
    ArchiveExpireTime,
    ArchiveInstanceType,
)
from apps.log_databus.exceptions import (
    ArchiveIndexSetInfoNotFound,
    ArchiveIndexSetStatusError,
    ArchiveNotFound,
    CollectorActiveException,
    CollectorConfigNotExistException,
    RestoreExpired,
    RestoreNotFound,
)
from apps.log_databus.models import ArchiveConfig, CollectorConfig, RestoreConfig
from apps.log_search.constants import (
    DEFAULT_TIME_FIELD,
    InnerTag,
    TimeFieldTypeEnum,
    TimeFieldUnitEnum,
)
from apps.log_search.handlers.index_set import IndexSetHandler
from apps.log_search.models import LogIndexSetData, Scenario
from apps.utils.db import array_group, array_hash
from apps.utils.function import ignored
from apps.utils.local import get_local_param, get_request_username
from apps.utils.thread import MultiExecuteFunc
from apps.utils.time_handler import (
    format_user_time_zone,
    format_user_time_zone_humanize,
)
from bkm_space.utils import bk_biz_id_to_space_uid

class ArchiveHandler:
    def __init__(self, archive_config_id=None):
        self.archive = None
        if archive_config_id is not None:
            try:
                self.archive: ArchiveConfig = ArchiveConfig.objects.get(archive_config_id=archive_config_id)
            except ArchiveConfig.DoesNotExist:
                raise ArchiveNotFound

    @classmethod
    def to_user_time_format(cls, time):
        return format_user_time_zone(time, get_local_param("time_zone"))

    @classmethod
    def list(cls, archives):
        """
        list
        @param archives:
        @return:
        """
        archive_group = array_group(archives, "archive_config_id", True)
        archive_config_ids = list(archive_group.keys())
        archive_objs_all = ArchiveConfig.objects.filter(archive_config_id__in=archive_config_ids)
        table_ids = list()
        index_set_ids = list()
        for _obj in archive_objs_all:
            if _obj.instance_type != ArchiveInstanceType.INDEX_SET.value:
                table_ids.append(_obj.table_id)
            else:
                index_set_ids.append(_obj.instance_id)
        log_index_set_data_objs = LogIndexSetData.objects.filter(index_set_id__in=index_set_ids)
        index_set_table_ids = [_obj.result_table_id for _obj in log_index_set_data_objs]

        table_ids.extend(index_set_table_ids)
        archive_detail = array_group(TransferApi.list_result_table_snapshot({"table_ids": table_ids}), "table_id", True)
        index_set_table_id_mapping = dict()
        for _obj in log_index_set_data_objs:
            if _obj.index_set_id not in index_set_table_id_mapping:
                index_set_table_id_mapping[_obj.index_set_id] = list()
            index_set_table_id_mapping[_obj.index_set_id].append(_obj.result_table_id)
        for archive in archive_objs_all:
            if archive.instance_type == ArchiveInstanceType.INDEX_SET.value:
                archive_group[archive.archive_config_id]["instance_name"] = archive.instance_name
                archive_group[archive.archive_config_id]["_index_set_id"] = archive.instance_id
                table_ids = index_set_table_id_mapping.get(archive.instance_id, [])
                if not table_ids:
                    for field in ["doc_count", "store_size", "index_count"]:
                        archive_group[archive.archive_config_id][field] = 0
                for table_id in table_ids:
                    for field in ["doc_count", "store_size", "index_count"]:
                        if field not in archive_group.get(archive.archive_config_id, {}):
                            archive_group[archive.archive_config_id][field] = 0
                        _num = int(archive_detail.get(table_id, {}).get(field, 0))
                        archive_group[archive.archive_config_id][field] += _num
            else:
                archive_group[archive.archive_config_id]["instance_name"] = archive.instance_name
                archive_group[archive.archive_config_id]["_collector_config_id"] = archive.collector_config_id
                for field in ["doc_count", "store_size", "index_count"]:
                    _num = int(archive_detail.get(archive.table_id, {}).get(field, 0))
                    archive_group[archive.archive_config_id][field] = _num
        return archives

    def retrieve(self, page, pagesize):
        """
        retrieve
        @param page:
        @param pagesize:
        @return:
        """
        if self.archive.instance_type == ArchiveInstanceType.INDEX_SET.value:
            table_ids = list(
                LogIndexSetData.objects.filter(index_set_id=self.archive.instance_id).values_list(
                    "result_table_id", flat=True
                )
            )
        else:
            table_ids = [self.archive.table_id]
        snapshot_info = TransferApi.list_result_table_snapshot_indices({"table_ids": table_ids})
        archive = model_to_dict(self.archive)
        indices = []
        for snapshot in snapshot_info:
            for _snapshot in snapshot:
                expired_time = (
                    ArchiveExpireTime.PERMANENT
                    if self.archive.snapshot_days == 0
                    else format_user_time_zone_humanize(_snapshot.get("expired_time"), get_local_param("time_zone"))
                )
                indices.extend(
                    [
                        {
                            **indice,
                            "start_time": self.to_user_time_format(indice.get("start_time")),
                            "end_time": self.to_user_time_format(indice.get("end_time")),
                            "expired_time": expired_time,
                            "state": _snapshot.get("state"),
                        }
                        for indice in _snapshot.get("indices", [])
                    ]
                )
        archive["indices"] = indices[page * pagesize : (page + 1) * pagesize]
        return archive

    @atomic
    def create_or_update(self, params):
        """
        create_or_update
        @param params:
        @return:
        """
        if self.archive:
            self.archive.snapshot_days = params.get("snapshot_days")
            self.archive.save()
            if self.archive.instance_type == ArchiveInstanceType.INDEX_SET.value:
                # 索引集归档需要查询当前索引集所关联的所有table_id
                index_set_data_objs = LogIndexSetData.objects.filter(index_set_id=self.archive.instance_id)
                if not index_set_data_objs:
                    raise ArchiveIndexSetInfoNotFound
                table_ids = list()
                for _obj in index_set_data_objs:
                    if _obj.apply_status != LogIndexSetData.Status.NORMAL:
                        raise ArchiveIndexSetStatusError
                    table_ids.append(_obj.result_table_id)

                multi_execute_func = MultiExecuteFunc()
                for table_id in table_ids:
                    multi_params = {"table_id": table_id, "snapshot_days": self.archive.snapshot_days}
                    multi_execute_func.append(
                        result_key=f"modify_result_table_snapshot_{table_id}",
                        func=TransferApi.modify_result_table_snapshot,
                        params=multi_params,
                    )

                multi_execute_func.run()

            else:
                meta_update_params = {"table_id": self.archive.table_id, "snapshot_days": self.archive.snapshot_days}
                TransferApi.modify_result_table_snapshot(meta_update_params)
            return model_to_dict(self.archive)
        # 只有采集项类型需要确认结果表状态
        if params["instance_type"] == ArchiveInstanceType.COLLECTOR_CONFIG.value:
            try:
                collector: CollectorConfig = CollectorConfig.objects.get(collector_config_id=params["instance_id"])
            except CollectorConfig.DoesNotExist:
                raise CollectorConfigNotExistException
            if not collector.is_active:
                raise CollectorActiveException

        if params["instance_type"] == ArchiveInstanceType.INDEX_SET.value:
            index_set_data_objs = LogIndexSetData.objects.filter(index_set_id=int(params["instance_id"]))
            if not index_set_data_objs:
                raise ArchiveIndexSetInfoNotFound
            table_ids = list()
            for _obj in index_set_data_objs:
                if _obj.apply_status != LogIndexSetData.Status.NORMAL:
                    raise ArchiveIndexSetStatusError
                table_ids.append(_obj.result_table_id)
            create_obj = ArchiveConfig.objects.create(**params)

            multi_execute_func = MultiExecuteFunc()
            for table_id in table_ids:
                multi_params = {
                    "table_id": table_id,
                    "target_snapshot_repository_name": create_obj.target_snapshot_repository_name,
                    "snapshot_days": create_obj.snapshot_days,
                }
                multi_execute_func.append(
                    result_key=f"create_result_table_snapshot_{table_id}",
                    func=TransferApi.create_result_table_snapshot,
                    params=multi_params,
                )

            multi_execute_func.run()
        else:
            create_obj = ArchiveConfig.objects.create(**params)
            meta_create_params = {
                "table_id": create_obj.table_id,
                "target_snapshot_repository_name": create_obj.target_snapshot_repository_name,
                "snapshot_days": create_obj.snapshot_days,
            }
            TransferApi.create_result_table_snapshot(meta_create_params)
        return model_to_dict(create_obj)

    @atomic
    def delete(self):
        self.archive.delete()
        if self.archive.instance_type != ArchiveInstanceType.INDEX_SET.value:
            TransferApi.delete_result_table_snapshot({"table_id": self.archive.table_id})
        else:
            table_ids = list(
                LogIndexSetData.objects.filter(index_set_id=self.archive.instance_id).values_list(
                    "result_table_id", flat=True
                )
            )
            multi_execute_func = MultiExecuteFunc()
            for table_id in table_ids:
                params = {"table_id": table_id}
                multi_execute_func.append(
                    result_key=f"delete_result_table_snapshot_{table_id}",
                    func=TransferApi.delete_result_table_snapshot,
                    params=params,
                )

            multi_execute_func.run()

    @atomic
    def restore(self, bk_biz_id, index_set_name, start_time, end_time, expired_time, notice_user):
        """
        restore
        @param bk_biz_id:
        @param index_set_name:
        @param start_time:
        @param end_time:
        @param expired_time:
        @param notice_user:
        @return:
        """
        index_set = self._create_index_set(index_set_name)
        result_errors = list()
        if self.archive.instance_type == ArchiveInstanceType.INDEX_SET.value:
            table_ids = list(
                LogIndexSetData.objects.filter(index_set_id=self.archive.instance_id).values_list(
                    "result_table_id", flat=True
                )
            )
            multi_execute_func = MultiExecuteFunc()
            for table_id in table_ids:
                params = {
                    "table_id": table_id,
                    "start_time": start_time,
                    "end_time": end_time,
                    "expired_time": expired_time,
                }
                multi_execute_func.append(
                    result_key=f"restore_{table_id}",
                    func=TransferApi.restore_result_table_snapshot,
                    params=params,
                )

            multi_result = multi_execute_func.run(return_exception=True)
            # 构建批量创建参数列表
            bulk_create_params = list()
            for table_id in table_ids:
                meta_restore_result = multi_result.get(f"restore_{table_id}", {})
                if isinstance(meta_restore_result, Exception):
                    result_errors.append({"table_id": table_id, "reason": str(meta_restore_result)})
                    continue
                params = {
                    "bk_biz_id": bk_biz_id,
                    "archive_config_id": self.archive.archive_config_id,
                    "start_time": start_time,
                    "end_time": end_time,
                    "expired_time": expired_time,
                    "index_set_name": index_set_name,
                    "notice_user": ",".join(notice_user),
                    "index_set_id": index_set.index_set_id,
                    "meta_restore_id": meta_restore_result.get("restore_id"),
                    "total_store_size": meta_restore_result.get("total_store_size"),
                    "total_doc_count": meta_restore_result.get("total_doc_count"),
                    "created_at": timezone.now(),
                    "created_by": get_request_username(),
                }
                bulk_create_params.append(RestoreConfig(**params))
            if bulk_create_params:
                RestoreConfig.objects.bulk_create(bulk_create_params)
                objs = RestoreConfig.objects.filter(
                    archive_config_id=self.archive.archive_config_id, index_set_name=index_set_name
                )
                restore_config_info = [model_to_dict(obj) for obj in objs]
            else:
                restore_config_info = list()
        else:
            create_restore_config = RestoreConfig.objects.create(
                **{
                    "bk_biz_id": bk_biz_id,
                    "archive_config_id": self.archive.archive_config_id,
                    "start_time": start_time,
                    "end_time": end_time,
                    "expired_time": expired_time,
                    "index_set_name": index_set_name,
                    "notice_user": ",".join(notice_user),
                }
            )

            meta_params = {
                "table_id": self.archive.table_id,
                "start_time": start_time,
                "end_time": end_time,
                "expired_time": expired_time,
            }
            meta_restore_result = TransferApi.restore_result_table_snapshot(meta_params)

            create_restore_config.index_set_id = index_set.index_set_id
            create_restore_config.meta_restore_id = meta_restore_result["restore_id"]
            create_restore_config.total_store_size = meta_restore_result["total_store_size"]
            create_restore_config.total_doc_count = meta_restore_result["total_doc_count"]
            create_restore_config.save()
            restore_config_info = [model_to_dict(create_restore_config)]

        res = {
            "index_set_id": index_set.index_set_id,
            "index_set_name": index_set.index_set_name,
            "restore_config_info": restore_config_info,
            "errors": result_errors,
        }

        return res

    def _create_index_set(self, index_set_name):
        index_set_name = _("[回溯]") + index_set_name
        if self.archive.instance_type == ArchiveInstanceType.INDEX_SET.value:
            table_ids = list(
                LogIndexSetData.objects.filter(index_set_id=self.archive.instance_id).values_list(
                    "result_table_id", flat=True
                )
            )
            indexes = [
                {
                    "bk_biz_id": self.archive.bk_biz_id,
                    "result_table_id": f"{RESTORE_INDEX_SET_PREFIX}*{table_id.replace('.', '_')}_*",
                    "result_table_name": self.archive.instance_name,
                    "time_field": DEFAULT_TIME_FIELD,
                }
                for table_id in table_ids
            ]
            # 索引集下所包含的物理索引一定存在一个集群 所以取第一个result_table_id去获取集群信息
            cluster_infos = TransferApi.get_result_table_storage(
                {"result_table_list": table_ids[0], "storage_type": "elasticsearch"}
            )
            cluster_info = cluster_infos.get(table_ids[0])
        else:
            indexes = [
                {
                    "bk_biz_id": self.archive.bk_biz_id,
                    "result_table_id": f"{RESTORE_INDEX_SET_PREFIX}*{self.archive.table_id.replace('.', '_')}_*",
                    "result_table_name": self.archive.instance_name,
                    "time_field": DEFAULT_TIME_FIELD,
                }
            ]

            cluster_infos = TransferApi.get_result_table_storage(
                {"result_table_list": self.archive.table_id, "storage_type": "elasticsearch"}
            )
            cluster_info = cluster_infos.get(self.archive.table_id)
        storage_cluster_id = cluster_info["cluster_config"]["cluster_id"]
        index_set = IndexSetHandler.create(
            index_set_name=index_set_name,
            space_uid=bk_biz_id_to_space_uid(self.archive.bk_biz_id),
            storage_cluster_id=storage_cluster_id,
            scenario_id=Scenario.ES,
            view_roles=[],
            indexes=indexes,
            category_id=self.archive.instance.category_id,
            collector_config_id=self.archive.collector_config_id,
            time_field=DEFAULT_TIME_FIELD,
            time_field_type=TimeFieldTypeEnum.DATE.value,
            time_field_unit=TimeFieldUnitEnum.MILLISECOND.value,
        )
        index_set.set_tag(index_set.index_set_id, InnerTag.RESTORING.value)
        return index_set

    @classmethod
    def list_archive(cls, bk_biz_id):
        """
        list_archive
        @param bk_biz_id:
        @return:
        """
        return [
            {
                "archive_config_id": archive.archive_config_id,
                "instance_name": archive.instance_name,
                "instance_id": archive.instance_id,
                "_collector_config_id": archive.collector_config_id,
            }
            for archive in ArchiveConfig.objects.filter(bk_biz_id=bk_biz_id)
        ]

    @classmethod
    def list_restore(cls, restore_list):
        """
        list_restore
        @param restore_list:
        @return:
        """
        ret = []
        instances = restore_list.serializer.instance
        for instance in instances:
            # archive config maybe delete so not show restore
            with ignored(ArchiveConfig.DoesNotExist):
                ret.append(
                    {
                        "restore_config_id": instance.restore_config_id,
                        "index_set_name": instance.index_set_name,
                        "index_set_id": instance.index_set_id,
                        "start_time": cls.to_user_time_format(instance.start_time),
                        "end_time": cls.to_user_time_format(instance.end_time),
                        "expired_time": cls.to_user_time_format(instance.expired_time),
                        "total_store_size": instance.total_store_size,
                        "instance_id": instance.archive.instance_id,
                        "instance_name": instance.archive.instance_name,
                        "instance_type": instance.archive.instance_type,
                        "_collector_config_id": instance.archive.collector_config_id,
                        "archive_config_id": instance.archive.archive_config_id,
                        "notice_user": instance.notice_user.split(","),
                        "is_expired": instance.is_expired(),
                    }
                )
        return ret

    @classmethod
    @atomic
    def update_restore(cls, restore_config_id, expired_time):
        """
        update_restore
        @param restore_config_id:
        @param expired_time:
        @return:
        """
        try:
            restore: RestoreConfig = RestoreConfig.objects.get(restore_config_id=restore_config_id)
        except RestoreConfig.DoesNotExist:
            raise RestoreNotFound
        if restore.is_expired():
            raise RestoreExpired

        # 归档任务是索引集的场景下 会存在多个归档回溯
        result_errors = list()
        if restore.archive.instance_type == ArchiveInstanceType.INDEX_SET.value:
            restore_objs = RestoreConfig.objects.filter(index_set_id=restore.index_set_id)
            restore_objs.update(expired_time=expired_time)
            meta_restore_ids = list(restore_objs.values_list("meta_restore_ids", flat=True))
            multi_execute_func = MultiExecuteFunc()
            for meta_restore_id in meta_restore_ids:
                params = {"restore_id": meta_restore_id, "expired_time": expired_time}
                multi_execute_func.append(
                    result_key=f"modify_restore_result_table_snapshot_{meta_restore_id}",
                    func=TransferApi.modify_restore_result_table_snapshot,
                    params=params,
                )
            multi_execute_result = multi_execute_func.run(return_exception=True)
            for meta_restore_id in meta_restore_ids:
                multi_result = multi_execute_result.get(f"modify_restore_result_table_snapshot_{meta_restore_id}", {})
                if isinstance(multi_result, Exception):
                    result_errors.append({"meta_restore_id": meta_restore_id, "reason": str(multi_result)})

            restore_config_info = [model_to_dict(obj) for obj in restore_objs]
        else:
            restore.expired_time = expired_time
            restore.save()
            TransferApi.modify_restore_result_table_snapshot(
                {"restore_id": restore.meta_restore_id, "expired_time": expired_time}
            )

            restore_config_info = [model_to_dict(restore)]

        return {"restore_config_info": restore_config_info, "errors": result_errors}

    @classmethod
    @atomic
    def delete_restore(cls, restore_config_id):
        """
        delete_restore
        @param restore_config_id:
        @return:
        """
        try:
            restore: RestoreConfig = RestoreConfig.objects.get(restore_config_id=restore_config_id)
        except RestoreConfig.DoesNotExist:
            raise RestoreNotFound
        index_set_handler = IndexSetHandler(restore.index_set_id)
        index_set_handler.stop()
        # 归档任务是索引集的场景下 会存在多个归档回溯
        if restore.archive.instance_type == ArchiveInstanceType.INDEX_SET.value:
            restore_objs = RestoreConfig.objects.filter(index_set_id=restore.index_set_id)
            meta_restore_ids = list(restore_objs.values_list("meta_restore_id", flat=True))
            restore_objs.delete()
            multi_execute_func = MultiExecuteFunc()
            for meta_restore_id in meta_restore_ids:
                params = {"restore_id": meta_restore_id}
                multi_execute_func.append(
                    result_key=f"delete_restore_result_table_snapshot_{meta_restore_id}",
                    func=TransferApi.delete_restore_result_table_snapshot,
                    params=params,
                )
            multi_execute_func.run()
        else:
            restore.delete()
            TransferApi.delete_restore_result_table_snapshot({"restore_id": restore.meta_restore_id})

    def archive_state(self):
        if self.archive.instance_type == ArchiveInstanceType.INDEX_SET.value:
            table_ids = list(
                LogIndexSetData.objects.filter(index_set_id=self.archive.instance_id).values_list(
                    "result_table_id", flat=True
                )
            )
        else:
            table_ids = [self.archive.table_id]
        return TransferApi.get_result_table_snapshot_state({"table_ids": table_ids})

    @staticmethod
    def batch_get_restore_state(restore_config_ids: list):
        """
        batch_get_restore_state
        @param restore_config_ids:
        @return:
        """
        restores = RestoreConfig.objects.filter(restore_config_id__in=restore_config_ids).values(
            "meta_restore_id", "restore_config_id"
        )
        meta_restore_ids = [restore["meta_restore_id"] for restore in restores]
        restore_hash = array_hash(restores, "meta_restore_id", "restore_config_id")
        meta_restore_states = TransferApi.get_restore_result_table_snapshot_state({"restore_ids": meta_restore_ids})
        for meta_restore_state in meta_restore_states:
            meta_restore_state["restore_config_id"] = restore_hash[meta_restore_state["restore_id"]]
        return meta_restore_states


__author__    = "Individual contributors (see AUTHORS file)"
__date__      = "$DATE$"
__rev__       = "$REV$"
__license__   = "AGPL v.3"
__copyright__ = """
This file is part of the ESP Web Site
Copyright (c) 2007 by the individual contributors
  (see AUTHORS file)

The ESP Web Site is free software; you can redistribute it and/or
modify it under the terms of the GNU Affero General Public License
as published by the Free Software Foundation; either version 3
of the License, or (at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU Affero General Public License for more details.

You should have received a copy of the GNU Affero General Public
License along with this program; if not, write to the Free Software
Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.

Contact information:
MIT Educational Studies Program
  84 Massachusetts Ave W20-467, Cambridge, MA 02139
  Phone: 617-253-4882
  Email: esp-webmasters@mit.edu
Learning Unlimited, Inc.
  527 Franklin St, Cambridge, MA 02139
  Phone: 617-379-0178
  Email: web-team@learningu.org
"""

import json
from datetime import datetime, timedelta

from django.db.models import Min
from django.db.models.query import Q
from django.http import HttpResponse
from django.shortcuts import get_object_or_404
from django.utils.safestring import mark_safe

from esp.users.models    import ESPUser, Record, ContactInfo, StudentInfo, K12School
from esp.program.models import RegistrationProfile
from esp.program.class_status import ClassStatus

from esp.program.modules.base import ProgramModuleObj, needs_onsite, needs_student, main_call, aux_call
from esp.program.models import ClassSubject, ClassSection, StudentRegistration, ScheduleMap, Program
from esp.utils.web import render_to_response
from esp.cal.models import Event
from argcache import cache_function
from esp.users.models import ESPUser, Record
from esp.resources.models import ResourceAssignment
from esp.utils.models import Printer, PrintRequest
from esp.utils.query_utils import nest_Q
from esp.tagdict.models import Tag
from esp.accounting.controllers import IndividualAccountingController
from esp.program.controllers.studentclassregmodule import RegistrationTypeController as RTC

class OnSiteClassList(ProgramModuleObj):
    doc = """Display lists of classes for onsite registration purposes."""

    @classmethod
    def module_properties(cls):
        return {
            "admin_title": "Show Open Classes at Onsite Registration",
            "link_title": "List of Open Classes",
            "module_type": "onsite",
            "seq": 32,
            "choosable": 1,
            }

    @cache_function
    def section_data(sec):
        sect = {}
        sect['id'] = sec.id
        sect['emailcode'] = sec.emailcode()
        sect['title'] = sec.title()
        sect['teachers'] = ', '.join([t.name() for t in list(sec.teachers)])
        sect['rooms'] = (' ,'.join(sec.prettyrooms()))[:12]
        return sect
    section_data.depend_on_model('resources.ResourceAssignment')
    section_data.depend_on_cache(ClassSubject.get_teachers, lambda **kwargs: {})
    section_data=staticmethod(section_data)

    """ Warning: for performance reasons, these views are not abstracted away from
        the models.  If the schema is changed this code will need to be updated.
    """

    @aux_call
    @needs_onsite
    def catalog_status(self, request, tl, one, two, module, extra, prog):
        resp = HttpResponse(content_type='application/json')
        #   Fetch a reduced version of the catalog to save time
        sections = list(ClassSection.objects.filter(parent_class__parent_program=prog, status__gt=0).extra({'event_ids':  """ARRAY(SELECT "cal_event"."id" FROM "cal_event", "program_classsection_meeting_times" WHERE ("program_classsection_meeting_times"."event_id" = "cal_event"."id" AND "program_classsection_meeting_times"."classsection_id" = "program_classsection"."id"))"""}).values('id', 'parent_class__id', 'enrolled_students', 'event_ids', 'registration_status'))
        for i in range(0, len(sections)):
            sections[i]['capacity'] = ClassSection.objects.get(id=sections[i]['id']).capacity
        data = {
            #   Todo: section current capacity ? (see ClassSection.get_capacity())
            'classes': list(ClassSubject.objects.filter(parent_program=prog, status__gt=0).extra({'teacher_names': """array_to_string(ARRAY(SELECT auth_user.first_name || ' ' || auth_user.last_name FROM auth_user,program_class_teachers WHERE program_class_teachers.classsubject_id=program_class.id AND auth_user.id=program_class_teachers.espuser_id), ', ')""", 'class_size_max_optimal': """SELECT program_classsizerange.range_max FROM program_classsizerange WHERE program_classsizerange.id = optimal_class_size_range_id"""}).values('id', 'class_size_max', 'class_size_max_optimal', 'class_info', 'prereqs', 'hardness_rating', 'grade_min', 'grade_max', 'title', 'teacher_names', 'category__symbol', 'category__id')),
            'sections': sections,
            'timeslots': list(prog.getTimeSlots().extra({'start_millis':"""EXTRACT(EPOCH FROM start) * 1000""",'label': """to_char("start", 'Dy HH:MI -- ') || to_char("end", 'HH:MI AM')"""}).values_list('id', 'label','start_millis').order_by("start")),
            'categories': list(prog.class_categories.all().order_by('-symbol').values('id', 'symbol', 'category')),
        }
        json.dump(data, resp)

        return resp

    @aux_call
    @needs_onsite
    def enrollment_status(self, request, tl, one, two, module, extra, prog):
        resp = HttpResponse(content_type='application/json')
        data = StudentRegistration.valid_objects().filter(section__status__gt=0, section__parent_class__status__gt=0, section__parent_class__parent_program=prog, relationship__name='Enrolled').values_list('user__id', 'section__id')
        json.dump(list(data), resp)
        return resp

    @aux_call
    @needs_onsite
    def students_status(self, request, tl, one, two, module, extra, prog):
        resp = HttpResponse(content_type='application/json')
        #   Try to ensure we don't miss anyone

        students_dict = self.program.students(QObjects=True)
        search_query = request.GET.get('q')

        student_types = ['student_profile']     #   You could add more list names here, but it would get very slow.
        students_Q = Q()

        for student_type in student_types:
            students_Q = students_Q | students_dict[student_type]

        students = ESPUser.objects.filter(students_Q)
        program_students_ids = []

        if search_query:
            #If user provided a search term then we want to expand search to the
            #entire student base

            search_tokens = search_query.split(' ')
            search_qset = Q()

            for token in search_tokens:
                search_qset = search_qset & (Q(last_name__icontains=token) | Q(first_name__icontains=token))

            program_students_ids = set(students.values_list('id', flat=True).distinct())
            students = ESPUser.objects.filter(search_qset)

        students = students.values_list('id', 'last_name', 'first_name') \
                           .distinct() \
                           .order_by('first_name', 'last_name')

        data = []
        for student in students:
            has_profile = not search_query or student[0] in program_students_ids
            data.append(list(student) + [has_profile])

        data.sort(key=lambda x: not x[3])
        if search_query:
            data = data[:20]

        json.dump(data, resp)
        return resp

    @aux_call
    @needs_onsite
    def register_student(self, request, tl, one, two, module, extra, prog):
        resp = HttpResponse(content_type='application/json')
        program = self.program
        success = False
        student = get_object_or_404(ESPUser,pk=request.POST.get("student_id"))

        registration_profile = RegistrationProfile.getLastForProgram(student,
                                                                program)
        success = registration_profile.student_info is not None

        if success:
            registration_profile.save()

            for extension in ['attended','med','liab','onsite']:
                Record.createBit(extension, program, student)

            IndividualAccountingController.updatePaid(self.program, student, paid=True)

        json.dump({'status':success}, resp)
        return resp

    @aux_call
    @needs_onsite
    def checkin_status(self, request, tl, one, two, module, extra, prog):
        resp = HttpResponse(content_type='application/json')
        students = prog.students(True)
        data = ESPUser.objects.filter(students['attended'] & ~students['checked_out']).distinct().values_list('id')
        json.dump(list(data), resp)
        return resp

    @aux_call
    @needs_onsite
    def counts_status(self, request, tl, one, two, module, extra, prog):
        resp = HttpResponse(content_type='application/json')
        data = ClassSection.objects.filter(status__gt=0, parent_class__status__gt=0, parent_class__parent_program=prog).values_list('id', 'enrolled_students', 'attending_students')
        json.dump(list(data), resp)
        return resp

    @aux_call
    @needs_onsite
    def full_status(self, request, tl, one, two, module, extra, prog):
        resp = HttpResponse(content_type='application/json')
        data = [[section.id, section.isFull(webapp=True)] for section in
                     ClassSection.objects.filter(status__gt=0, parent_class__status__gt=0,
                                                 parent_class__parent_program=prog)]
        json.dump(data, resp)
        return resp

    @aux_call
    @needs_onsite
    def rooms_status(self, request, tl, one, two, module, extra, prog):
        resp = HttpResponse(content_type='application/json')
        data = ClassSection.objects.filter(status__gt=0, parent_class__status__gt=0, parent_class__parent_program=prog, resourceassignment__resource__res_type__name="Classroom").select_related('resourceassignment__resource__name').values_list('id', 'resourceassignment__resource__name', 'resourceassignment__resource__num_students')
        json.dump(list(data), resp)
        return resp

    @aux_call
    @needs_onsite
    def get_schedule_json(self, request, tl, one, two, module, extra, prog):
        resp = HttpResponse(content_type='application/json')
        result = {'user': None, 'user_grade': 0, 'sections': [], 'messages': []}
        try:
            result['user'] = int(request.GET['user'])
        except:
            result['messages'].append('Error: no user specified.')
        if result['user']:
            result['user_grade'] = ESPUser.objects.get(id=result['user']).getGrade(program=prog)
            result['sections'] = list(ClassSection.objects.filter(nest_Q(StudentRegistration.is_valid_qobject(), 'studentregistration'), status__gt=0, parent_class__status__gt=0, parent_class__parent_program=prog, studentregistration__relationship__name='Enrolled', studentregistration__user__id=result['user']).values_list('id', flat=True).distinct())
        json.dump(result, resp)
        return resp

    @aux_call
    @needs_onsite
    def update_schedule_json(self, request, tl, one, two, module, extra, prog):
        resp = HttpResponse(content_type='application/json')
        result = {'user': None, 'sections': [], 'messages': []}
        try:
            user = ESPUser.objects.get(id=int(request.GET['user']))
        except:
            user = None
            result['messages'].append('Error: could find user %s' % request.GET.get('user', None))
        try:
            desired_sections = json.loads(request.GET['sections'])
        except:
            result['messages'].append('Error: could not parse requested sections %s' % request.GET.get('sections', None))
            desired_sections = None

        #   Check in student if not currently checked in, since if they're using this view they must be onsite
        if not prog.isCheckedIn(user) and request.GET.get('check_in') == 'true':
            rec = Record(user=user, program=prog, event='attended')
            rec.save()

        if user and desired_sections is not None:
            override_full = (request.GET.get("override", "") == "true")

            current_sections = list(ClassSection.objects.filter(nest_Q(StudentRegistration.is_valid_qobject(), 'studentregistration'), status__gt=0, parent_class__status__gt=0, parent_class__parent_program=prog, studentregistration__relationship__name='Enrolled', studentregistration__user__id=user.id).values_list('id', flat=True).order_by('id').distinct())
            sections_to_remove = ClassSection.objects.filter(id__in=list(set(current_sections) - set(desired_sections)))
            sections_to_add = ClassSection.objects.filter(id__in=list(set(desired_sections) - set(current_sections)))

            failed_add_sections = []
            for sec in sections_to_add:
                if sec.isFull(webapp=True) and not override_full:
                    result['messages'].append('Failed to add %s (%s) to %s: %s (%s).  Error was: %s' % (user.name(), user.id, sec.emailcode(), sec.title(), sec.id, 'Class is currently full.'))
                    failed_add_sections.append(sec.id)

            if len(failed_add_sections) == 0:
                verbs = RTC.getVisibleRegistrationTypeNames(prog)
                #   Remove sections the student wants out of
                for sec in sections_to_remove:
                    sec.unpreregister_student(user, verbs)
                    result['messages'].append('Removed %s (%s) from %s: %s (%s)' % (user.name(), user.id, sec.emailcode(), sec.title(), sec.id))

                #   Remove sections that conflict with those the student wants into
                sec_times = sections_to_add.select_related('meeting_times__id').values_list('id', 'meeting_times__id').order_by('meeting_times__id').distinct()
                sm = ScheduleMap(user, prog)
                existing_sections = []
                for (sec, ts) in sec_times:
                    if ts and ts in sm.map and len(sm.map[ts]) > 0:
                        #   We found something we need to remove
                        for sm_sec in sm.map[ts]:
                            if sm_sec.id not in sections_to_add:
                                sm_sec.unpreregister_student(user, verbs)
                                result['messages'].append('Removed %s (%s) from %s: %s (%s)' % (user.name(), user.id, sm_sec.emailcode(), sm_sec.title(), sm_sec.id))
                            else:
                                existing_sections.append(sm_sec)

                #   Add the sections the student wants
                for sec in sections_to_add:
                    if sec not in existing_sections and sec.id not in failed_add_sections:
                        error = sec.cannotAdd(user, not override_full, webapp=True)
                        if not error:
                            reg_result = sec.preregister_student(user, overridefull=override_full, webapp=True)
                            if not reg_result:
                                error = 'Class is currently full.'
                        else:
                            reg_result = False
                        if reg_result:
                            result['messages'].append('Added %s (%s) to %s: %s (%s)' % (user.name(), user.id, sec.emailcode(), sec.title(), sec.id))
                        else:
                            result['messages'].append('Failed to add %s (%s) to %s: %s (%s).  Error was: %s' % (user.name(), user.id, sec.emailcode(), sec.title(), sec.id, error))

            result['user'] = user.id
            result['sections'] = list(ClassSection.objects.filter(nest_Q(StudentRegistration.is_valid_qobject(), 'studentregistration'), status__gt=0, parent_class__status__gt=0, parent_class__parent_program=prog, studentregistration__relationship__name='Enrolled', studentregistration__user__id=result['user']).values_list('id', flat=True).distinct())

        json.dump(result, resp)
        return resp

    """ End of highly model-dependent JSON views    """

    @aux_call
    @needs_onsite
    def printschedule_status(self, request, tl, one, two, module, extra, prog):
        resp = HttpResponse(content_type='application/json')

        result = {}

        try:
            user = int(request.GET.get('user', None))
            user_obj = ESPUser.objects.get(id=user)
        except:
            result['message'] = "Could not find user %s." % request.GET.get('user', None)

        printer = request.GET.get('printer',None)
        if printer is not None:
            # we could check that it exists and is unique first, but if not, that should be an error anyway, and it isn't the user's fault unless they're trying to mess with us, so a 500 is reasonable and gives us better debugging output.
            printer = Printer.objects.get(name=printer)
        req = PrintRequest.objects.create(user=user_obj, printer=printer)
        result['message'] = "Submitted %s's schedule for printing (print request #%s)." % (user_obj.name(), req.id)

        json.dump(result, resp)
        return resp

    @aux_call
    @needs_onsite
    def classchange_grid(self, request, tl, one, two, module, extra, prog):
        context = {}
        context['timeslots'] = prog.getTimeSlots()
        context['printers'] = Printer.objects.all().values_list('name', flat=True)
        context['initial_student'] = request.GET.get('student_id', '')
        context['check_in_default'] = datetime.today().date() in prog.dates()

        open_class_category = prog.open_class_category
        open_class_category = dict( [ (k, getattr( open_class_category, k )) for k in ['id','symbol','category'] ] )
        context['open_class_category'] = mark_safe(json.dumps(open_class_category))

        return render_to_response(self.baseDir()+'ajax_status.html', request, context)

    @aux_call
    @needs_onsite
    def classList(self, request, tl, one, two, module, extra, prog):
        return self.classList_base(request, tl, one, two, module, extra, prog, template_name='classlist.html')

    @aux_call
    @needs_student
    def classlist_public(self, request, tl, one, two, module, extra, prog):
        return self.classList_base(request, tl, one, two, module, extra, prog, options={}, template_name='allclass_fragment.html')

    def classList_base(self, request, tl, one, two, module, extra, prog, options=None, template_name='classlist.html'):
        """ Display a list of all classes that still have space in them """

        #   Allow options to be supplied as an argument to the view function, in lieu
        #   of request.GET (used by the public view just above)
        if options is None:
            options = request.GET.copy()

            #   Display options-selection page if this page is requested with no GET variables
            if len(options.keys()) == 0:
                return render_to_response(self.baseDir() + 'classlist_options.html', request, {'prog': prog})

        context = {}
        defaults = {'refresh': 120, 'scrollspeed': 1}
        for key_option in defaults.keys():
            if key_option in options:
                context[key_option] = options[key_option]
            else:
                context[key_option] = defaults[key_option]

        time_now = datetime.now()

        start_id = int(options.get('start', -1))
        if start_id != -1:
            curtime = Event.objects.filter(id=start_id)
        else:
            window_start = time_now + timedelta(-1, 85200)  # 20 minutes ago
            curtime = Event.objects.filter(start__gte=window_start, event_type__description='Class Time Block').order_by('start')
            # If there are no events after the current time, just pick the first event of the program
            if curtime.count() == 0:
                curtime = Event.objects.filter(program=self.program, event_type__description='Class Time Block').order_by('start')

        end_id = int(options.get('end', -1))
        if end_id != -1:
            endtime = Event.objects.filter(id=end_id)
        else:
            endtime = None

        sort_spec = options.get('sorting', None)
        if sort_spec is None:
            sort_spec = extra

        #   Enforce a maximum refresh speed to avoid server overload.
        min_refresh = int(Tag.getTag('onsite_classlist_min_refresh'))
        if int(context['refresh']) < min_refresh:
            context['refresh'] = min_refresh

        classes = []
        if curtime:
            curtime = curtime[0]
            if endtime:
                endtime = endtime[0]
                classes = self.program.sections().annotate(begin_time=Min("meeting_times__start")).filter(
                    status=ClassStatus.ACCEPTED, parent_class__status=ClassStatus.ACCEPTED,
                    begin_time__gte=curtime.start, begin_time__lte=endtime.start
                    )
            else:
                 classes = self.program.sections().annotate(begin_time=Min("meeting_times__start")).filter(
                     status=ClassStatus.ACCEPTED, parent_class__status=ClassStatus.ACCEPTED,
                     begin_time__gte=curtime.start
                     )
            if sort_spec == 'unsorted':
                classes = classes.order_by('begin_time', 'id').distinct()
            elif sort_spec == 'by_time':
                classes = classes.order_by('begin_time', 'parent_class__category', 'id').distinct()
            else:
                classes = classes.order_by('parent_class__category', 'begin_time', 'id').distinct()

        context.update({'prog': prog, 'current_time': curtime, 'classes': classes, 'one': one, 'two': two})

        if sort_spec == 'unsorted':
            context['use_categories'] = False
        else:
            context['use_categories'] = True

        return render_to_response(self.baseDir()+template_name, request, context)

    @main_call
    @needs_onsite
    def allClassList(self, request, tl, one, two, module, extra, prog):
        """ Display a list of all classes that still have space in them """

        #   This view still uses classes, not sections.  The templates show information
        #   for each section of each class.
        classes = [(i.num_students()/(i.class_size_max + 1), i) for i in self.program.classes() if i.class_size_max]
        classes.sort()
        classes = [i[1] for i in classes]

        categories = {}
        for cls in classes:
            categories[cls.category_id] = {'id': cls.category_id, 'category': cls.category.category}

        printers = [ x.name for x in Printer.objects.all() ]

        return render_to_response(self.baseDir()+'allclasslist.html', request,
            {'classes': classes, 'prog': self.program, 'one': one, 'two': two, 'categories': categories.values(), 'printers': printers})

    def makeLink(self):
        calls = [("classchange_grid","Grid-based Class Changes Interface"), ("classList","Scrolling Class List"), (self.get_main_view(),self.module.link_title)]
        strings = [u'<a href="%s" title="%s" class="vModuleLink" >%s</a>' % \
                ('/' + self.module.module_type + '/' + self.program.url + '/' + call[0], call[1], call[1]) for call in calls]
        return "</li><li>".join(strings)

    def makeButtonLink(self):
        calls = [("classchange_grid","Grid-based Class Changes Interface"), ("classList","Scrolling Class List"), (self.get_main_view(),self.module.link_title)]
        strings = [u'<a href="%s"><button type="button" class="module_link_large btn btn-default btn-lg"><div class="module_link_main">%s</div></button></a>' % \
                ('/' + self.module.module_type + '/' + self.program.url + '/' + call[0], call[1]) for call in calls]
        return "<br>".join(strings)

    class Meta:
        proxy = True
        app_label = 'modules'


import datetime
import time
hh, mm, ss = [int(i) for i in input().split()]
print(hh, mm, ss)
now = datetime.datetime.now()
deadline = now + datetime.timedelta(hours=hh, minutes=mm, seconds=ss)
print(deadline.hour, deadline.minute, deadline.second)
while True:
    time.sleep(1)
    now = datetime.datetime.now()
    if deadline.hour <= now.hour and deadline.minute <= now.minute and deadline.second <= now.second:
        print("Time is up")
        break


from django.db import models
from account.models import *

class Category(models.Model):
    category_name = models.CharField(verbose_name='Category name', max_length=255)
    category_cover_image = models.ImageField(verbose_name='Category Cover Image', upload_to='category_cover_images/')
    category_description = models.TextField(verbose_name='Category description', blank=True, null=True)
    category_follow_users = models.ManyToManyField(CustomUser, verbose_name='Are you interested in this category?', blank=True, null=True)

    def __str__(self):
        return str(self.category_name)

    class Meta:
        verbose_name = 'Category'
        verbose_name_plural = 'Categories'
        db_table = 'category'

class Company(models.Model):

    GOAL_CHOICES = (
        ('Get more visitors', 'Get more visitors'),
        ('Get more money', 'Get more money'),
        ('Get more messages', 'Get more messages'),
        ('Get more likes', 'Get more likes')
    )

    company_name = models.CharField(verbose_name='Company Name', max_length=255)
    company_logo = models.ImageField(verbose_name='Company Logo', upload_to='company_logos/', blank=True, null=True)
    company_goal = models.CharField(verbose_name='Company Goal', max_length=255, choices=GOAL_CHOICES)
    company_description = models.TextField(verbose_name='Company Description', blank=True, null=True)
    company_email = models.EmailField(verbose_name='Company Email')
    company_phone_number = models.CharField(verbose_name='Company Phone Number', max_length=255)
    company_address = models.CharField(verbose_name='Company Address', max_length=255)
    company_address_second = models.CharField(verbose_name='Company Address Second', max_length=255, blank=True, null=True)
    company_city = models.CharField(verbose_name='Company City', max_length=255)
    company_post_code = models.PositiveIntegerField(verbose_name='Company Post Code')

    def __str__(self):
        return str(self.company_name) + ' ' + str(self.company_city)

    class Meta:
        verbose_name = 'Company'
        verbose_name_plural = 'Companies'
        db_table = 'company'

class Activity(models.Model):

    ACTIVITY_TYPE = (
        ('Personal Activity', 'Personal Activity'),
        ('Team Activity', 'Team Activity')
    )

    GENDER_PERMISSIONS = (
        ('All', 'All'),
        ('Male', 'Male'),
        ('Female', 'Female')
    )

    ACTIVITY_FORMAT = (
        ('Offline', 'Offline'),
        ('Online', 'Online')
    )

    PRICE_TYPE_ACTIVITY = (
        ('Per Hour', 'Per Hour'),
        ('Per Day', 'Per Day'),
        ('Per Week', 'Per Week'),
        ('Per Month', 'Per Month'),
        ('Per Year', 'Per Year'),
        ('For Activity', 'For Activity')
    )

    DURATION_TYPE_ACTIVITY = (
        ('Per Day', 'Per Day'),
        ('Per Week', 'Per Week'),
        ('Per Month', 'Per Month'),
        ('Per Year', 'Per Year'),
        ('Depends on person', 'Depends on person')
    )

    company_id = models.ForeignKey(Company, verbose_name='Company Activation', on_delete=models.CASCADE, related_name='activity_company')
    category_type = models.ForeignKey(Category, verbose_name='Category Activation', on_delete=models.CASCADE, related_name='activity_category_type')
    activity_name = models.CharField(verbose_name='Activity Name', max_length=255)
    activity_cover_image = models.ImageField(verbose_name='Activity Cover Image', upload_to='activity_cover_images/', blank=True, null=True)
    activity_description = models.TextField(verbose_name='Activity Description', blank=True, null=True)
    activity_type = models.CharField(verbose_name='Activity Type', max_length=255, choices=ACTIVITY_TYPE)
    activity_gender_permissions = models.CharField(verbose_name='Activity Gender_Permissions', max_length=255, choices=GENDER_PERMISSIONS)
    activity_age_permissions_lower = models.PositiveIntegerField(verbose_name='Activity Age Permissions Lower', default=9)
    activity_age_permissions_upper = models.PositiveIntegerField(verbose_name='Activity Age Permissions Upper', default=25)
    activity_format = models.CharField(verbose_name='Activity Format', max_length=255, choices=ACTIVITY_FORMAT)
    activity_price = models.PositiveIntegerField(verbose_name='Activity Price', default=0)
    activity_price_type = models.CharField(verbose_name='Activity Price Type', max_length=255, choices=PRICE_TYPE_ACTIVITY)
    activity_started_date = models.DateField(verbose_name='Activity Started Date')
    activity_duration = models.PositiveIntegerField(verbose_name='Activity Duration', blank=True, null=True)
    activity_duration_type = models.CharField(verbose_name='Activity Duration Type', max_length=255, choices=DURATION_TYPE_ACTIVITY)

    def __str__(self):
        return str(self.activity_name) + '' + str(self.activity_type) + '' + str(self.activity_format)

    class Meta:
        verbose_name = 'Activity'
        verbose_name_plural = 'Activities'
        db_table = 'activity'