import os
import warnings
import json
import whisper
import pandas as pd
import numpy as np
from langchain_core.documents import Document

try: import polars as pl
except ImportError: pl = None

try: from sqlalchemy import create_engine
except ImportError: create_engine = None
try: from pymongo import MongoClient
except ImportError: MongoClient = None
try: from neo4j import GraphDatabase
except ImportError: GraphDatabase = None
try: from langchain_community.document_loaders import WebBaseLoader
except ImportError: WebBaseLoader = None
try: from langchain_community.document_loaders import YoutubeLoader
except ImportError: YoutubeLoader = None

try: from rostaing_ocr import ocr_extractor; HAS_OCR = True
except ImportError: HAS_OCR = False

from langchain_community.document_loaders import (
    TextLoader, UnstructuredExcelLoader, 
    UnstructuredPowerPointLoader, UnstructuredWordDocumentLoader, 
    UnstructuredMarkdownLoader, UnstructuredHTMLLoader,
    UnstructuredXMLLoader, UnstructuredEPubLoader
)

warnings.filterwarnings("ignore")

class UniversalLoader:
    def __init__(self, logger=None):
        self.logger = logger
        self._log("Whisper Audio Initialization...")
        try: self.whisper_model = whisper.load_model("base")
        except: self.whisper_model = None

    def _log(self, msg, level="info"):
        if self.logger: self.logger(msg, level)
        else: print(msg)

    # ==========================
    # DATAFRAMES (Pandas & Polars) & SQL
    # ==========================
    def process_dataframe_object(self, df, source_name="InMemory_DataFrame"):
        try:
            # Polars Handling (Safe extraction sans PyArrow forcé)
            if pl and isinstance(df, pl.DataFrame):
                try:
                    schema_dict = dict(df.schema)
                except:
                    schema_dict = str(df.schema)
                summary = f"DATASET: {source_name}\nSchema: {schema_dict}\nRows: {df.height}\nSample:\n{df.head(5)}"
            else:
                # Pandas Handling
                if hasattr(df, "to_pandas"): df = df.to_pandas()
                summary = f"DATASET: {source_name}\nColumns: {list(df.columns)}\nRows: {len(df)}\nSample:\n{df.head(5).to_markdown()}"
                
            return [Document(page_content=summary, metadata={"source": source_name, "type": "dataframe"})]
        except Exception as e:
            self._log(f"DataFrame ingestion Error: {e}", "error")
            return []

    def process_sql(self, config, return_raw=False):
        if not create_engine: return [], None
        uri = config.get("connection_string")
        query = config.get("query", "SELECT * FROM users")
        try:
            engine = create_engine(uri)
            with engine.connect() as connection:
                df = pd.read_sql(query, connection)
                docs = self.process_dataframe_object(df, source_name="sql_query")
                return (docs, df) if return_raw else docs
        except Exception as e:
            self._log(f"SQL Error: {e}", "error")
            return [], None

    # ==========================
    # WEB & YOUTUBE
    # ==========================
    def process_web(self, url):
        if "youtube.com" in url or "youtu.be" in url:
            if not YoutubeLoader: return []
            try:
                loader = YoutubeLoader.from_youtube_url(url, add_video_info=True, language=["fr", "en", "fr-FR"])
                docs = loader.load()
            except:
                try:
                    loader = YoutubeLoader.from_youtube_url(url, add_video_info=False, language=["fr", "en"])
                    docs = loader.load()
                except: return []
            for d in docs: d.metadata.update({"type": "youtube_video", "source": url})
            return docs
        
        if not WebBaseLoader: return []
        try:
            loader = WebBaseLoader(url)
            docs = loader.load()
            for d in docs: d.metadata.update({"source": url, "type": "web"})
            return docs
        except Exception as e: 
            self._log(f"Web Load Error: {e}", "error")
            return []

    def process_mongodb(self, config):
        if not MongoClient: return []
        try:
            client = MongoClient(config["uri"])
            db = client[config["db"]]
            col = db[config["collection"]]
            cursor = col.find().limit(config.get("limit", 100))
            text_content = f"MONGODB SAMPLE ({config['collection']}):\n"
            for doc in cursor: 
                doc.pop("_id", None)
                text_content += json.dumps(doc, default=str) + "\n"
            return [Document(page_content=text_content, metadata={"source": "mongodb"})]
        except: return []

    def process_neo4j(self, config):
        if not GraphDatabase: return []
        try:
            driver = GraphDatabase.driver(config["uri"], auth=(config["user"], config["password"]))
            query = config.get("query", "MATCH (n) RETURN n LIMIT 100")
            
            with driver.session() as session:
                count_res = session.run("MATCH (n) RETURN count(n) as count")
                total_nodes = count_res.single()["count"]
                res = session.run(query)
                results = [json.dumps(record.data(), default=str) for record in res]
            
            driver.close()
            
            text_content = (f"NEO4J GRAPH STATISTICS:\nTotal Nodes: {total_nodes}\n\nRESULTS SAMPLE:\n" + "\n".join(results))
            return [Document(page_content=text_content, metadata={"source": "neo4j", "type": "graph", "uri": config["uri"]})]
        except Exception as e:
            self._log(f"Neo4j Error: {e}", "error")
            return []

    # ==========================
    # FILES
    # ==========================
    def _process_with_rostaing_ocr(self, file_path):
        if not HAS_OCR: return []
        try:
            self._log(f"OCR Processing: {os.path.basename(file_path)}")
            extractor = ocr_extractor(file_path, print_to_console=False, save_file=False)
            if extractor.status == "Success":
                return [Document(page_content=extractor.extracted_text, metadata={"source": file_path, "type": "pdf_ocr"})]
            return []
        except: return []

    def _transcribe_audio(self, p):
        if not self.whisper_model: return ""
        try: return self.whisper_model.transcribe(p)["text"]
        except: return ""

    def load_file(self, file_path, return_raw=False):
        ext = os.path.splitext(file_path)[1].lower()
        
        if ext == ".csv":
            try: 
                df = pd.read_csv(file_path)
                docs = self.process_dataframe_object(df, os.path.basename(file_path))
                return (docs, df) if return_raw else docs
            except: pass
        elif ext == ".parquet":
            try:
                df = pd.read_parquet(file_path)
                docs = self.process_dataframe_object(df, os.path.basename(file_path))
                return (docs, df) if return_raw else docs
            except: pass
        elif ext in [".xlsx", ".xls"]:
            try:
                df = pd.read_excel(file_path)
                docs = self.process_dataframe_object(df, os.path.basename(file_path))
                return (docs, df) if return_raw else docs
            except: pass

        if ext == ".pdf":
            docs = self._process_with_rostaing_ocr(file_path)
            return (docs, None) if return_raw else docs

        if ext in [".jpg", ".png", ".jpeg", ".bmp", ".tiff", ".webp"]:
            return ([Document(page_content="[IMAGE FILE]", metadata={"source": file_path})], None) if return_raw else []

        if ext in [".mp3", ".mp4", ".wav", ".m4a", ".ogg", ".flac", ".webm", ".avi", ".mov", ".mkv"]:
            t = self._transcribe_audio(file_path)
            docs = [Document(page_content=t, metadata={"source": file_path})] if t else []
            return (docs, None) if return_raw else docs
            
        try:
            if ext in [".docx", ".doc"]: loader = UnstructuredWordDocumentLoader(file_path)
            elif ext in [".pptx", ".ppt"]: loader = UnstructuredPowerPointLoader(file_path)
            elif ext in [".html", ".htm"]: loader = UnstructuredHTMLLoader(file_path)
            elif ext == ".xml": loader = UnstructuredXMLLoader(file_path)
            elif ext == ".epub": loader = UnstructuredEPubLoader(file_path)
            elif ext == ".md": loader = UnstructuredMarkdownLoader(file_path)
            else: loader = TextLoader(file_path, encoding="utf-8", autodetect_encoding=True)
            
            docs = loader.load()
            return (docs, None) if return_raw else docs
        except Exception as e: 
            self._log(f"Load Error {file_path}: {e}", "error")
            return ([], None) if return_raw else []

