import json
import os
import traceback

# --- Конфигурация ---
# Базовый путь к директории, где Zeppelin хранит все ноутбуки.
# Убедитесь, что этот путь корректен для вашей установки Zeppelin.
# Возможно, понадобится полный абсолютный путь, например, /zeppelin/notebook
# Или относительный путь, если ваш скрипт запускается из определенной директории.

ZEPPELIN_BASE_DIR = "/notebook" # Используем /notebook как пример базовой директории

def find_notebook_path(notebook_path_prefix, base_dir=ZEPPELIN_BASE_DIR):
    """
    Находит полный путь к файлу ноутбука Zeppelin по его префиксу пути.

    Args:
        notebook_path_prefix (str): Путь к ноутбуку без ID и расширения,
                                     например "Libraries/main_config" или "MyFolder/SubFolder/MyNotebook".
        base_dir (str): Базовая директория, где Zeppelin хранит ноутбуки.

    Returns:
        str: Полный путь к найденному файлу ноутбука .zpln.

    Raises:
        FileNotFoundError: Если ноутбук не найден или директория не существует.
        ValueError: Если найдено несколько ноутбуков, соответствующих префиксу.
    """
    # Нормализуем префикс пути (убираем '/' в начале/конце, если есть)
    normalized_prefix = notebook_path_prefix.strip('/')

    # Разделяем путь на директорию и базовое имя файла
    if '/' in normalized_prefix:
        parts = normalized_prefix.rsplit('/', 1)
        notebook_dir_relative = parts[0]
        notebook_base_name = parts[1]
    else:
        # Если путь без '/', значит ноутбук в корневой директории относительно base_dir
        notebook_dir_relative = ""
        notebook_base_name = normalized_prefix

    # Формируем полный путь к директории, где должен лежать ноутбук
    target_dir = os.path.join(base_dir, notebook_dir_relative)

    print(f"Ищем ноутбук с базовым именем '{notebook_base_name}' в директории: {target_dir}")

    if not os.path.isdir(target_dir):
        raise FileNotFoundError(f"Директория не найдена: {target_dir}")

    matches = []
    try:
        for filename in os.listdir(target_dir):
            # Проверяем, что файл имеет расширение .zpln
            if filename.endswith(".zpln"):
                # Убираем расширение
                name_without_ext = filename[:-5] # Длина ".zpln" = 5
                # Ищем последний символ '_' который разделяет имя и ID
                last_underscore_index = name_without_ext.rfind('_')

                # Проверяем, что '_' найден и часть до него совпадает с искомым базовым именем
                if last_underscore_index != -1 and name_without_ext[:last_underscore_index] == notebook_base_name:
                     # Проверяем, что ID после '_' не пустой (на всякий случай)
                     if last_underscore_index < len(name_without_ext) - 1:
                         matches.append(os.path.join(target_dir, filename))

    except OSError as e:
        raise FileNotFoundError(f"Ошибка доступа к директории {target_dir}: {e}")

    # Анализируем результаты поиска
    if not matches:
        raise FileNotFoundError(f"Ноутбук с префиксом '{notebook_path_prefix}' не найден в '{target_dir}'.")
    elif len(matches) == 1:
        print(f"Найден один совпадающий ноутбук: {matches[0]}")
        return matches[0]
    else:
        # Найдено несколько совпадений (на случай, если заметку создали не из UI: в UI стоит проверка на одинаковые имена)
        error_message = (
            f"Найдено несколько ноутбуков, соответствующих префиксу '{notebook_path_prefix}' в '{target_dir}':\n"
            f"{chr(10).join(matches)}\n" # chr(10) это перевод строки \n
            "Пожалуйста, укажите полный путь к файлу ноутбука, включая его ID и расширение .zpln.\n"
            "ID ноутбука обычно виден в его URL в адресной строке браузера (например, .../notebook/ID_НОУТБУКА)."
        )
        raise ValueError(error_message)

# def import_zeppelin_notebook_from_path(full_notebook_path):
#     """
#     Выполняет Python-код из указанного файла ноутбука Zeppelin
#     в текущем глобальном пространстве имен. Игнорирует параграфы не на Python.

#     Args:
# full_notebook_path (str): Полный путь к файлу .zpln ноутбука.

#     Returns:
#         bool: True если выполнение прошло успешно (или не было Python кода), False при ошибке.
#     """
#     print(f"Пытаюсь прочитать и выполнить код из: {full_notebook_path}")

#     if not os.path.exists(full_notebook_path):
#         print(f"ОШИБКА: Файл ноутбука не найден: {full_notebook_path}")
#         # В новой логике эта проверка дублируется с find_notebook_path,
#         # но оставим на случай прямого вызова этой функции
#         return False

#     try:
#         with open(full_notebook_path, 'r', encoding='utf-8') as f:
#             notebook_data = json.load(f)

#         python_code_to_execute = ""

#         # Проходим по всем параграфам
#         for i, paragraph in enumerate(notebook_data.get("paragraphs", [])):
#             code = paragraph.get("text", "")

