from google.appengine.ext import ndb

from consts.event_type import EventType
from database.database_query import DatabaseQuery
from database.dict_converters.district_converter import DistrictConverter
from models.district import District
from models.event import Event

# These district codes are equiavlent due to FIRST renaming districts
CODE_MAP = {
    # Old to new
    'mar': 'fma',
    'nc': 'fnc',
    # New to old
    'fma': 'mar',
    'fnc': 'nc',
}
def get_equivalent_keys(district_key):
    # Returns a list of equivalent district keys
    year = int(district_key[:4])
    code = district_key[4:]

    keys = [district_key]
    if code in CODE_MAP:
        keys.append('{}{}'.format(year, CODE_MAP[code]))
    return keys


class DistrictQuery(DatabaseQuery):
    CACHE_VERSION = 1
    CACHE_KEY_FORMAT = 'district_{}'  # (district_key)
    DICT_CONVERTER = None

    @ndb.tasklet
    def _query_async(self):
        district_key = self._query_args[0]
        # Fetch all equivalent keys
        keys = get_equivalent_keys(district_key)
        districts = yield ndb.get_multi_async([ndb.Key(District, key) for key in keys])
        for district in districts:
            if district:
                # Return first key that exists
                raise ndb.Return(district)
        raise ndb.Return(None)


class DistrictChampsInYearQuery(DatabaseQuery):
    CACHE_VERSION = 0
    CACHE_KEY_FORMAT = 'district_list_{}'  # (year)
    DICT_CONVERTER = None  # Not exposed in API, not needed

    @ndb.tasklet
    def _query_async(self):
        year = self._query_args[0]
        all_cmp_event_keys = yield Event.query(
            Event.year == int(year),
            Event.event_type_enum == EventType.DISTRICT_CMP).fetch_async(keys_only=True)
        events = yield ndb.get_multi_async(all_cmp_event_keys)
        raise ndb.Return(events)


class DistrictsInYearQuery(DatabaseQuery):
    CACHE_VERSION = 0
    CACHE_KEY_FORMAT = "districts_in_year_{}"  # (year)
    DICT_CONVERTER = DistrictConverter

    @ndb.tasklet
    def _query_async(self):
        year = self._query_args[0]
        district_keys = yield District.query(District.year == year).fetch_async(keys_only=True)
        districts = yield ndb.get_multi_async(district_keys)
        raise ndb.Return(districts)


class DistrictHistoryQuery(DatabaseQuery):
    CACHE_VERSION = 0
    CACHE_KEY_FORMAT = "district_history_{}"  # (abbreviation)
    DICT_CONVERTER = DistrictConverter

    @ndb.tasklet
    def _query_async(self):
        abbreviation = self._query_args[0]
        district_keys = yield District.query(District.abbreviation == abbreviation).fetch_async(keys_only=True)
        districts = yield ndb.get_multi_async(district_keys)
        raise ndb.Return(districts)

@@ -453,16 +453,18 @@
 'nc',%0A%7D%0A
+%0A%0A
 def get_
@@ -474,33 +474,26 @@
 ivalent_
-keys(district_key
+codes(code
 ):%0A    #
@@ -535,56 +535,187 @@
 ict 
-key
+code
 s%0A    
-year = int(district_key%5B:4%5D)%0A    code =
+codes = %5Bcode%5D%0A    if code in CODE_MAP:%0A        codes.append(CODE_MAP%5Bcode%5D)%0A    return codes%0A%0A%0Adef get_equivalent_keys(district_key):%0A    # Returns a list of equivalent
  dis
@@ -723,30 +723,25 @@
 rict
-_key%5B4:%5D%0A
+ keys
 %0A    
-keys
+year
  = 
-%5B
 dist
@@ -752,107 +752,130 @@
 _key
+%5B:4
 %5D%0A    
-if 
 code 
-in CODE_MAP:%0A        keys.append('%7B%7D%7B%7D'.format(year, CODE_MAP%5Bcode%5D))%0A    return keys
+= district_key%5B4:%5D%0A    return %5B'%7B%7D%7B%7D'.format(year, equiv_code) for equiv_code in get_equivalent_codes(code)%5D
 %0A%0A%0Ac
@@ -2483,33 +2483,33 @@
 CACHE_VERSION = 
-0
+1
 %0A    CACHE_KEY_F
@@ -2750,20 +2750,41 @@
 eviation
- == 
+.IN(get_equivalent_codes(
 abbrevia
@@ -2788,16 +2788,18 @@
 viation)
+))
 .fetch_a

Fix district history not returning equivalent code districts (#2407)