import uuid
from aworld.logs.util import logger
from aworld.trace.opentelemetry.memory_storage import SpanModel
from aworld.trace.constants import RunType, SPAN_NAME_PREFIX_EVENT_AGENT
from aworld.trace.instrumentation import semconv


def build_trace_tree(spans: list[SpanModel]):
    spans_dict = {span.span_id: span.dict() for span in spans}
    for span in list(spans_dict.values()):
        parent_id = span['parent_id'] if span['parent_id'] else None
        if parent_id:
            parent_span = spans_dict.get(parent_id)
            if not parent_span:
                logger.warning(f"span[{parent_id}] not be exported")
                parent_span = {
                    'span_id': parent_id,
                    'trace_id': span['trace_id'],
                    'name': 'Pengding-Span',
                    'start_time': span['start_time'],
                    'end_time': span['end_time'],
                    'duration_ms': span['duration_ms'],
                    'attributes': {},
                    'status': {},
                    'parent_id': None,
                    'run_type': 'OTHER'
                }
                spans_dict[parent_id] = parent_span
            if 'children' not in parent_span:
                parent_span['children'] = []
            parent_span['children'].append(span)

    root_spans = [span for span in spans_dict.values()
                  if span['parent_id'] is None]
    return root_spans


def _get_agent_show_name(span: dict):
    agent_name_prefix = SPAN_NAME_PREFIX_EVENT_AGENT
    name = span.get("name")
    if name and name.startswith(agent_name_prefix):
        name = name[len(agent_name_prefix):]
    if name and '---' in name:
        name = name.split('---', 1)[0]
    return name


def _remove_span_detail(root_spans: list):
    keys_to_keep = {'span_id', 'show_name', 'task_group_id', 'event_id'}
    for span in root_spans:
        keys_to_remove = [key for key in span.keys() if key not in keys_to_keep]
        for key in keys_to_remove:
            span.pop(key, None)
        if 'children' in span:
            _remove_span_detail(span['children'])


def _get_top_task_nodes(spans_dict):
    task_nodes = [
        span for span in spans_dict.values()
        if span.get('name', '').startswith('task.')
    ]
    top_task_nodes = []
    for task_node in task_nodes:
        parent_id = task_node.get('parent_id')
        is_top = True

        while parent_id:
            parent_span = spans_dict.get(parent_id)
            if not parent_span:
                break

            if parent_span.get('name', '').startswith('task.'):
                is_top = False
                break

            parent_id = parent_span.get('parent_id')

        if is_top:
            top_task_nodes.append(task_node)

    return top_task_nodes


def _get_root_nodes(edges):
    sources = set()
    targets = set()
    for edge in edges:
        sources.add(edge['source'])
        targets.add(edge['target'])
    return sources - targets


def _build_graph(root_spans: list):
    nodes = []
    edges = []

    group_id_counter = 0

    def __process_group_span(parent_spans, group_id, group_spans):
        nonlocal group_id_counter
        group_id_counter += 1
        # add group node
        group_node = {
            'span_id': f'group_{group_id_counter}',
            'group_id': group_id,
            'show_name': 'Task Group'
        }
        nodes.append(group_node)

        # add edges from parent_spans to group node
        for parent_span in parent_spans:
            edges.append({
                'source': parent_span['span_id'],
                'target': group_node['span_id']
            })

        # add edges from group node to children spans
        last_spans = []
        for child in group_spans:
            edges.append({
                'source': group_node['span_id'],
                'target': child['span_id']
            })
            last_spans.extend(__process_span(child))
        return last_spans

    def __process_span(span):
        nonlocal group_id_counter
        nodes.append(span)
        if 'children' in span:
            groups = {}
            for child in span['children']:
                group_id = child.get('task_group_id', id(child))
                if group_id not in groups:
                    groups[group_id] = []
                groups[group_id].append(child)

            last_spans = [span]  # The leaf nodes of the current subtree
            for group_id, group_spans in groups.items():
                if len(group_spans) > 1:
                    parent_spans = last_spans
                    last_spans = __process_group_span(parent_spans, group_id, group_spans)
                else:
                    child_span = group_spans[0]
                    # add edges from last_spans to child
                    for prev_node in last_spans:
                        edges.append({
                            'source': prev_node['span_id'],
                            'target': child_span['span_id']
                        })
                    last_spans = __process_span(child_span)
        return last_spans

    for span in root_spans:
        __process_span(span)

    return {
        'nodes': nodes,
        'edges': edges
    }


def get_agent_flow(trace_id):
    from aworld.trace.server import get_trace_server

    storage = get_trace_server().get_storage()
    spans = storage.get_all_spans(trace_id)
    spans_dict = {span.span_id: span.dict() for span in spans}
    children_spans = []
    top_task_nodes = _get_top_task_nodes(spans_dict)

    filtered_spans = {}
    for span_id, span in spans_dict.items():
        if span.get('is_event', False) and span.get('run_type') == RunType.AGNET.value:
            span['show_name'] = _get_agent_show_name(span)
            span['event_id'] = span.get('attributes', {}).get('event.id')
            filtered_spans[span_id] = span

    sub_task_spans = []
    for span in list(filtered_spans.values()):
        skip_this_span = False
        parent_id = span['parent_id'] if span['parent_id'] else None

        while parent_id and parent_id not in filtered_spans:
            parent_span = spans_dict.get(parent_id)
            if parent_span and parent_span.get('run_type') == RunType.TASK.value:
                span['task_group_id'] = parent_span['attributes'].get(semconv.TASK_GROUP_ID)
            parent_id = parent_span['parent_id'] if parent_span and parent_span['parent_id'] else None

        if skip_this_span:
            continue
        if parent_id:
            parent_span = filtered_spans.get(parent_id)
            if not parent_span:
                continue

            if 'children' not in parent_span:
                parent_span['children'] = []
            parent_span['children'].append(span)
            children_spans.append(span)

    filtered_span_list = [span for span in filtered_spans.values() if span not in sub_task_spans]
    root_spans = [span for span in filtered_span_list
                  if span not in children_spans]

    data = _build_graph(root_spans)
    _remove_span_detail(data["nodes"])

    # add query start node
    _add_query_node(data, top_task_nodes)
    return data

def _add_query_node(data, top_task_nodes):
    top_task_node = top_task_nodes[0] if top_task_nodes else None
    if top_task_node:
        start_node_span_id = f'{uuid.uuid4().hex[:16]}'
        data['nodes'].append({
            'span_id': start_node_span_id,
            'show_name': top_task_node['attributes'].get(semconv.TASK_INPUT),
        })
        root_nodes = _get_root_nodes(data['edges'])
        if root_nodes:
            for root_node in root_nodes:
                data['edges'].append({
                    'source': start_node_span_id,
                    'target': root_node,
                })
        else:
            for root_node in data['nodes']:
                if root_node['span_id'] != start_node_span_id:
                    data['edges'].append({
                        'source': start_node_span_id,
                        'target': root_node['span_id'],
                    })
