import base64
import html
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 numpy as np
import pandas as pd

logger = logging.getLogger(__name__)


class DrawIOParser:
    def __init__(self, file_path=None):
        """Инициализация парсера для файлов DrawIO.

        Args:
            file_path (str, optional): Путь к файлу .drawio.xml
        """
        self.file_path = file_path
        self.root = None
        self.df = None
        self.expanded_df = None
        self.tables_content = None
        self.cleaned_dict = None
        self.edges = None

        if file_path:
            self.load_file(file_path)

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

        Args:
            file_path (str): Путь к файлу .drawio.xml

        Returns:
            self: Для поддержки цепочки вызовов
        """
        self.file_path = file_path
        xml_content = self._decompress_drawio(file_path)
        self.root = ET.fromstring(xml_content)
        logger.debug('Файл успешно загружен и распакован')
        return self

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

        Args:
            file_path (str): Путь к файлу .drawio.xml

        Returns:
            str: Распакованное содержимое XML
        """
        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("Элемент 'diagram' не найден в XML файле")
        except Exception as e:
            logger.error(f'Ошибка при декомпрессии файла: {e}')
        return None

    def parse_to_dataframe(self):
        """Преобразование XML в pandas DataFrame.

        Returns:
            self: Для поддержки цепочки вызовов
        """
        if self.root is None:
            logger.error('Необходимо сначала загрузить файл с помощью load_file()')
            raise ValueError('Необходимо сначала загрузить файл с помощью load_file()')

        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)} строками')
        return self

    def expand_style_column(self):
        """Расширение DataFrame за счёт разбора колонки "style".

        Returns:
            self: Для поддержки цепочки вызовов
        """
        if self.df is None:
            logger.error('Необходимо сначала выполнить parse_to_dataframe()')
            raise ValueError('Необходимо сначала выполнить parse_to_dataframe()')

        df = deepcopy(self.df)

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

            result = {}
            parts = style_str.split(';')
            for part in parts:
                if '=' in part:
                    key, value = part.split('=', 1)
                    result[key.strip()] = value.strip()
                else:
                    # Если нет '=', то это значение для колонки "shape"
                    result['shape'] = part.strip()
            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

        self.expanded_df = result_df
        logger.debug(f'Колонка style разобрана, добавлено {len(all_keys)} новых атрибутов')
        return self

    def extract_tables(self):
        """Извлечение названий таблиц и списка их полей.

        Returns:
            self: Для поддержки цепочки вызовов
        """
        if self.expanded_df is None:
            logger.error('Необходимо сначала выполнить expand_style_column()')
            raise ValueError('Необходимо сначала выполнить expand_style_column()')

        table_names = self.expanded_df.loc[
            (self.expanded_df['shape'] == 'text') & (self.expanded_df['fontColor'] == '#ffffff'),
            'value',
        ].values.tolist()

        logger.debug(f'Найдено {len(table_names)} потенциальных таблиц')

        tables_content = defaultdict(list)

        for table_name in table_names:
            try:
                cell = self.expanded_df.loc[self.expanded_df['value'] == table_name, 'parent'].values[0]
                fields = self.expanded_df.loc[self.expanded_df['parent'] == cell, 'value'].values
                tables_content[table_name] += list(fields)
                logger.debug(f"Извлечено {len(fields)} полей для таблицы '{table_name}'")
            except IndexError:
                # Пропускаем, если не можем найти родительскую ячейку
                logger.warning(f"Не удалось найти родительскую ячейку для таблицы '{table_name}'")
                continue

        self.tables_content = tables_content

        return self

    def _clean_string(self, input_string):
        """Очистка строки от разметочного мусора.

        Args:
            input_string (str): Исходная строка

        Returns:
            list: Список очищенных частей строки
        """
        # Разбиваем строку по <br/>
        parts = input_string.split('<br/>')

        # Очищаем каждую часть от escape-символов (например, &#x9;)
        cleaned_parts = [re.sub(r'&#x[0-9A-Fa-f]+;', '', part).strip() for part in parts]
        # Декодируем HTML-сущности
        cleaned_parts = [html.unescape(part) for part in cleaned_parts]

        # Убираем пустые строки
        cleaned_parts = [part for part in cleaned_parts if part]

        return cleaned_parts

    def _clean_dict(self, input_dict):
        """Очистка всего словаря от разметочного мусора.

        Args:
            input_dict (dict): Исходный словарь

        Returns:
            dict: Очищенный словарь
        """
        cleaned_dict = defaultdict(list)
        for key, values in input_dict.items():
            cleaned_values = []
            for value in values:
                if value is None:
                    continue

                # Если значение содержит <br/>, разбиваем его
                if '<br/>' in value:
                    cleaned_values.extend(self._clean_string(value))
                else:
                    # Если нет <br/>, просто очищаем от лишних символов
                    cleaned_value = re.sub(r'&#x[0-9A-Fa-f]+;', '', value).strip()
                    cleaned_value = html.unescape(cleaned_value)
                    if cleaned_value:
                        cleaned_values.append(cleaned_value)
            cleaned_dict[key] = cleaned_values

        logger.debug(f'Словарь очищен от HTML-разметки и управляющих символов')
        return cleaned_dict

    def _remove_keys_from_values(self, input_dict):
        """Удаление из значений ключей, если они дублируются.

        Args:
            input_dict (dict): Исходный словарь

        Returns:
            dict: Словарь с удаленными дубликатами
        """
        result_dict = {}

        for key, values in input_dict.items():
            # Удаляем из значений элементы, которые совпадают с ключом
            cleaned_values = [value for value in values if value != key]
            result_dict[key] = cleaned_values

        return result_dict

    def clean_tables(self):
        """Очистка словаря таблиц от разметочного мусора и дубликатов.

        Returns:
            self: Для поддержки цепочки вызовов
        """
        if self.tables_content is None:
            logger.error('Необходимо сначала выполнить extract_tables()')
            raise ValueError('Необходимо сначала выполнить extract_tables()')

        cleaned_data = self._clean_dict(self.tables_content)
        self.cleaned_dict = self._remove_keys_from_values(cleaned_data)
        return self

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

        Returns:
            self: Для поддержки цепочки вызовов
        """
        if self.expanded_df is None:
            logger.error('Необходимо сначала выполнить expand_style_column()')
            raise ValueError('Необходимо сначала выполнить expand_style_column()')

        edges = defaultdict(list)

        links_df = self.expanded_df[self.expanded_df['source'].notna() & self.expanded_df['target'].notna()]
        logger.debug(f'Найдено {len(links_df)} потенциальных связей между таблицами')

        for source, target in zip(links_df['source'], links_df['target']):
            try:
                # Находим имена таблиц по cell ID
                svalue = self.expanded_df.loc[
                    (self.expanded_df['parent'] == source) & (self.expanded_df['fontColor'] == '#ffffff'),
                    'value',
                ].values[0]

                tvalue = self.expanded_df.loc[
                    (self.expanded_df['parent'] == target) & (self.expanded_df['fontColor'] == '#ffffff'),
                    'value',
                ].values[0]

                edges[svalue].append(tvalue)
                logger.debug(f"Найдена связь: '{svalue}' -> '{tvalue}'")
            except IndexError:
                # Пропускаем, если не можем найти имя таблицы
                logger.warning(f'Не удалось определить связь между таблицами для source={source}, target={target}')
                continue

        self.edges = edges
        return self

    def get_tables(self):
        """Основной метод 1: Получение очищенного словаря с названиями таблиц и
        их полями.

        Returns:
            dict: Словарь с таблицами и их полями
        """
        if self.cleaned_dict is None:
            # Если данные еще не были обработаны, запускаем полный пайплайн
            self.parse_to_dataframe().expand_style_column().extract_tables().clean_tables()

        return self.cleaned_dict

    def get_relationships(self):
        """Основной метод 2: Получение словаря со списком смежности (связи
        между таблицами)

        Returns:
            dict: Словарь связей между таблицами
        """
        if self.edges is None:
            # Если связи еще не были извлечены, запускаем необходимые шаги
            if self.expanded_df is None:
                self.parse_to_dataframe().expand_style_column()
            self.extract_relationships()

        return self.edges

    def process_all(self):
        """Запуск полного пайплайна обработки данных.

        Returns:
            self: Для поддержки цепочки вызовов
        """
        result = self.parse_to_dataframe().expand_style_column().extract_tables().clean_tables().extract_relationships()
        return result