# ------------------------------------------------------------------------
# import os
# import warnings
# import json
# import whisper
# import pandas as pd
# import numpy as np
# from langchain_core.documents import Document

# try: import polars as pl
# except ImportError: pl = None

# try: from sqlalchemy import create_engine
# except ImportError: create_engine = None
# try: from pymongo import MongoClient
# except ImportError: MongoClient = None
# try: from neo4j import GraphDatabase
# except ImportError: GraphDatabase = None
# try: from langchain_community.document_loaders import WebBaseLoader
# except ImportError: WebBaseLoader = None
# try: from langchain_community.document_loaders import YoutubeLoader
# except ImportError: YoutubeLoader = None

# try: from rostaing_ocr import ocr_extractor; HAS_OCR = True
# except ImportError: HAS_OCR = False

# from langchain_community.document_loaders import (
#     TextLoader, UnstructuredExcelLoader, 
#     UnstructuredPowerPointLoader, UnstructuredWordDocumentLoader, 
#     UnstructuredMarkdownLoader, UnstructuredHTMLLoader,
#     UnstructuredXMLLoader, UnstructuredEPubLoader
# )

# warnings.filterwarnings("ignore")

# class UniversalLoader:
#     def __init__(self, logger=None):
#         self.logger = logger
#         self._log("Whisper Audio Initialization...")
#         try: self.whisper_model = whisper.load_model("base")
#         except: self.whisper_model = None

