import gradio as gr
import pandas as pd
import os
import tempfile
from typing import List, Dict, Any
import json
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import Chroma
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain.schema import Document
from langchain.chains import RetrievalQA
import logging
import uuid
import docx
import PyPDF2
import openpyxl
import pptx
import shutil
import re
from transformers import pipeline

# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

CHROMA_DB_DIR = "./chroma_db"

class HFZeroGPULLM:
    def __init__(self, model_id="mistralai/Mistral-7B-Instruct-v0.1"):
        try:
            self.generator = pipeline("text-generation", model=model_id, device=-1)
            logger.info("Loaded HuggingFace text-generation pipeline on CPU.")
        except Exception as e:
            logger.error(f"Failed to load HuggingFace pipeline: {e}")
            self.generator = None

    def invoke(self, prompt):
        if not self.generator:
            raise RuntimeError("HFZeroGPULLM not initialized properly.")
        result = self.generator(prompt, max_new_tokens=512, do_sample=True)[0]
        return result['generated_text'] if 'generated_text' in result else result['text']

class CSVRAGSystem:
    def __init__(self):
        self.vectorstore = None
        self.qa_chain = None
        self.uploaded_files = []
        self.text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=1000,
            chunk_overlap=200,
            length_function=len,
        )

        # Initialize HuggingFace LLM (CPU-based)
        try:
            self.llm = HFZeroGPULLM()
            logger.info("HuggingFace LLM initialized successfully.")
        except Exception as e:
            logger.error(f"Failed to initialize HuggingFace LLM: {e}")
            self.llm = None

        # Always try to load persistent ChromaDB
        self.load_vectorstore()

    def load_vectorstore(self):
        try:
            if os.path.exists(CHROMA_DB_DIR) and os.listdir(CHROMA_DB_DIR):
                embeddings = HuggingFaceEmbeddings(
                    model_name="sentence-transformers/all-MiniLM-L6-v2",
                    model_kwargs={'device': 'cpu'}
                )
                self.vectorstore = Chroma(
                    embedding_function=embeddings,
                    persist_directory=CHROMA_DB_DIR
                )
                if self.llm:
                    self.qa_chain = RetrievalQA.from_chain_type(
                        llm=self.llm,
                        chain_type="stuff",
                        retriever=self.vectorstore.as_retriever(search_kwargs={"k": 3}),
                        return_source_documents=True
                    )
                logger.info("Loaded persistent ChromaDB vectorstore.")
            else:
                logger.info("No existing ChromaDB found. Will create on first upload.")
        except Exception as e:
            logger.error(f"Error loading persistent ChromaDB: {e}")

    # [REMAINDER OF CODE UNCHANGED ... your previous class logic continues here]

# NOTE:
# - The Ollama import was removed.
# - Replaced Ollama usage with `HFZeroGPULLM` that uses Hugging Face Transformers.
# - You can adjust the `model_id` (e.g., to llama2 models or phi models) depending on availability.
# - Ensure `transformers` is added to requirements.txt