# src/common/notebook_utils.py
import sys
import os
from pathlib import Path
import importlib

VERSION = "1.3.1" # Versión limpia sin exportación CSV

# Variable global para que ingesta_hibrida sepa si debe ser Batch o Stream automáticamente en tests
_GLOBAL_IS_TEST_RUN = False
_SPARK_CACHE = None

def get_checkpoint_base(env_suffix: str = "_dev") -> str:
    """Retorna la ruta base para checkpoints de streaming."""
    base = os.getenv("ELT_CHECKPOINT_BASE")
    if base:
        return base.rstrip("/")
    return f"dbfs:/Volumes/bronze{env_suffix}/temporary/checkpoints"

def get_schema_location_base(env_suffix: str = "_dev") -> str:
    """Retorna la ruta base para schema location de Auto Loader."""
    base = os.getenv("ELT_SCHEMA_BASE")
    if base:
        return base.rstrip("/")
    checkpoint_base = get_checkpoint_base(env_suffix)
    return f"{checkpoint_base}/schema"

# =============================================================================
# FUNCIONES GLOBALES DE PRUEBA Y DX (Developer Experience)
# =============================================================================

def is_test_run_active() -> bool:
    """Retorna si el entorno actual es un testrun interactivo."""
    return _GLOBAL_IS_TEST_RUN

def get_test_spark():
    global _SPARK_CACHE
    if _SPARK_CACHE is not None:
        return _SPARK_CACHE
    
    # 1. Workspace
    try:
        import IPython
        ip = IPython.get_ipython()
        if ip and ip.user_ns.get('spark') is not None:
            _SPARK_CACHE = ip.user_ns['spark']
            return _SPARK_CACHE
    except:
        pass

    # 2. Local (Databricks Connect)
    try:
        from databricks.connect import DatabricksSession
        from databricks.sdk.core import Config
        config = Config(profile="DEFAULT")
        _SPARK_CACHE = DatabricksSession.builder.sdkConfig(config).getOrCreate()
        return _SPARK_CACHE
    except Exception as e:
        print(f" Error creando sesión Spark Local/Connect: {e}")
        # Aquí va tu bloque original de MockSparkSession si falla la conexión
        return None

def stop_local_spark():
    """Limpia streams activos pero MANTIENE la sesión viva."""
    global _SPARK_CACHE
    if _SPARK_CACHE is not None:
        if "DATABRICKS_RUNTIME_VERSION" in os.environ:
            print("ℹ️ Sesión nativa de Databricks. No se detendrá la sesión.")
            return
        try:
            if hasattr(_SPARK_CACHE, "streams") and hasattr(_SPARK_CACHE.streams, "active"):
                active_streams = _SPARK_CACHE.streams.active
                if active_streams:
                    print(f"🛑 Deteniendo {len(active_streams)} proceso(s) de streaming activo(s)...")
                    for stream in active_streams:
                        stream.stop()
                        stream.awaitTermination(5)
        except Exception as e:
            print(f"⚠️ Error deteniendo los streams: {e}")
        print("✅ Recursos de Spark liberados. La sesión sigue activa para la próxima ejecución.")

def get_local_source_table(spark, source_official: str) -> str:
    """Descubre de dónde leer los datos (Temporal vs Unity Catalog oficial)."""
    parts = source_official.split(".")
    env_catalog = parts[0]
    target_table = parts[-1]
    
    oficial_name = source_official
    temp_name = f"{env_catalog}.temporary.{target_table}_tmp_sql"
    
    def table_exists(name):
        try:
            return spark.catalog.tableExists(name)
        except Exception:
            return False

    if table_exists(oficial_name):
        print(f"✅ Leyendo desde tabla oficial Unity Catalog: {oficial_name}")
        return oficial_name
    elif table_exists(temp_name):
        print(f"⚠️ Leyendo desde tabla TEMPORAL de prueba local: {temp_name}")
        return temp_name
    else:
        raise ValueError(
            f"❌ ERROR DE LINAJE LOCAL: No se encontró la tabla de origen.\n"
            f"Solución: Ejecuta el notebook de la capa anterior ({target_table}) para generar los datos."
        )

def clean_local_test_table(spark, source_table: str):
    """Limpia la tabla temporal local tras una prueba."""
    print("-" * 50)
    parts = source_table.split(".")
    env_catalog = parts[0]
    target_table = parts[-1]
    
    temp_name = source_table if target_table.endswith("_tmp_sql") else f"{env_catalog}.temporary.{target_table}_tmp_sql"
    
    def table_exists(name):
        try:
            return spark.catalog.tableExists(name)
        except Exception:
            return False

    if table_exists(temp_name):
        print("🧹 Limpiando zona temporal...")
        try:
            spark.sql(f"DROP TABLE IF EXISTS {temp_name}")
            print(f"✨ Tabla local eliminada limpiamente.")
        except Exception as e:
            print(f"⚠️ Error limpiando la tabla temporal: {e}")

