import gradio as gr
import openpyxl
import PyPDF2
import pandas as pd
from PIL import Image
import pytesseract  # Replaced EasyOCR
import io
import os
from huggingface_hub import InferenceClient

# Access the Hugging Face token from the environment variable
hf_token = os.environ.get("HF_TOKEN")

def reconcile_statements(erp_file, bank_file):
    yield "⏳ Processing your request... Please wait.", ""

    # your existing code block...

    try:
        # File parsing...
        # Extract ERP statement
        erp_statement = ""
        erp_filename = erp_file.name

        if erp_filename.endswith((".xlsx", ".xls")):
            workbook = openpyxl.load_workbook(erp_filename)
            sheet = workbook.active
            for row in sheet.iter_rows():
                for cell in row:
                    erp_statement += str(cell.value) + "\t"
                erp_statement += "\n"
        elif erp_filename.endswith(".pdf"):
            pdf_reader = PyPDF2.PdfReader(erp_filename)
            for page in pdf_reader.pages:
                erp_statement += page.extract_text() or ""
        elif erp_filename.endswith((".jpg", ".jpeg", ".png")):
            image = Image.open(io.BytesIO(erp_file.read()))
            erp_statement = pytesseract.image_to_string(image)  # Tesseract OCR
        elif erp_filename.endswith(".csv"):
            df = pd.read_csv(erp_filename)
            erp_statement = df.to_string()
        else:
            raise ValueError("Unsupported ERP file format.")

        # Extract bank statement (similar logic as above)
        bank_statement = ""
        bank_filename = bank_file.name

        if bank_filename.endswith((".xlsx", ".xls")):
            workbook = openpyxl.load_workbook(bank_filename)
            sheet = workbook.active
            for row in sheet.iter_rows():
                for cell in row:
                    bank_statement += str(cell.value) + "\t"
                bank_statement += "\n"
        elif bank_filename.endswith(".pdf"):
            pdf_reader = PyPDF2.PdfReader(bank_filename)
            for page in pdf_reader.pages:
                bank_statement += page.extract_text() or ""
        elif bank_filename.endswith((".jpg", ".jpeg", ".png")):
            image = Image.open(io.BytesIO(bank_file.read()))
            bank_statement = pytesseract.image_to_string(image)  # Tesseract OCR
        elif bank_filename.endswith(".csv"):
            df = pd.read_csv(bank_filename)
            bank_statement = df.to_string()
        else:
            raise ValueError("Unsupported bank file format.")

        # Hugging Face request...
        prompt = f"Reconcile these statements:\nERP:\n{erp_statement}\nBank:\n{bank_statement}"

        client = InferenceClient(provider="together", api_key=hf_token)
        completion = client.chat.completions.create(
            model="deepseek-ai/DeepSeek-R1",
            messages=[{"role": "user", "content": prompt}],
        )

        if completion.choices:
            reconciliation_results = completion.choices[0].message.get('content', '')
        else:
            reconciliation_results = "⚠️ No response received from the model."

        output = f"""
        <div style="font-family: 'Segoe UI', ...">
            <h2>🔍 Reconciliation Results</h2>
            <div style="...">
                <pre>{reconciliation_results}</pre>
            </div>
        </div>
        """
        yield "✅ Processing complete!", output

    except Exception as e:
        yield f"❌ Error: {e}", f"<h1>Error</h1><p>{e}</p>"

with gr.Blocks(css="""
#company-logo {
    width: 25%;
    margin: auto;
    display: block;
}
""") as iface:
    gr.Image("logo_Icon.png", elem_id="company-logo", label="Beiing Human")
    status_text = gr.Markdown("👋 Upload your files to begin reconciliation.")
    with gr.Row():
        erp_input = gr.File(label="📂 Upload ERP Statement", type="filepath")
        bank_input = gr.File(label="📂 Upload Bank Statement", type="filepath")
    submit_btn = gr.Button("🔄 Start Reconciliation")
    result_output = gr.HTML()

    submit_btn.click(
        fn=reconcile_statements,
        inputs=[erp_input, bank_input],
        outputs=[status_text, result_output]
    )
    
if __name__ == "__main__":
    iface.launch(debug=True)