import base64
import logging
import re
import xml.etree.ElementTree as ET
import zlib
from collections import defaultdict
from copy import deepcopy
from urllib.parse import unquote

import networkx as nx
import numpy as np
import pandas as pd

logger = logging.getLogger(__name__)


class DrawIOParser:
    """Класс для парсинга, обработки и визуализации диаграмм из DrawIO."""

    def __init__(self, file_path=None):
        """Инициализация парсера.

        Параметры:
        ----------
        file_path : str, опционально
            Путь к файлу DrawIO XML
        """
        self.file_path = file_path
        self.xml_content = None
        self.root = None
        self.df = None
        self.df_full = None
        self.expanded_df = None
        self.result = None
        self.edges = None
        self.graph = None

        if file_path:
            self.load_file(file_path)

    def load_file(self, file_path):
        """Загрузка и декомпрессия XML-файла DrawIO.

        Параметры:
        ----------
        file_path : str
            Путь к файлу DrawIO XML
        """
        self.file_path = file_path
        self.xml_content = self._decompress_drawio(file_path)
        if self.xml_content:
            self.root = ET.fromstring(self.xml_content)
            self._parse_xml_to_dataframe()
            self._clean_dataframe()
            self._expand_style_column()
            self.debug_dataframe_structure()
            self._extract_tables_and_fields()
            self._extract_relationships()
            self._transform_tables_dictionary()
            self._build_graph()
            return True
        return False

    def _decompress_drawio(self, file_path):
        """Распаковка сжатого содержимого файла DrawIO.

        Параметры:
        ----------
        file_path : str
            Путь к файлу DrawIO XML

        Возвращает:
        ----------
        str или None
            Распакованное содержимое XML или None в случае ошибки
        """
        try:
            tree = ET.parse(file_path)
            
            # Проверяем формат файла
            # Если это несжатый файл, ищем корневой элемент mxGraphModel
            root_graph_model = tree.find('.//mxGraphModel')
            if root_graph_model is not None:
                # Это несжатый формат - просто возвращаем XML в строковом виде
                logger.debug('Обнаружен несжатый формат DrawIO')
                return ET.tostring(tree.getroot(), encoding='unicode')
            
            # Если это сжатый формат, обрабатываем как раньше
            diagram_element = tree.find('diagram')
            if diagram_element is not None:
                logger.debug('Обнаружен сжатый формат DrawIO, выполняем декомпрессию')
                encoded_data = diagram_element.text
                decoded_data = base64.b64decode(encoded_data)
                decompressed_data = zlib.decompress(decoded_data, wbits=-15)
                return unquote(decompressed_data.decode('utf-8'))
            
            # Если не обнаружен ни один из форматов
            logger.warning('Не обнаружена структура DrawIO в файле')
            
        except Exception as e:
            logger.error(f'Ошибка при обработке файла: {e}')
        return None

    def _parse_xml_to_dataframe(self):
        """Преобразование XML-контента в DataFrame."""
        data_dict = {'cell': [], 'value': [], 'parent': [], 'style': [], 'source': [], 'target': []}

        for cell in self.root.findall('.//mxCell'):
            data_dict['cell'] += [cell.get('id')]
            data_dict['value'] += [cell.get('value')]
            data_dict['parent'] += [cell.get('parent')]
            data_dict['style'] += [cell.get('style')]
            data_dict['source'] += [cell.get('source')]
            data_dict['target'] += [cell.get('target')]

        self.df = pd.DataFrame(data_dict)
        logger.debug(f'XML содержимое преобразовано в DataFrame с {len(self.df)} строками')

    def _clean_dataframe(self):
        """Очистка строк датафрейма от HTML-разметки."""
        self.df_full = deepcopy(self.df)
        self.df_full['value'] = self.df_full['value'].apply(self._remove_html_tags)
        logger.debug('DataFrame очищен от HTML-тегов')

    @staticmethod
    def _remove_html_tags(text):
        """Удаление HTML-тегов из текста.

        Параметры:
        ----------
        text : str
            Текст с HTML-тегами

        Возвращает:
        ----------
        str
            Очищенный текст
        """
        if text:
            return re.sub(r'<[^>]+>', '', text)
        else:
            return ''

    def _expand_style_column(self):
        """Парсинг колонки style и расширение датафрейма значениями."""
        self.expanded_df = self._parse_style_column(self.df_full)
        logger.debug('Колонка style разобрана и добавлена к DataFrame')

    def _parse_style_column(self, _df):
        """Разбор столбца style и создание новых столбцов на основе его
        содержимого.

        Параметры:
        ----------
        _df : DataFrame
            Исходный DataFrame

        Возвращает:
        ----------
        DataFrame
            Расширенный DataFrame с колонками из style
        """
        df = deepcopy(_df)

        # Разбор строки style
        def parse_style(style_str):
            if pd.isna(style_str):
                return {}

            result = {}
            for pair in style_str.split(';'):
                if '=' in pair:
                    key, value = pair.split('=', 1)
                    result[key] = value
            return result

        parsed_styles = df['style'].apply(parse_style)

        # Собираем все уникальные ключи
        all_keys = set()
        for d in parsed_styles:
            all_keys.update(d.keys())

        new_columns = defaultdict(lambda: np.repeat(None, len(df)))

        for idx, style_dict in enumerate(parsed_styles):
            for key, value in style_dict.items():
                new_columns[key][idx] = value

        result_df = df.copy()

        # Добавляем новые столбцы
        for key, values in new_columns.items():
            result_df[key] = values

        return result_df

    def debug_dataframe_structure(self):
        """Отладочный метод для анализа структуры DataFrame."""
        if self.expanded_df is None:
            logger.error('DataFrame не создан')
            return
        
        logger.debug("=== АНАЛИЗ СТРУКТУРЫ DATAFRAME ===")
        
        # Анализ всех уникальных shape
        shapes = self.expanded_df['shape'].value_counts()
        logger.debug(f"Уникальные shapes: \n{shapes}")
        
        # Анализ таблиц
        tables = self.expanded_df.loc[self.expanded_df['shape'] == 'table']
        logger.debug(f"Найдено {len(tables)} таблиц:")
        for _, table in tables.iterrows():
            logger.debug(f"  Таблица: {table['value']} (ID: {table['cell']})")
            
            # Анализ строк таблицы
            rows = self.expanded_df.loc[
                (self.expanded_df['parent'] == table['cell']) & 
                (self.expanded_df['shape'] == 'tableRow')
            ]
            logger.debug(f"    Строк в таблице: {len(rows)}")
            
            for _, row in rows.iterrows():
                # Анализ ячеек строки
                cells = self.expanded_df.loc[
                    (self.expanded_df['parent'] == row['cell']) & 
                    (self.expanded_df['shape'] == 'partialRectangle')
                ]
                cell_values = [str(val) for val in cells['value'].values]
                logger.debug(f"      Строка {row['cell']}: {cell_values}")
        
        logger.debug("=== КОНЕЦ АНАЛИЗА ===")

    def _extract_tables_and_fields(self):
        """Извлечение названий таблиц и их полей."""
        table_list = self.expanded_df.loc[self.expanded_df['shape'] == 'table', 'value'].values

        nodes = defaultdict(list)

        for table in table_list:
            parent_1 = self.expanded_df.loc[self.expanded_df['value'] == table, 'cell'].values[0]
            for cell in self.expanded_df.loc[self.expanded_df['parent'] == parent_1, 'cell']:
                if cell:
                    parent_2 = self.expanded_df.loc[self.expanded_df['parent'] == cell, 'value'].values
                    nodes[table].append(list(parent_2))

        self.nodes = nodes
        logger.debug(f'Извлечено {len(nodes)} таблиц с полями')

    def _extract_relationships(self):
        """Извлечение связей между таблицами."""
        links_df = self.expanded_df[self.expanded_df['source'].notna() & self.expanded_df['target'].notna()]

        edges = defaultdict(list)

        for source, target in zip(links_df['source'], links_df['target']):
            svalue = self.expanded_df.loc[self.expanded_df['cell'] == source, 'value'].values[0]
            if svalue == '':
                parent = self.expanded_df.loc[self.expanded_df['cell'] == source, 'parent'].values[0]
                svalue = self.expanded_df.loc[self.expanded_df['cell'] == parent, 'value'].values[0]

            rvalue = self.expanded_df.loc[self.expanded_df['cell'] == target, 'value'].values[0]
            if rvalue == '':
                rparent = self.expanded_df.loc[self.expanded_df['cell'] == target, 'parent'].values[0]
                rvalue = self.expanded_df.loc[self.expanded_df['cell'] == rparent, 'value'].values[0]

            edges[self._clean_string(svalue)].append(self._clean_string(rvalue))

        self.edges = edges
        logger.debug(f'Извлечены связи между таблицами: {len(edges)} исходных таблиц')

    @staticmethod
    def _clean_string(s):
        """Очистка строк от специальных символов и лишних пробелов.

        Параметры:
        ----------
        s : str
            Исходная строка

        Возвращает:
        ----------
        str
            Очищенная строка
        """
        # Удаляем \xa0, \n и лишние пробелы
        s = s.replace('\xa0', ' ')
        s = s.replace('&nbsp;', ' ')
        s = re.sub(r'\s+', ' ', s)
        return s.strip()

    def _transform_tables_dictionary(self):
        """Преобразование словаря таблиц: добавление префиксов PK и FK к
        полям."""
        result = defaultdict(list)

        for key, value in self.nodes.items():
            # Очищаем ключ
            cleaned_key = self._clean_string(key)

            # Очищаем значения
            transformed_values = []
            for item in value:
                if len(item) == 0:  # Пропускаем пустые списки
                    continue
                elif len(item) == 2:  # Если список содержит два элемента
                    prefix, content = item
                    if prefix in ('PK', 'FK'):
                        transformed_values.append(f'{prefix} {self._clean_string(content)}')
                    else:
                        transformed_values.append(self._clean_string(content))
                else:
                    transformed_values.append('-------------------- НЕВЕРНАЯ СТРОКА--------------------')
                    logger.warning(f'Обнаружена неверная строка в таблице {cleaned_key}: {item}')

            result[cleaned_key] = transformed_values

        self.result = result
        logger.debug('Словарь таблиц преобразован с добавлением префиксов')

    def get_tables_without_relations(self):
        """Получение списка таблиц, не имеющих связей.

        Возвращает:
        ----------
        set
            Множество таблиц без связей
        """
        unique_strings = set()

        unique_strings.update(self.edges.keys())

        for value_list in self.edges.values():
            unique_strings.update(value_list)

        tables_without_relations = set(self.result.keys()) - unique_strings
        logger.debug(f'Обнаружено {len(tables_without_relations)} таблиц без связей')
        return tables_without_relations

    def _build_graph(self):
        """Построение графа связей между таблицами."""
        # Создаем направленный граф
        G = nx.DiGraph()

        # Добавляем вершины из словаря таблиц
        for table, tags in self.result.items():
            G.add_node(table, tags=tags)  # Добавляем атрибуты (тэги) к вершинам

        # Добавляем рёбра на основе словаря связей
        for source, targets in self.edges.items():
            for target in targets:
                G.add_edge(source, target)

        self.graph = G
        logger.debug(f'Построен граф связей с {len(G.nodes())} вершинами и {len(G.edges())} рёбрами')

    def get_tables_and_fields(self):
        """Получение словаря таблиц и их полей.

        Возвращает:
        ----------
        dict
            Словарь, где ключи - названия таблиц, значения - списки полей
        """
        return dict(self.result)

    def get_table_relationships(self):
        """Получение словаря связей между таблицами.

        Возвращает:
        ----------
        dict
            Словарь, где ключи - исходные таблицы, значения - списки целевых таблиц
        """
        return dict(self.edges)

    @staticmethod
    def parse_schema(json_data):
        result = {}

        for table_name, columns in json_data.items():
            result[table_name] = {}

            for column_info in columns:
                # Проверяем, содержит ли строка разделитель ": "
                if ': ' in column_info:
                    parts = column_info.split(': ', 1)
                    column_name = parts[0]
                    column_type = parts[1]
                else:
                    column_name = column_info
                    column_type = ""

                result[table_name][column_name] = column_type

        return result
