src/finee/ui.py

"""
FinEE UI - Gradio Interface
============================

Beautiful web UI for financial entity extraction with:
- Message extraction demo
- Batch processing
- File upload (PDF/Image)
- Analytics dashboard
- Chat interface

Author: Ranjit Behera
"""

import json
from typing import Optional

try:
    import gradio as gr
except ImportError:
    raise ImportError("Please install gradio: pip install gradio")


# ============================================================================
# EXTRACTION FUNCTIONS
# ============================================================================

def extract_entities(
    message: str,
    use_rag: bool = True,
    use_llm: bool = False
) -> tuple:
    """Extract entities from a message."""
    if not message.strip():
        return "{}", "", "Please enter a message"
    
    try:
        # Try to import FinEE
        try:
            from finee import FinancialExtractor
            from finee.rag import RAGEngine
            
            extractor = FinancialExtractor(use_llm=use_llm)
            result = extractor.extract(message)
            
            # RAG enhancement
            rag_info = ""
            if use_rag:
                rag = RAGEngine()
                context = rag.retrieve(message)
                
                if context.merchant_info:
                    if not result.get("merchant"):
                        result["merchant"] = context.merchant_info["name"]
                    if not result.get("category"):
                        result["category"] = context.merchant_info["category"]
                    
                    rag_info = f"""
### RAG Context
- **Merchant**: {context.merchant_info['name']}
- **Category**: {context.merchant_info['category']}
- **Confidence Boost**: +{context.confidence_boost:.0%}
"""
            
            # Format output
            json_output = json.dumps(result, indent=2, ensure_ascii=False)
            
            # Create summary
            summary_parts = []
            if result.get("amount"):
                summary_parts.append(f"💰 **Amount**: ₹{result['amount']:,.2f}")
            if result.get("type"):
                emoji = "📤" if result["type"] == "debit" else "📥"
                summary_parts.append(f"{emoji} **Type**: {result['type'].upper()}")
            if result.get("merchant"):
                summary_parts.append(f"🏪 **Merchant**: {result['merchant']}")
            if result.get("beneficiary"):
                summary_parts.append(f"👤 **Beneficiary**: {result['beneficiary']}")
            if result.get("category"):
                summary_parts.append(f"📂 **Category**: {result['category']}")
            if result.get("bank"):
                summary_parts.append(f"🏦 **Bank**: {result['bank']}")
            if result.get("reference"):
                summary_parts.append(f"🔗 **Reference**: {result['reference']}")
            
            summary = "\n".join(summary_parts) if summary_parts else "No entities extracted"
            summary += rag_info
            
            return json_output, summary, "✅ Extraction successful!"
            
        except ImportError:
            # Fallback mock extraction
            import re
            result = {}
            
            # Amount
            amount_match = re.search(r'Rs\.?\s*([\d,]+(?:\.\d{2})?)', message)
            if amount_match:
                result['amount'] = float(amount_match.group(1).replace(',', ''))
            
            # Type
            if any(w in message.lower() for w in ['debit', 'debited', 'paid']):
                result['type'] = 'debit'
            elif any(w in message.lower() for w in ['credit', 'credited']):
                result['type'] = 'credit'
            
            # Bank
            banks = ['HDFC', 'ICICI', 'SBI', 'Axis', 'Kotak']
            for bank in banks:
                if bank.lower() in message.lower():
                    result['bank'] = bank
                    break
            
            json_output = json.dumps(result, indent=2)
            summary = f"💰 Amount: ₹{result.get('amount', 'N/A')}\n📤 Type: {result.get('type', 'N/A')}"
            
            return json_output, summary, "⚠️ Using mock extractor (install finee for full functionality)"
    
    except Exception as e:
        return "{}", "", f"❌ Error: {str(e)}"