#     def _log(self, msg, level="info"):
#         if self.logger: self.logger(msg, level)
#         else: print(msg)

#     # ==========================
#     # DATAFRAMES (Pandas & Polars) & SQL
#     # ==========================
#     def process_dataframe_object(self, df, source_name="InMemory_DataFrame"):
#         try:
#             # Polars Handling
#             if pl and isinstance(df, pl.DataFrame):
#                 summary = f"DATASET: {source_name}\nSchema: {df.schema}\nRows: {df.height}\nSample:\n{df.head(5).to_pandas().to_markdown()}"
#             else:
#                 # Pandas Handling
#                 if hasattr(df, "to_pandas"): df = df.to_pandas()
#                 summary = f"DATASET: {source_name}\nColumns: {list(df.columns)}\nRows: {len(df)}\nSample:\n{df.head(5).to_markdown()}"
                
#             return [Document(page_content=summary, metadata={"source": source_name, "type": "dataframe"})]
#         except Exception as e:
#             self._log(f"DataFrame ingestion Error: {e}", "error")
#             return []

#     def process_sql(self, config, return_raw=False):
#         if not create_engine: return [], None
#         uri = config.get("connection_string")
#         query = config.get("query", "SELECT * FROM users")
#         try:
#             engine = create_engine(uri)
#             with engine.connect() as connection:
#                 df = pd.read_sql(query, connection)
#                 docs = self.process_dataframe_object(df, source_name="sql_query")
#                 return (docs, df) if return_raw else docs
#         except Exception as e:
#             self._log(f"SQL Error: {e}", "error")
#             return [], None

#     # ==========================
#     # WEB & YOUTUBE
#     # ==========================
#     def process_web(self, url):
#         if "youtube.com" in url or "youtu.be" in url:
#             if not YoutubeLoader: return []
#             try:
#                 loader = YoutubeLoader.from_youtube_url(url, add_video_info=True, language=["fr", "en", "fr-FR"])
#                 docs = loader.load()
#             except:
#                 try:
#                     loader = YoutubeLoader.from_youtube_url(url, add_video_info=False, language=["fr", "en"])
#                     docs = loader.load()
#                 except: return []
#             for d in docs: d.metadata.update({"type": "youtube_video", "source": url})
#             return docs
        
#         if not WebBaseLoader: return []
#         try:
#             loader = WebBaseLoader(url)
#             docs = loader.load()
#             for d in docs: d.metadata.update({"source": url, "type": "web"})
#             return docs
#         except Exception as e: 
#             self._log(f"Web Load Error: {e}", "error")
#             return []

#     def process_mongodb(self, config):
#         if not MongoClient: return []
#         try:
#             client = MongoClient(config["uri"])
#             db = client[config["db"]]
#             col = db[config["collection"]]
#             cursor = col.find().limit(config.get("limit", 100)) # Increased limit for Big Data handling
#             text_content = f"MONGODB SAMPLE ({config['collection']}):\n"
#             for doc in cursor: 
#                 doc.pop("_id", None)
#                 text_content += json.dumps(doc, default=str) + "\n"
#             return [Document(page_content=text_content, metadata={"source": "mongodb"})]
#         except: return []

