import gradio as gr
import pandas as pd
from openai import OpenAI
import os
import pdfplumber
import pytesseract
from PIL import Image

# Set your OpenAI API key using environment variables for security
OPENAI_API_KEY = os.getenv('OPEN_API_KEY')
    
client = OpenAI(
    api_key=OPENAI_API_KEY
    # os.getenv('OPENAI_API_KEY'),
)

# Function to extract text from PDFs
def extract_text_from_pdf(file_path):
    text = ""
    try:
        with pdfplumber.open(file_path) as pdf:
            for page in pdf.pages:
                text += page.extract_text() + "\n"
    except Exception as e:
        print(f"Error extracting text from PDF: {str(e)}")
    return text

# Function to extract text from images using OCR
def extract_text_from_image(image):
    try:
        text = pytesseract.image_to_string(image)
        return text
    except Exception as e:
        print(f"Error extracting text from image: {str(e)}")
        return ""

# Function to infer table and column names from extracted text
def infer_table_columns_from_text(text):
    messages = [
        {"role": "system", "content": "You are a helpful assistant that can interpret text to identify table and column names."},
        {"role": "user", "content": f"The following text is extracted from a document: '{text}'. Identify any table names and corresponding column names in the text."}
    ]

    response = client.chat.completions.create(
        model="gpt-4",
        messages=messages,
        max_tokens=300,
        temperature=0.5
    )
    
    return response.choices[0].message.content.strip()

# Function to classify columns using GPT and pre-defined classifications
def classify_columns_from_text(text, classification_file='classifications.xlsx'):
    # Load classification rules
    classification_df = pd.read_excel(classification_file)
    # print("Classification DataFrame Columns:\n", classification_df.columns)
    
    # Infer table and column names using GPT-4
    inferred_structure = infer_table_columns_from_text(text)
    # print("GPT Inferred Structure:\n", inferred_structure)
    
    classified_columns = []
    inferred_lines = inferred_structure.split("\n")
    
    table_name = None

    for line in inferred_lines:
        line = line.strip()
        
        if "typically be found in a" in line and "table" in line:
            table_name = line.split("typically be found in a ")[-1].split(" table")[0].strip()
            print(f"Inferred Table Name: {table_name}")
        
        elif line and line[0].isdigit():
            column_name = line.split(".")[1].strip()
            print(f"Checking for Column: '{column_name}' in Table: '{table_name}'")

            # Find the closest match using GPT-4 to compare the inferred column name with the classification file's columns
            best_match = None
            best_similarity = 0

            for col in classification_df['Column names'].unique():
                prompt_messages = [
                   {"role": "system", "content": "You are a helpful assistant that can compare and identify similarity between two strings."},
                   {"role": "user", "content": f"On a scale of 0 to 1, how similar are these two column names? Column 1: '{column_name}', Column 2: '{col}'. Respond with only the similarity score."}
                ]

                response = client.chat.completions.create(
                    model="gpt-4",
                    messages=prompt_messages,
                    max_tokens=10,
                    temperature=0
                )

                raw_response = response.choices[0].message.content.strip()
                # print("Raw GPT-4 Response:", raw_response)

                try:
                    similarity_score = float(response.choices[0].message.content.strip())
                except ValueError:
                    print("Error: Unable to convert GPT-4 response to a float. Response was:", response.choices[0].message['content'].strip())
                    similarity_score = 0.0  # Default to 0 if conversion fails

                if similarity_score > best_similarity:
                    best_similarity = similarity_score
                    best_match = col

                # print(f"Best match for '{column_name}' is '{best_match}' with similarity {best_similarity}")
    
                if best_match:
                    classification_match = classification_df[classification_df['Column names'].str.strip().str.lower() == best_match.strip().lower()]['Classification']
                    
                    if not classification_match.empty:
                        classification = classification_match.values[0]
                        print(f"Match Found: '{column_name}' classified as '{classification}'")
                    else:
                        classification = "General Use"
                        print(f"No Match Found for Column: '{column_name}' - Defaulting to 'General Use'")
                else:
                    classification = "General Use"
    
                classified_columns.append({
                    'Table Name': table_name,
                    'Column Name': column_name,
                    'Classification': classification
                })

    result_df = pd.DataFrame(classified_columns)
    print("Resulting Classified Columns DataFrame:\n", result_df)
    return result_df

def process_file(file):
    text = ""
    file_extension = file.name.split(".")[-1].lower()
    
    if file_extension in ["pdf"]:
        # Extract text from PDF
        text = extract_text_from_pdf(file)
    elif file_extension in ["jpg", "jpeg", "png", "tiff"]:
        # Extract text from image
        image = Image.open(file)
        text = extract_text_from_image(image)
    else:
        # Assume it's an Excel file
        input_df = pd.read_excel(file)
        text = "\n".join(input_df.apply(lambda row: " ".join(row.values.astype(str)), axis=1))
    
    # Classify columns based on the extracted text
    result_df = classify_columns_from_text(text)
    return result_df

# Create a Gradio interface
iface = gr.Interface(
    fn=process_file,
    inputs=gr.File(label="Upload Excel, PDF, or Image File"),
    outputs=gr.Dataframe(label="Classified Data"),
    title="Data Classification using GPT and Pre-defined Rules",
    description="Upload an Excel, PDF, or Image file to classify its columns based on predefined rules and GPT-4."
)

# Launch the Gradio app with a public link
iface.launch(share=True)