app.py

import gradio as gr
import os
import json
from utils import (
    get_salesforce_client, get_salesforce_objects, get_object_fields,
    extract_text_from_pdf, extract_key_value_pairs,
    create_record, attach_pdf, log_failure
)

from ai_mapper import ai_map_fields  # ✅ NEW: AI-based mapper

# Ensure uploads directory exists
os.makedirs("uploads", exist_ok=True)

# --- Gradio Interface Functions ---

def upload_pdfs(files):
    uploaded_files = []
    if isinstance(files, list):
        for file in files:
            if file is not None:
                file_name = os.path.basename(file.name)
                file_path = f"uploads/{file_name}"
                try:
                    if hasattr(file, 'read'):
                        content = file.read()
                    else:
                        content = str(file).encode('utf-8')
                    with open(file_path, "wb") as f:
                        f.write(content)
                    uploaded_files.append(file_path)
                except Exception as e:
                    return f"Error processing {file_name}: {str(e)}"
    else:
        if files is not None:
            file_name = os.path.basename(files.name)
            file_path = f"uploads/{file_name}"
            try:
                if hasattr(files, 'read'):
                    content = files.read()
                else:
                    content = str(files).encode('utf-8')
                with open(file_path, "wb") as f:
                    f.write(content)
                uploaded_files.append(file_path)
            except Exception as e:
                return f"Error processing {file_name}: {str(e)}"
    return f"Uploaded {len(uploaded_files)} PDF(s): {', '.join(uploaded_files)}"

def fetch_objects():
    sf, error = get_salesforce_client()
    if error:
        return gr.update(choices=[]), f"Error: {error}"
    objects, error = get_salesforce_objects(sf)
    if error:
        return gr.update(choices=[]), f"Error: {error}"
    return gr.update(choices=objects), "Objects fetched successfully"

def fetch_fields(object_name):
    sf, error = get_salesforce_client()
    if error:
        return gr.update(choices=[]), f"Error: {error}"
    fields, error = get_object_fields(sf, object_name)
    if error:
        return gr.update(choices=[]), f"Error: {error}"
    return gr.update(choices=fields), "Fields fetched successfully"

def process_pdf(pdf_paths):
    if not pdf_paths or not isinstance(pdf_paths, list) or not pdf_paths[0]:
        return "Error: No valid PDF file provided"
    pdf_path = pdf_paths[0]
    text_data, error = extract_text_from_pdf(pdf_path)
    if error:
        return f"Error: {error}"
    kv_pairs, error = extract_key_value_pairs(pdf_path)
    if error:
        return f"Error: {error}"
    return f"Text:\n{text_data}\n\nKey-Value Pairs:\n{kv_pairs}"

def display_mappings(pdf_paths, object_name):
    if not pdf_paths or not isinstance(pdf_paths, list) or not pdf_paths[0]:
        return "Error: No valid PDF file provided"
    pdf_path = pdf_paths[0]
    sf, error = get_salesforce_client()
    if error:
        return f"Error: {error}"
    fields, error = get_object_fields(sf, object_name)
    if error:
        return f"Error: {error}"
    extracted_data, error = extract_key_value_pairs(pdf_path)
    if error:
        return f"Error: {error}"
    mappings, confidence_scores, error = ai_map_fields(extracted_data[0]["keys"], fields)  # ✅ updated
    if error:
        return f"Error: {error}"
    output = ""
    for key, field in mappings.items():
        output += f"{key} -> {field} (Confidence: {confidence_scores[key]})\n"
    return output

