from itertools import groupby, chain
import time

from octostar.utils.ontology import multiquery_ontology
from octostar.utils.exceptions import StopAsyncIterationWithResult

from inheritance import is_child_concept

REL_FETCHED_FIELDS = '`os_relationship_name`,`os_relationship_label`,`os_entity_uid_from`,`os_entity_uid_to`,`os_entity_uid`,`os_entity_type_from`,`os_entity_type_to`'
TARGETS_FETCHED_FIELDS = '`os_entity_uid`,`os_concept`,`entity_label`'

class ExecutionMetrics:
    def __init__(self):
        self.n_queries = 0
        self.n_relationships = 0
        self.n_target_entities = 0
        self.exec_time = 0.0
        self.exec_times = {}
        self.n_cancelled_queries = 0
        self.relationship_names = set()
        self.timeout = 0

    def print_metrics(self):
        print("Execution Metrics:")
        print(f"  Execution Time: {self.exec_time}")
        print(f"  Per-Type Execution Times: {self.exec_times}")
        print(f"  Number of Queries: {self.n_queries}")
        print(f"  Number of Cancelled Queries: {self.n_cancelled_queries}")
        print(f"  Timeout per Query: {self.timeout}")
        print(f"  Number of Relationships: {self.n_relationships}")
        print(f"  Number of Target Entities: {self.n_target_entities}")
        print(f"  Relationship Names: {self.relationship_names}")

async def _get_stream_result(stream):
    result = None
    try:
        async for _ in stream:
            pass
    except StopAsyncIterationWithResult as e:
        result = e.value
    return result

def _left_join(left_list, right_list, left_keys, right_keys):
    left_keys = left_keys if isinstance(left_keys, (list, tuple)) else [left_keys]
    right_keys = right_keys if isinstance(right_keys, (list, tuple)) else [right_keys]
    def _make_composite_key(item, keys):
        return tuple(item[key] for key in keys)
    left_list.sort(key=lambda item: _make_composite_key(item, left_keys))
    right_list.sort(key=lambda item: _make_composite_key(item, right_keys))
    left_groups = groupby(left_list, key=lambda item: _make_composite_key(item, left_keys))
    right_groups = groupby(right_list, key=lambda item: _make_composite_key(item, right_keys))
    right_dict = {key: list(group) for key, group in right_groups}
    result = []
    for left_key_val, left_items in left_groups:
        left_items = list(left_items)
        associated_rights = right_dict.get(left_key_val, [])
        for left_item in left_items:
            result.append((left_item, associated_rights))
    return result

async def expand_otm(expanded_entities, entities_by_concept, ontology_rels, concepts_to_otm_rels, metrics, client, timeout, limit):
    start_time = time.time()
    otm_queries = []
    for concept_name, rels in concepts_to_otm_rels.items():
        for rel_name in rels:
            rel = ontology_rels[rel_name]
            source_properties = rel['source_properties'].split(",")
            source_properties_values = [[entity[p] for p in source_properties] for entity in entities_by_concept[concept_name]]
            source_properties_values = [values for values in source_properties_values if all(value is not None for value in values)]
            if source_properties_values:
                target_prop_names = rel['target_properties'].split(",")
                target_prop_names = ",".join(["`" + prop + "`" for prop in target_prop_names])
                source_properties_values = ",".join(["(" + ",".join(["'" + value + "'" for value in values]) + ")" for values in source_properties_values])
                query = f"""SELECT {TARGETS_FETCHED_FIELDS} FROM `etimbr`.`{rel['target_concept']}` WHERE ({target_prop_names}) IN ({source_properties_values}) LIMIT {limit}"""
                otm_queries.append({
                    'concept_name': concept_name,
                    'query': query,
                    'relationship_name': rel_name,
                    'source_properties': rel['source_properties'].split(","),
                    'target_properties': rel['target_properties'].split(",")
                })
    otm_queries = {str(i): otm_queries[i] for i in range(len(otm_queries))}
    otm_stream = multiquery_ontology.streaming(
        sql_queries={query_id: query_data['query'] for query_id, query_data in otm_queries.items()},
        client=client, timeout=timeout
    )
    metrics.n_queries += len(otm_queries)
    all_results = await _get_stream_result(otm_stream)
    for query_id, data in all_results.items():
        if data == None:
            metrics.n_cancelled_queries += 1
        if not data:
            continue
        query_data = otm_queries[query_id]
        data = _left_join(
            [e[0] for e in expanded_entities.values() if e[0]['os_concept'] == query_data['concept_name']], data,
            query_data['source_properties'], query_data['target_properties']
        )
        data = {e[0]['os_entity_uid']: (e[0], {r['os_entity_uid']: ({'os_relationship_name': query_data['relationship_name']}, {r['os_entity_uid']: r}) for r in e[1]}) for e in data}
        for entity_id, entities in data.items():
            if entity_id not in expanded_entities:
                expanded_entities[entity_id] = [entities[0], {}]
            for sub_entity_id, sub_entities in entities[1].items():
                if sub_entity_id not in expanded_entities[entity_id][1]:
                    expanded_entities[entity_id][1][sub_entity_id] = [sub_entities[0], {}]
                expanded_entities[entity_id][1][sub_entity_id][1].update(sub_entities[1])
    metrics.exec_times['otm'] = time.time() - start_time
    return expanded_entities