#             # Проверяем, является ли параграф Python-параграфом
#             # (%python, %spark.pyspark, %flink.pyflink, %jdbc(python) и т.д.)
#             # Используем более гибкую проверку
#             if code and code.strip().startswith(("%python", "%spark.pyspark", "%flink.pyflink")):
#                 # Убираем строку с магической командой и лишние пробелы/переводы строк
#                 lines = code.splitlines()
#                 if lines:
#                     actual_code = "\n".join(lines[1:]) # Пропускаем первую строку
#                     if actual_code.strip():
#                         print(f"--- Добавляю код из параграфа {i+1} ---")
#                         python_code_to_execute += actual_code + "\n\n" # Добавляем код
#                 else:
#                      print(f"--- Параграф {i+1} пуст (после удаления магической команды) ---")

#             # Можно добавить обработку %jdbc(python) если нужно
#             elif code and code.strip().startswith("%jdbc(python)"):
#                  lines = code.splitlines()
#                  if len(lines) > 1: # Обычно первая строка - конфигурация JDBC
#                      actual_code = "\n".join(lines[1:])
#                      if actual_code.strip():
#                         print(f"--- Добавляю код из параграфа %jdbc(python) {i+1} ---")
#                         python_code_to_execute += actual_code + "\n\n"


#         if python_code_to_execute:
#             print("-" * 20)
#             print("Итоговый Python код для выполнения:")
#             print(python_code_to_execute.strip())
#             print("-" * 20)

#             # Выполняем собранный код в глобальном пространстве имен текущего ноутбука
#             exec(python_code_to_execute, globals())
#             print(f"Выполнение кода из {os.path.basename(full_notebook_path)} завершено.")
#             return True
#         else:
#             print(f"В ноутбуке {os.path.basename(full_notebook_path)} не найдено параграфов с исполняемым Python кодом.")
#             return True # Считаем успешным, т.к. ошибок не было, просто нечего было выполнять

#     except json.JSONDecodeError:
#         print(f"ОШИБКА: Не удалось распарсить JSON файл: {full_notebook_path}")
#         traceback.print_exc()
#         return False
#     except Exception as e:
#         print(f"ОШИБКА: Произошла непредвиденная ошибка при обработке {full_notebook_path}")
#         traceback.print_exc()
#         return False

# # --- Оберточная функция ---
# def import_note(notebook_path_prefix, base_dir=ZEPPELIN_BASE_DIR):
#     """
#     Находит и выполняет Python-код из другого ноутбука Zeppelin,
#     указанного по префиксу пути.

#     Args:
#         notebook_path_prefix (str): Префикс пути к ноутбуку, например "Libraries/main_config".
#         base_dir (str): Базовая директория ноутбуков Zeppelin.

#     Returns:
#         bool: Результат выполнения import_zeppelin_notebook_from_path.
#     """
#     try:
#         full_path = find_notebook_path(notebook_path_prefix, base_dir)
#         return import_zeppelin_notebook_from_path(full_path)
#     except (FileNotFoundError, ValueError) as e:
#         print(f"ОШИБКА ИМПОРТА: {e}")
#         return False
#     except Exception as e:
#         print(f"НЕПРЕДВИДЕННАЯ ОШИБКА при поиске или импорте '{notebook_path_prefix}':")
#         traceback.print_exc()
#         return False

def import_zeppelin_notebook_from_path(full_notebook_path):
    try:
        with open(full_notebook_path, 'r', encoding='utf-8') as f:
            content = f.read()

        python_code_blocks = []
        in_code_block = False
        current_code_block = ""

        for line in content.splitlines():
            if line.startswith("```python"):
                in_code_block = True
                current_code_block = ""
            elif line.startswith("```"):
                if in_code_block:
                    python_code_blocks.append(current_code_block.strip())
                in_code_block = False
            elif in_code_block:
                current_code_block += line + "\n"

        python_code_to_execute = "\n".join(python_code_blocks)

        defined_functions = {}
        if python_code_to_execute:
            local_namespace = {}
            try:
                exec(python_code_to_execute, globals(), local_namespace)
                for name, obj in local_namespace.items():
                    if callable(obj):
                        defined_functions[name] = obj
            except Exception as e:
                print(f"Ошибка при выполнении кода из '{full_notebook_path}': {e}")
                traceback.print_exc()

        return defined_functions

    except FileNotFoundError:
        raise FileNotFoundError(f"Файл не найден: {full_notebook_path}")
    except Exception as e:
        print(f"НЕПРЕДВИДЕННАЯ ОШИБКА при чтении '{full_notebook_path}':")
        traceback.print_exc()
        return {}

def import_note(notebook_path_prefix, base_dir=ZEPPELIN_BASE_DIR):
    try:
        full_path = find_notebook_path(notebook_path_prefix, base_dir)
        return import_zeppelin_notebook_from_path(full_path)
    except (FileNotFoundError, ValueError) as e:
        print(f"ОШИБКА ИМПОРТА: {e}")
        return {}
    except Exception as e:
        print(f"НЕПРЕДВИДЕННАЯ ОШИБКА при поиске или импорте '{notebook_path_prefix}':")
        traceback.print_exc()
        return {}