def batch_extract(messages_text: str, use_rag: bool = True) -> str:
    """Extract from multiple messages."""
    if not messages_text.strip():
        return "Please enter messages (one per line)"
    
    messages = [m.strip() for m in messages_text.split('\n') if m.strip()]
    results = []
    
    for i, msg in enumerate(messages, 1):
        json_out, summary, status = extract_entities(msg, use_rag, False)
        try:
            data = json.loads(json_out)
            data['_message'] = msg[:50] + '...' if len(msg) > 50 else msg
            results.append(data)
        except:
            results.append({'error': status, '_message': msg[:50]})
    
    return json.dumps(results, indent=2, ensure_ascii=False)


def analyze_transactions(transactions_json: str) -> str:
    """Analyze transactions and generate insights."""
    try:
        transactions = json.loads(transactions_json)
        
        if not isinstance(transactions, list):
            return "Please provide a list of transactions"
        
        # Calculate stats
        total_debit = sum(t.get('amount', 0) for t in transactions if t.get('type') == 'debit')
        total_credit = sum(t.get('amount', 0) for t in transactions if t.get('type') == 'credit')
        
        # Category breakdown
        categories = {}
        for t in transactions:
            cat = t.get('category', 'other')
            if cat not in categories:
                categories[cat] = {'total': 0, 'count': 0}
            categories[cat]['total'] += t.get('amount', 0)
            categories[cat]['count'] += 1
        
        # Format report
        report = f"""
## 📊 Transaction Analysis

### Summary
- **Total Transactions**: {len(transactions)}
- **Total Debits**: ₹{total_debit:,.2f}
- **Total Credits**: ₹{total_credit:,.2f}
- **Net Flow**: ₹{total_credit - total_debit:,.2f}

### Category Breakdown
"""
        
        sorted_cats = sorted(categories.items(), key=lambda x: x[1]['total'], reverse=True)
        for cat, data in sorted_cats:
            pct = (data['total'] / total_debit * 100) if total_debit > 0 else 0
            report += f"- **{cat.title()}**: ₹{data['total']:,.2f} ({pct:.1f}%) - {data['count']} txns\n"
        
        return report
    
    except json.JSONDecodeError:
        return "❌ Invalid JSON. Please paste valid transaction data."
    except Exception as e:
        return f"❌ Analysis error: {str(e)}"


def chat_response(message: str, history: list) -> str:
    """Handle chat messages."""
    if not message.strip():
        return ""
    
    message_lower = message.lower()
    
    # Intent detection
    if any(w in message_lower for w in ['extract', 'parse', 'analyze this']):
        # Try to extract from the message
        _, summary, _ = extract_entities(message)
        return f"I extracted these entities:\n\n{summary}"
    
    elif any(w in message_lower for w in ['how much', 'spent', 'spending']):
        return "To analyze your spending, please share your transaction messages or paste extracted JSON data."
    
    elif any(w in message_lower for w in ['help', 'what can']):
        return """I can help you with:

1. **Extract Entities** - Paste a bank SMS/email and I'll extract the details
2. **Batch Processing** - Process multiple messages at once
3. **Analyze Spending** - Get insights from your transactions
4. **Categorize** - Understand your spending categories

Try pasting a bank message like:
`HDFC Bank: Rs.2,500 debited from A/c XX1234 on 12-Jan-26. UPI:swiggy@ybl`"""
    
    elif 'hello' in message_lower or 'hi' in message_lower:
        return "Hello! 👋 I'm FinEE, your financial entity extraction assistant. Paste a bank message and I'll extract the details!"
    
    else:
        # Try extraction as default
        json_out, summary, status = extract_entities(message)
        if summary and "No entities" not in summary:
            return f"I found these in your message:\n\n{summary}"
        else:
            return "I'm not sure what you mean. Try pasting a bank SMS or email, or type 'help' for more options."


# ============================================================================
# GRADIO UI
# ============================================================================