async def expand_mtm_mixed(expanded_entities, entities_by_concept, ontology, concepts_to_mixed_rels, ontology_rels, metrics, client, timeout, limit):
    start_time = time.time()
    mixed_queries = []
    match_patterns = {}
    os_workspace_relationship_fields = [prop['property_name'] for prop in ontology['concepts']['os_workspace_relationship']['allProperties']]
    for concept_name, rels in concepts_to_mixed_rels.items():
        if concept_name not in match_patterns:
            match_patterns[concept_name] = []
        match_patterns[concept_name].extend([entity['os_entity_uid'] for entity in entities_by_concept[concept_name]])
    for concept_name, entity_ids in match_patterns.items():
        rels = concepts_to_mixed_rels[concept_name]
        entity_ids = ",".join(["'" + uid + "'" for uid in entity_ids])
        for rel_name in rels:
            rel = ontology_rels[rel_name]
            target_concept = rel['target_concept']
            relationship_fields = ",".join([f"`{rel_name}[{target_concept}]_" + prop  + "` AS " + "`rel__" + prop + "`" for prop in os_workspace_relationship_fields])
            target_concept_fields = [prop['property_name'] for prop in ontology['concepts'][rel['target_concept']]['allProperties']]
            target_concept_fields = ",".join([f"`{rel_name}[{target_concept}]." + prop + "` AS " + "`tgt__" + prop + "`" for prop in target_concept_fields])
            mixed_queries.append({
                'query': f"SELECT {relationship_fields}, {target_concept_fields}, `os_entity_uid` FROM `dtimbr`.`{concept_name}` WHERE `os_entity_uid` IN ({entity_ids}) LIMIT {limit}"
            })
    mixed_queries = {str(i): mixed_queries[i] for i in range(len(mixed_queries))}
    metrics.n_queries += len(mixed_queries)
    local_stream = multiquery_ontology.streaming(
        sql_queries={query_id: query_data['query'] for query_id, query_data in mixed_queries.items()},
        client=client, timeout=timeout/2.0
    )
    all_results = await _get_stream_result(local_stream)
    for query_id, data in all_results.items():
        if data == None:
            metrics.n_cancelled_queries += 1
        if not data:
            continue
        data = _left_join(
            [e[0] for e in expanded_entities.values()], data,
            'os_entity_uid', 'os_entity_uid'
        )
        data = {e[0]['os_entity_uid']: (e[0], {r['os_entity_uid']: ({}, {r['os_entity_uid']: r}) for r in e[1]}) for e in data}
        for entity_id, entities in data.items(): ### CHECK
            if entity_id not in expanded_entities:
                expanded_entities[entity_id] = [entities[0], {}]
            for sub_entity_id, sub_entities in entities[1].items():
                if sub_entity_id not in expanded_entities[entity_id][1]:
                    expanded_entities[entity_id][1][sub_entity_id] = [sub_entities[0], {}]
                expanded_entities[entity_id][1][sub_entity_id][1].update(sub_entities[1])
    metrics.exec_times['mtm_mixed'] = time.time() - start_time
    return expanded_entities

