from odoo import models, fields, api
import networkx as nx
import matplotlib.pyplot as plt
import io
import base64
import logging

class HgIndex(models.Model):
    _inherit = 'hg.index'

    def generate_graph_image(self):
        hg_link = self.env['hg.link'].search([('source_id','=', self.node_id.id)])
        nodes = self.node_id | hg_link.mapped('target_ids')
        internal_code = self.internal_code_id
        Index = self.env['hg.index'].sudo()
        Index_by_code = Index.search([('internal_code_id','=', internal_code.id)])


        G = nx.DiGraph()

        # Добавляем все вершины: node_id и related_ids
        # Используем node_id.id как ключ, но для отображения можно имя
        for rec in hg_link:
            if rec.source_id:
                node_key = rec.source_id.name
                # Добавляем узел с именем, например, названием
                label_value = sum(Index_by_code.mapped('current_value'))
                G.add_node(node_key, label=label_value)

                # Добавляем связанные узлы как отдельные вершины и ребра
                for related_node in rec.target_ids:
                    related_key = related_node.name
                    label_values = Index_by_code.filtered(lambda r: r.node_id == related_node)
                    label_value = sum(label_values.mapped('current_value'))
                    G.add_node(related_key, label=label_value)
                    G.add_edge(node_key, related_key)

        # Определяем расположение узлов
        pos = nx.spring_layout(G, seed=1)

        plt.figure(figsize=(10, 8))

        # Рисуем граф без автоподписей (with_labels=False)
        nx.draw(G, pos, node_size=3000, node_color='lightblue', with_labels=False)

        # Формируем словарь меток с безопасным приведением к числу
        labels = {}
        for n, d in G.nodes(data=True):
            try:
                label_value = float(d.get("label", 0))
                labels[n] = f'{n}\n{label_value:.2f}'
            except (ValueError, TypeError):
                labels[n] = f'{n}\n{d.get("label", "")}'

        # Отрисовываем метки отдельно
        nx.draw_networkx_labels(G, pos, labels)

        # Сохраняем рисунок в буфер
        buf = io.BytesIO()
        plt.savefig(buf, format='png')
        plt.close()
        buf.seek(0)

        # Кодируем изображение в base64 для отображения в вебе
        image_data = base64.b64encode(buf.read()).decode('utf-8')
        return image_data