def create_ui():
    """Create the Gradio interface."""
    
    # Custom CSS
    custom_css = """
    .gradio-container {
        font-family: 'Inter', sans-serif !important;
    }
    .main-header {
        text-align: center;
        margin-bottom: 20px;
    }
    .output-json {
        font-family: 'Monaco', 'Menlo', monospace !important;
    }
    """
    
    with gr.Blocks(
        title="FinEE - Financial Entity Extractor",
        theme=gr.themes.Soft(
            primary_hue="blue",
            secondary_hue="gray",
        ),
        css=custom_css
    ) as demo:
        
        # Header
        gr.Markdown("""
        # 🏦 FinEE - Financial Entity Extractor
        
        Extract structured data from Indian banking SMS, emails, and statements.
        
        [![GitHub](https://img.shields.io/badge/GitHub-Repo-blue)](https://github.com/Ranjitbehera0034/Finance-Entity-Extractor)
        [![HuggingFace](https://img.shields.io/badge/🤗-Dataset-yellow)](https://huggingface.co/datasets/Ranjit0034/finee-dataset)
        [![PyPI](https://img.shields.io/badge/PyPI-finee-orange)](https://pypi.org/project/finee/)
        """)
        
        with gr.Tabs():
            # Tab 1: Single Extraction
            with gr.Tab("🔍 Extract", id="extract"):
                with gr.Row():
                    with gr.Column(scale=1):
                        input_message = gr.Textbox(
                            label="Bank Message",
                            placeholder="Paste your bank SMS, email, or notification here...",
                            lines=4,
                        )
                        
                        with gr.Row():
                            use_rag = gr.Checkbox(label="Use RAG", value=True, info="Context-aware extraction")
                            use_llm = gr.Checkbox(label="Use LLM", value=False, info="For complex cases")
                        
                        extract_btn = gr.Button("Extract Entities", variant="primary")
                        
                        # Examples
                        gr.Examples(
                            examples=[
                                "HDFC Bank: Rs.2,500 debited from A/c XX1234 on 12-Jan-26. UPI:swiggy@ybl. Ref:123456789012",
                                "SBI: Rs.15,000 credited to A/c XX4567 from rahul.sharma@oksbi. NEFT Ref: N987654321",
                                "ICICI: Your EMI of Rs.12,500 for Loan A/c XX8901 debited on 01-Jan-26",
                                "Axis Bank: Rs.999 debited for Netflix subscription. Card XX5678",
                                "Kotak: Rs.50,000 transferred to Zerodha Broking. Ref: 456789012345",
                            ],
                            inputs=input_message,
                        )
                    
                    with gr.Column(scale=1):
                        status_output = gr.Textbox(label="Status", interactive=False)
                        summary_output = gr.Markdown(label="Summary")
                        json_output = gr.Code(label="JSON Output", language="json")
                
                extract_btn.click(
                    extract_entities,
                    inputs=[input_message, use_rag, use_llm],
                    outputs=[json_output, summary_output, status_output]
                )
            
            # Tab 2: Batch Processing
            with gr.Tab("📋 Batch", id="batch"):
                with gr.Row():
                    with gr.Column():
                        batch_input = gr.Textbox(
                            label="Messages (one per line)",
                            placeholder="Paste multiple messages, one per line...",
                            lines=10,
                        )
                        batch_rag = gr.Checkbox(label="Use RAG", value=True)
                        batch_btn = gr.Button("Process All", variant="primary")
                    
                    with gr.Column():
                        batch_output = gr.Code(label="Results", language="json", lines=20)
                
                batch_btn.click(
                    batch_extract,
                    inputs=[batch_input, batch_rag],
                    outputs=batch_output
                )
            
            # Tab 3: Analytics
            with gr.Tab("📊 Analytics", id="analytics"):
                with gr.Row():
                    with gr.Column():
                        analytics_input = gr.Code(
                            label="Transaction Data (JSON)",
                            language="json",
                            lines=15,
                            value="""[
  {"amount": 2500, "type": "debit", "category": "food", "merchant": "Swiggy"},
  {"amount": 15000, "type": "credit", "category": "transfer"},
  {"amount": 999, "type": "debit", "category": "entertainment", "merchant": "Netflix"},
  {"amount": 5000, "type": "debit", "category": "shopping", "merchant": "Amazon"}
]"""
                        )
                        analyze_btn = gr.Button("Analyze", variant="primary")
                    
                    with gr.Column():
                        analytics_output = gr.Markdown(label="Analysis Report")
                
                analyze_btn.click(
                    analyze_transactions,
                    inputs=analytics_input,
                    outputs=analytics_output
                )
            
            # Tab 4: Chat
            with gr.Tab("💬 Chat", id="chat"):
                chatbot = gr.Chatbot(
                    label="FinEE Assistant",
                    height=400,
                    placeholder="Ask me to extract entities or analyze your transactions..."
                )
                
                with gr.Row():
                    chat_input = gr.Textbox(
                        label="Message",
                        placeholder="Type a message or paste a bank SMS...",
                        scale=4
                    )
                    send_btn = gr.Button("Send", variant="primary", scale=1)
                
                def respond(message, history):
                    if not message.strip():
                        return "", history
                    response = chat_response(message, history)
                    history.append((message, response))
                    return "", history
                
                send_btn.click(respond, [chat_input, chatbot], [chat_input, chatbot])
                chat_input.submit(respond, [chat_input, chatbot], [chat_input, chatbot])
            
            # Tab 5: About
            with gr.Tab("ℹ️ About", id="about"):
                gr.Markdown("""
                ## About FinEE
                
                **FinEE (Financial Entity Extractor)** is a specialized NLP tool for extracting 
                structured information from Indian banking messages.
                
                ### Features
                
                - ✅ **Multi-Bank Support**: HDFC, ICICI, SBI, Axis, Kotak, and 20+ banks
                - ✅ **All Transaction Types**: UPI, NEFT, IMPS, Credit Card, EMI
                - ✅ **Multilingual**: English, Hindi, Tamil, Telugu, Bengali, Kannada
                - ✅ **RAG Enhanced**: Context-aware extraction with merchant knowledge base
                - ✅ **High Accuracy**: 95%+ on standard benchmarks
                
                ### Output Schema
                
                | Field | Type | Description |
                |-------|------|-------------|
                | amount | float | Transaction amount in INR |
                | type | string | "debit" or "credit" |
                | merchant | string | Business name (P2M) |
                | beneficiary | string | Person name (P2P) |
                | category | string | Transaction category |
                | bank | string | Bank name |
                | reference | string | UPI/NEFT reference |
                | vpa | string | UPI VPA address |
                
                ### Links
                
                - 📦 **PyPI**: `pip install finee`
                - 🤗 **Dataset**: [Ranjit0034/finee-dataset](https://huggingface.co/datasets/Ranjit0034/finee-dataset)
                - 💻 **GitHub**: [Ranjitbehera0034/Finance-Entity-Extractor](https://github.com/Ranjitbehera0034/Finance-Entity-Extractor)
                
                ### Author
                
                Built by **Ranjit Behera** | MIT License
                """)
        
        # Footer
        gr.Markdown("""
        ---
        <center>
        Made with ❤️ for the Indian fintech ecosystem
        </center>
        """)
    
    return demo


# ============================================================================
# MAIN
# ============================================================================

def launch(share: bool = False, port: int = 7860):
    """Launch the Gradio app."""
    demo = create_ui()
    demo.launch(share=share, server_port=port)


if __name__ == "__main__":
    import argparse
    
    parser = argparse.ArgumentParser(description="FinEE Gradio UI")
    parser.add_argument("--share", action="store_true", help="Create public link")
    parser.add_argument("--port", type=int, default=7860, help="Port to run on")
    
    args = parser.parse_args()
    launch(share=args.share, port=args.port)