async def expand_mtm_local(expanded_entities, entities_by_concept, ontology_rels, concepts_to_local_rels, metrics, client, timeout, limit):
    start_time = time.time()
    local_queries = {
        'from': [],
        'to': []
    }
    match_patterns = {}
    for concept_name in concepts_to_local_rels.keys():
        if concept_name not in match_patterns:
            match_patterns[concept_name] = []
        match_patterns[concept_name].extend([entity['os_entity_uid'] for entity in entities_by_concept[concept_name]])
    for concept_name, entity_ids in match_patterns.items():
        rel_names = concepts_to_local_rels[concept_name]
        inverse_names = [ontology_rels[rel_name]['inverse_name'] for rel_name in rel_names]
        rel_names = ",".join(["'" + rel_name + "'" for rel_name in rel_names])
        inverse_names = ",".join(["'" + rel_name + "'" for rel_name in inverse_names])
        entity_ids = ",".join(["'" + uid + "'" for uid in entity_ids])
        local_queries['from'].append({
            'query': f"SELECT {REL_FETCHED_FIELDS} FROM `timbr`.`os_workspace_relationship` WHERE `os_relationship_name` IN ({rel_names}) AND `os_entity_uid_from` IN ({entity_ids}) LIMIT {limit}"
        })
        local_queries['to'].append({
            'query': f"SELECT {REL_FETCHED_FIELDS} FROM `timbr`.`os_workspace_relationship` WHERE `os_relationship_name` IN ({inverse_names}) AND `os_entity_uid_to` IN ({entity_ids}) LIMIT {limit}"
        })
    local_queries = {
        **{'from_' + str(i): local_queries['from'][i] for i in range(len(local_queries['from']))},
        **{'to_' + str(i): local_queries['to'][i] for i in range(len(local_queries['to']))},
    }
    metrics.n_queries += len(local_queries)
    local_stream = multiquery_ontology.streaming(
        sql_queries={query_id: query_data['query'] for query_id, query_data in local_queries.items()},
        client=client, timeout=timeout/2.0
    )
    all_results = await _get_stream_result(local_stream)
    middle_entities = []
    for query_id, data in all_results.items():
        if data == None:
            metrics.n_cancelled_queries += 1
        if not data:
            continue
        query_id = query_id.split("_")
        data = _left_join(
            [e[0] for e in expanded_entities.values()], data, 'os_entity_uid', 'os_entity_uid_' + query_id[0]
        )
        inverse_direction = 'to' if query_id[0] == 'from' else 'from'
        data = {e[0]['os_entity_uid']: (e[0], {r['os_entity_uid']: (r, {}) for r in e[1]}) for e in data}
        middle_entities.extend([(inverse_direction, r) for e in data.values() for r in e[1].values()])
        for entity_id, entities in data.items():
            if entity_id not in expanded_entities:
                expanded_entities[entity_id] = [entities[0], {}]
            for sub_entity_id, sub_entities in entities[1].items():
                if sub_entity_id not in expanded_entities[entity_id][1]:
                    expanded_entities[entity_id][1][sub_entity_id] = [sub_entities[0], {}]
                expanded_entities[entity_id][1][sub_entity_id][1].update(sub_entities[1])
    middle_entities = [{
        'entity_id': r[1][0]['os_entity_uid_' + r[0]],
        'concept_name': r[1][0]['os_entity_type_' + r[0]],
        'direction': r[0],
        'relationship': r[1]
    } for r in middle_entities]
    middle_entities = sorted(middle_entities, key=lambda x: x['concept_name'])
    middle_entities = groupby(middle_entities, key=lambda x: x['concept_name'])
    middle_entities = {e[0]: list(e[1]) for e in middle_entities}
    local_queries_2 = []
    for concept_name, entities in middle_entities.items():
        entity_ids = ",".join(["'" + e['entity_id'] + "'" for e in entities])
        local_queries_2.append({
            'query': f"SELECT {TARGETS_FETCHED_FIELDS} FROM `timbr`.`{concept_name}` WHERE `os_entity_uid` IN ({entity_ids}) LIMIT {limit}"
        })
    local_queries_2 = {str(i): local_queries_2[i] for i in range(len(local_queries_2))}
    metrics.n_queries += len(local_queries_2)
    local_stream_2 = multiquery_ontology.streaming(
        sql_queries={query_id: query_data['query'] for query_id, query_data in local_queries_2.items()},
        client=client, timeout=timeout/2.0
    )
    all_results = await _get_stream_result(local_stream_2)
    middle_entities = list(chain(*middle_entities.values()))
    for query_id, data in all_results.items():
        if data == None:
            metrics.n_cancelled_queries += 1
        if not data:
            continue
        data = _left_join(
            middle_entities, data, 'entity_id', 'os_entity_uid'
        )
        for entry in data:
            relationship_targets = entry[0]['relationship']
            for target_entity in entry[1]:
                relationship_targets[1][target_entity['os_entity_uid']] = target_entity
    if data:
        for rel in data:
            targets = rel[0]['relationship'][1]
            rel_uid = rel[0]['relationship'][0]['os_entity_uid']
            for entity in expanded_entities.values():
                entity_rels = entity[1]
                if rel_uid in entity_rels:
                    entity_rels[rel_uid][1].update(targets)
    metrics.exec_times['mtm_local'] = time.time() - start_time
    return expanded_entities