def migrate_to_salesforce(pdf_paths, object_name):
    if not pdf_paths or not isinstance(pdf_paths, list) or not pdf_paths[0]:
        return "Error: No valid PDF file provided"
    pdf_path = pdf_paths[0]
    sf, error = get_salesforce_client()
    if error:
        log_failure(pdf_path, object_name, error)
        return f"Error: {error}"
    extracted_data, error = extract_key_value_pairs(pdf_path)
    if error:
        log_failure(pdf_path, object_name, error)
        return f"Error: {error}"
    fields, error = get_object_fields(sf, object_name)
    if error:
        log_failure(pdf_path, object_name, error)
        return f"Error: {error}"
    mappings, _, error = ai_map_fields(extracted_data[0]["keys"], fields)  # ✅ updated
    if error:
        log_failure(pdf_path, object_name, error)
        return f"Error: {error}"
    data = {mappings[key]: value for key, value in zip(extracted_data[0]["keys"], extracted_data[0]["values"])}

    # ✅ Required Salesforce IDs added as per Option 2
    data["AccountId"] = "001dL00001ASyPbQAL"
    data["OwnerId"] = "005dL00000f9B0l"

    record_id, error = create_record(sf, object_name, data)
    if error:
        log_failure(pdf_path, object_name, error)
        return f"Error: {error}"
    attach_status, error = attach_pdf(sf, record_id, pdf_path)
    if error:
        log_failure(pdf_path, object_name, error)
        return f"Error: {error}"
    return f"✅ Record Created: {record_id}\n📎 Attachment: {attach_status}"

def display_failures():
    try:
        with open("failures.json", "r") as f:
            failures = [json.loads(line) for line in f]
        output = ""
        for idx, failure in enumerate(failures):
            output += f"Failure {idx + 1}: PDF={failure['pdf']}, Object={failure['object']}, Error={failure['error']}\n"
        return output
    except FileNotFoundError:
        return "No failures logged"

def retry_migration(pdf_path, object_name):
    return migrate_to_salesforce([pdf_path], object_name)

# --- Gradio App Layout ---

with gr.Blocks(css="footer {display: none !important;}") as app:
    gr.Markdown("# 🧠 Smart Contract Migrator")

    with gr.Tab("📄 Upload & Process PDF"):
        with gr.Row():
            pdf_upload = gr.File(label="Upload Contract PDFs", file_count="multiple", file_types=[".pdf"])
            upload_status = gr.Textbox(label="Status", interactive=False)
        pdf_upload.change(upload_pdfs, inputs=pdf_upload, outputs=upload_status)

        with gr.Row():
            process_button = gr.Button("Extract Info")
            process_output = gr.Textbox(label="Text & Key-Value Output", lines=6, interactive=False)
        process_button.click(process_pdf, inputs=pdf_upload, outputs=process_output)

    with gr.Tab("🔁 Salesforce Integration"):
        with gr.Row():
            fetch_objects_button = gr.Button("🔄 Get Objects")
            object_dropdown = gr.Dropdown(label="Salesforce Object")
            object_status = gr.Textbox(label="Status", interactive=False)
        fetch_objects_button.click(fetch_objects, outputs=[object_dropdown, object_status])

        with gr.Row():
            fetch_fields_button = gr.Button("🔄 Get Fields")
            field_checkboxes = gr.CheckboxGroup(label="Fields")
            field_status = gr.Textbox(label="Status", interactive=False)
        fetch_fields_button.click(fetch_fields, inputs=object_dropdown, outputs=[field_checkboxes, field_status])

        with gr.Row():
            map_button = gr.Button("🔗 Map Fields")
            mapping_output = gr.Textbox(label="Mappings", lines=6, interactive=False)
        map_button.click(display_mappings, inputs=[pdf_upload, object_dropdown], outputs=mapping_output)

        with gr.Row():
            migrate_button = gr.Button("🚀 Migrate to Salesforce")
            migrate_output = gr.Textbox(label="Migration Result", lines=4, interactive=False)
        migrate_button.click(migrate_to_salesforce, inputs=[pdf_upload, object_dropdown], outputs=migrate_output)

    with gr.Tab("🧾 Reconciliation"):
        with gr.Row():
            show_failures_button = gr.Button("📋 Show Failures")
            failures_output = gr.Textbox(label="Failures", lines=6, interactive=False)
        show_failures_button.click(display_failures, outputs=failures_output)

        with gr.Row():
            pdf_path_input = gr.Textbox(label="PDF Path to Retry")
            retry_button = gr.Button("🔁 Retry")
            retry_output = gr.Textbox(label="Retry Status", interactive=False)
        retry_button.click(retry_migration, inputs=[pdf_path_input, object_dropdown], outputs=retry_output)

# --- Launch App ---
app.launch()