#     def process_neo4j(self, config):
#         if not GraphDatabase: return []
#         try:
#             driver = GraphDatabase.driver(config["uri"], auth=(config["user"], config["password"]))
#             query = config.get("query", "MATCH (n) RETURN n LIMIT 100")
            
#             with driver.session() as session:
#                 count_res = session.run("MATCH (n) RETURN count(n) as count")
#                 total_nodes = count_res.single()["count"]
#                 res = session.run(query)
#                 results = [json.dumps(record.data(), default=str) for record in res]
            
#             driver.close()
            
#             text_content = (f"NEO4J GRAPH STATISTICS:\nTotal Nodes: {total_nodes}\n\nRESULTS SAMPLE:\n" + "\n".join(results))
#             return [Document(page_content=text_content, metadata={"source": "neo4j", "type": "graph", "uri": config["uri"]})]
#         except Exception as e:
#             self._log(f"Neo4j Error: {e}", "error")
#             return []

#     # ==========================
#     # FILES
#     # ==========================
#     def _process_with_rostaing_ocr(self, file_path):
#         if not HAS_OCR: return []
#         try:
#             self._log(f"OCR Processing: {os.path.basename(file_path)}")
#             extractor = ocr_extractor(file_path, print_to_console=False, save_file=False)
#             if extractor.status == "Success":
#                 return [Document(page_content=extractor.extracted_text, metadata={"source": file_path, "type": "pdf_ocr"})]
#             return []
#         except: return []

#     def _transcribe_audio(self, p):
#         if not self.whisper_model: return ""
#         try: return self.whisper_model.transcribe(p)["text"]
#         except: return ""

#     def load_file(self, file_path, return_raw=False):
#         ext = os.path.splitext(file_path)[1].lower()
        
#         # TABULAR
#         if ext == ".csv":
#             try: 
#                 df = pd.read_csv(file_path)
#                 docs = self.process_dataframe_object(df, os.path.basename(file_path))
#                 return (docs, df) if return_raw else docs
#             except: pass
#         elif ext == ".parquet":
#             try:
#                 df = pd.read_parquet(file_path)
#                 docs = self.process_dataframe_object(df, os.path.basename(file_path))
#                 return (docs, df) if return_raw else docs
#             except: pass
#         elif ext in [".xlsx", ".xls"]:
#             try:
#                 df = pd.read_excel(file_path)
#                 docs = self.process_dataframe_object(df, os.path.basename(file_path))
#                 return (docs, df) if return_raw else docs
#             except: pass

#         # PDF
#         if ext == ".pdf":
#             docs = self._process_with_rostaing_ocr(file_path)
#             return (docs, None) if return_raw else docs

#         # IMAGES
#         if ext in [".jpg", ".png", ".jpeg", ".bmp", ".tiff", ".webp"]:
#             return ([Document(page_content="[IMAGE FILE]", metadata={"source": file_path})], None) if return_raw else []

#         # AUDIO/VIDEO
#         if ext in [".mp3", ".mp4", ".wav", ".m4a", ".ogg", ".flac", ".webm", ".avi", ".mov", ".mkv"]:
#             t = self._transcribe_audio(file_path)
#             docs = [Document(page_content=t, metadata={"source": file_path})] if t else []
#             return (docs, None) if return_raw else docs
            
#         # TEXT/DOCS
#         try:
#             if ext in [".docx", ".doc"]: loader = UnstructuredWordDocumentLoader(file_path)
#             elif ext in [".pptx", ".ppt"]: loader = UnstructuredPowerPointLoader(file_path)
#             elif ext in [".html", ".htm"]: loader = UnstructuredHTMLLoader(file_path)
#             elif ext == ".xml": loader = UnstructuredXMLLoader(file_path)
#             elif ext == ".epub": loader = UnstructuredEPubLoader(file_path)
#             elif ext == ".md": loader = UnstructuredMarkdownLoader(file_path)
#             else: loader = TextLoader(file_path, encoding="utf-8", autodetect_encoding=True)
            
#             docs = loader.load()
#             return (docs, None) if return_raw else docs
#         except Exception as e: 
#             self._log(f"Load Error {file_path}: {e}", "error")
#             return ([], None) if return_raw else []