async def expand_entities(entities, ontology, relationship_mappings_info, client, relationship_names_by_entity_type=None, batch_size = 10000, avg_limit_per_entity = 20, timeout=10.0):
    metrics = ExecutionMetrics()
    metrics.timeout = timeout
    start_time = time.time()
    ordered_entities = sorted(entities, key=lambda x: x['entity_type'])
    ordered_entities = list({e['os_entity_uid']: e for e in entities}.values())
    entity_batches = [ordered_entities[i:i+batch_size] for i in range(0, len(ordered_entities), batch_size)]
    expanded_entities = {}
    for batch in entity_batches:
        new_entities = await expand_entities_batch(batch, ontology, relationship_mappings_info, relationship_names_by_entity_type, metrics, client, avg_limit_per_entity, timeout)
        expanded_entities.update(new_entities)
    metrics.exec_time = time.time() - start_time
    for _, entity_data in expanded_entities.items():
        for rel in entity_data[1].values():
            if rel[0]:
                metrics.n_relationships += 1
                metrics.relationship_names.add(rel[0]['os_relationship_name'])
            metrics.n_target_entities += len(rel[1])
    metrics.relationship_names = list(metrics.relationship_names)
    expanded_entities = [(e[0], [(r[0], list(r[1].values())[0]) for r in e[1].values()]) for e in expanded_entities.values()]
    return expanded_entities, metrics

async def expand_entities_batch(entities, ontology, relationship_mappings, relationship_names_by_entity_type, metrics, client, avg_limit_per_entity = 20, timeout=10.0):
    limit = avg_limit_per_entity * len(entities)
    expanded_entities = {e['os_entity_uid']: [e, {}] for e in entities}
    entities_by_concept = sorted(entities, key=lambda x: x['entity_type'])
    entities_by_concept = groupby(entities_by_concept, key=lambda x: x['entity_type'])
    entities_by_concept = {e[0]: list(e[1]) for e in entities_by_concept}
    ontology_rels = {r['relationship_name']: r for r in ontology['relationships']}
    concepts_to_rels = {cn: c['relationships'] for cn, c in ontology['concepts'].items()}
    concepts_to_mixed_rels = {}
    concepts_to_otm_rels = {}
    concepts_to_local_rels = {}
    for concept_name in entities_by_concept.keys():
        relationship_names = None
        if relationship_names_by_entity_type:
            relationship_names = set()
            for cn_name in relationship_names_by_entity_type.keys():
                if is_child_concept(concept_name, cn_name, ontology):
                    relationship_names = relationship_names.union(set(relationship_names_by_entity_type[cn_name]))
        rels = concepts_to_rels[concept_name]
        if relationship_names:
            rels = [r for r in rels if r in relationship_names]
        filtered_mtm_rels = []
        filtered_otm_rels = []
        filtered_local_mtm_rels = []
        for rel in rels:
            if not ontology_rels[rel]['is_mtm']:
                filtered_otm_rels.append(rel)
            elif rel in relationship_mappings['unmapped']:
                continue
            elif rel in relationship_mappings['local_only']:
                filtered_local_mtm_rels.append(rel)
            else:
                filtered_mtm_rels.append(rel)
        concepts_to_mixed_rels[concept_name] = filtered_mtm_rels
        concepts_to_otm_rels[concept_name] = filtered_otm_rels
        concepts_to_local_rels[concept_name] = filtered_local_mtm_rels
    # OTM QUERIES: fetch target entities directly from the target tables with an IN statement
    expanded_entities = await expand_otm(expanded_entities, entities_by_concept, ontology_rels, concepts_to_otm_rels, metrics, client, timeout, limit)
    # LOCAL MTM QUERIES: fetch relationships directly from the os_workspace_relationship table, then fetch target entities
    expanded_entities = await expand_mtm_local(expanded_entities, entities_by_concept, ontology_rels, concepts_to_local_rels, metrics, client, timeout, limit)
    # MIXED MTM QUERIES: query via timbr to make the most of JOINs
    expanded_entities = await expand_mtm_mixed(expanded_entities, entities_by_concept, ontology, concepts_to_mixed_rels, ontology_rels, metrics, client, timeout, limit)
    return expanded_entities
    
