app.py

import os
import gradio as gr
from datetime import datetime
from utils import extract_kyc_fields
from simple_salesforce import Salesforce
from paddleocr import PaddleOCR

# ------------------------
# CONFIGURATION
# ------------------------
SF_USERNAME = os.getenv("SF_USERNAME", "[email protected]")
SF_PASSWORD = os.getenv("SF_PASSWORD", "Lic@2025@")
SF_TOKEN = os.getenv("SF_TOKEN", "AmmfRcd6IiYaRtSGntBnzNMQU")

# Limit Paddle threads to avoid container crashes
os.environ["OMP_NUM_THREADS"] = "1"

# Initialize OCR
ocr = PaddleOCR(use_angle_cls=True, lang='en')

# ------------------------
# CONNECT TO SALESFORCE
# ------------------------
def connect_salesforce():
    try:
        sf = Salesforce(
            username=SF_USERNAME,
            password=SF_PASSWORD,
            security_token=SF_TOKEN,
            domain="login"  # Production
        )
        print("✅ Connected to Salesforce (Production)")
        return sf
    except Exception as e:
        print("❌ Production login failed:", e)
        try:
            sf = Salesforce(
                username=SF_USERNAME,
                password=SF_PASSWORD,
                security_token=SF_TOKEN,
                domain="test"  # Sandbox
            )
            print("✅ Connected to Salesforce (Sandbox)")
            return sf
        except Exception as e2:
            print("❌ Sandbox login failed:", e2)
            return None

sf = connect_salesforce()

# ------------------------
# CREATE KYC RECORD IN SALESFORCE
# ------------------------
def create_kyc_record(sf, agent_id, kyc_data, file_name=None):
    try:
        if not sf:
            return {"status": "error", "message": "Salesforce connection failed"}

        record = {
            "Agent__c": agent_id,
            "Name": kyc_data.get("name", "Unknown KYC"),
            "Document_Number__c": kyc_data.get("pan_number") or kyc_data.get("aadhaar_number"),
            "Document_Type__c": kyc_data.get("card_type", "UNKNOWN"),
            "Upload_Date__c": datetime.now().isoformat()
        }

        # Optional: store file name / URL if needed
        if file_name:
            record["KYC_Document_URL__c"] = f"Uploaded File: {file_name}"

        response = sf.KYC_Document__c.create(record)
        return {"status": "success", "record_id": response.get("id")}
    except Exception as e:
        return {"status": "error", "message": str(e)}

# ------------------------
# PROCESS IMAGE FUNCTION
# ------------------------
def process_image(image_file, agent_id):
    if not sf:
        return {"status": "error", "message": "Salesforce connection failed"}

    if not image_file:
        return {"status": "error", "message": "No image uploaded"}

    if not agent_id:
        return {"status": "error", "message": "Agent ID is required"}

    try:
        # Extract KYC data using utils.py
        kyc_data = extract_kyc_fields(image_file)
        # Create Salesforce record
        sf_response = create_kyc_record(sf, agent_id, kyc_data, os.path.basename(image_file))
        return {
            "ocr_result": kyc_data,
            "salesforce_result": sf_response
        }
    except Exception as e:
        return {"status": "error", "message": f"OCR/Salesforce error: {str(e)}"}

# ------------------------
# GRADIO INTERFACE
# ------------------------
iface = gr.Interface(
    fn=process_image,
    inputs=[
        gr.Image(type="filepath", label="Upload PAN / Aadhaar"),
        gr.Textbox(label="Agent Salesforce ID", placeholder="Enter Agent record ID")
    ],
    outputs=gr.JSON(),
    title="AI KYC OCR & Salesforce Integration",
    description="Upload PAN or Aadhaar card images to extract text and automatically create KYC records in Salesforce."
)

if __name__ == "__main__":
    iface.launch(server_name="0.0.0.0", server_port=7860)