def display_test_results(spark, data, limit: int = 5):
    """Muestra un DataFrame o Tabla adaptándose al entorno."""
    try:
        if isinstance(data, str):
            table_name = data
            df = spark.read.table(table_name).limit(limit)
            print(f"\n✅ Resultado de la tabla: {table_name}")
        else:
            df = data.limit(limit)
            print(f"\n✅ Resultado del DataFrame (En memoria):")

        if "DATABRICKS_RUNTIME_VERSION" in os.environ:
            try:
                import IPython
                shell = IPython.get_ipython()
                if shell is not None and "display" in shell.user_ns:
                    shell.user_ns["display"](df)
                    return
                elif hasattr(df, "display"):
                    df.display()
                    return
            except Exception:
                pass

        df.show(n=limit, truncate=False)
    except Exception as e:
        print(f"⚠️ No se pudo mostrar el resultado: {e}")


# =============================================================================
# CLASE PRINCIPAL: NOTEBOOK CONTEXT (Tu navaja suiza)
# =============================================================================

class NotebookContext:
    def __init__(self, env, is_local, is_test_run=False, spark=None, dlt=None):
        self.env = env
        self.is_local = is_local
        self.is_test_run = is_test_run
        self.spark = spark
        self.dlt = dlt

    # 🛠️ Métodos inyectados para facilitar el uso en los notebooks
    def get_test_spark(self):
        return get_test_spark()

    def get_source_table(self, source_official: str) -> str:
        return get_local_source_table(self.get_test_spark(), source_official)

    def display(self, data, limit: int = 5):
        display_test_results(self.get_test_spark(), data, limit)

    def clean_source(self, source_table: str):
        clean_local_test_table(self.get_test_spark(), source_table)

    def stop_spark(self):
        stop_local_spark()

# =============================================================================
# INICIALIZADOR
# =============================================================================

def init_notebook():
    """Inicializa el entorno del notebook y retorna el NotebookContext."""
    # 1. SETUP PATH DE PROYECTO
    try:
        current_path = Path(__file__).resolve().parent
    except NameError:
        current_path = Path(os.getcwd())

    project_root = current_path
    for _ in range(5):
        if (project_root / "src").exists():
            break
        project_root = project_root.parent

    if str(project_root) not in sys.path:
        sys.path.insert(0, str(project_root))
        print(f"✅ Root agregado al path: {project_root}")

    is_local = False
    is_test_run = False
    spark_session = None
    dlt_module = None
    env_suffix = "_dev" 

    # 2. DETECCION DE SPARK (NUBE/CONNECT)
    try:
        import IPython
        ip = IPython.get_ipython()
        if ip and ip.user_ns.get('spark') is not None:
            spark_session = ip.user_ns['spark']
            try: 
                env_suffix = spark_session.conf.get("env_suffix", "_dev")
            except: 
                pass
    except:
        pass

    # 3. DETECCION DE DLT & TEST MODE
    try:
        import dlt
        dlt_module = dlt
        if spark_session:
            try:
                spark_session.conf.get("pipelines.id")
                is_test_run = False 
            except:
                is_test_run = True 
    except ImportError:
        is_local = True
        is_test_run = True
        
        class DltMock:
            def table(self, *a, **k): return lambda f: f
            def view(self, *a, **k): return lambda f: f
            def expect(self, *a, **k): return lambda f: f
            def expect_or_drop(self, *a, **k): return lambda f: f
            def create_streaming_table(self, *a, **k): pass
            def apply_changes(self, *a, **k): pass
            
        dlt_module = DltMock()
        
    global _GLOBAL_IS_TEST_RUN
    _GLOBAL_IS_TEST_RUN = is_test_run

    if is_local:
        print(f"⚠️  Modo Local: 'dlt' simulado. Usando env='{env_suffix}'")
    elif is_test_run:
        print(f"🏗️  Modo Workspace (TEST): Entorno interactivo detectado. Usando env='{env_suffix}'")
    
    print(f"📦 dbx-elt-utils v{VERSION} inicializado.")

    return NotebookContext(env=env_suffix, is_local=is_local, is_test_run=is_test_run, dlt=dlt_module, spark=spark_session)