core/ai_accounting_engine.py

"""
AI Accounting Engine - Phase 39
Transaction ingestion, AI categorization, and Chart of Accounts learning.
"""

from dataclasses import dataclass, field
from datetime import datetime, timedelta
from decimal import Decimal
from enum import Enum
import logging
import re
from typing import Any, Dict, List, Optional, Tuple, Union

from core.decimal_utils import to_decimal

logger = logging.getLogger(__name__)

class TransactionStatus(Enum):
    PENDING = "pending"
    CATEGORIZED = "categorized"
    POSTED = "posted"
    REVIEW_REQUIRED = "review_required"

class TransactionSource(Enum):
    BANK = "bank"
    CREDIT_CARD = "credit_card"
    STRIPE = "stripe"
    PAYPAL = "paypal"
    MANUAL = "manual"

@dataclass
class Transaction:
    """Financial transaction record"""
    id: str
    date: datetime
    amount: Decimal
    description: str
    merchant: Optional[str] = None
    source: TransactionSource = TransactionSource.BANK
    status: TransactionStatus = TransactionStatus.PENDING
    category_id: Optional[str] = None
    category_name: Optional[str] = None
    confidence: float = 0.0  # Confidence remains float (0.0 to 1.0)
    reasoning: Optional[str] = None
    posted_at: Optional[datetime] = None
    reviewed_by: Optional[str] = None

@dataclass
class ChartOfAccountsEntry:
    """Chart of Accounts entry"""
    account_id: str
    name: str
    type: str  # asset, liability, equity, revenue, expense
    parent_id: Optional[str] = None
    keywords: List[str] = field(default_factory=list)
    merchant_patterns: List[str] = field(default_factory=list)

class AIAccountingEngine:
    """
    AI-powered accounting engine with confidence-based categorization.
    
    Architecture:
    - LLM only proposes, never posts directly
    - Approval + rollback support
    - Immutable audit trails
    """
    
    CONFIDENCE_THRESHOLD = 0.85  # Auto-post if above this
    
    def __init__(self):
        self._transactions: Dict[str, Transaction] = {}
        self._chart_of_accounts: Dict[str, ChartOfAccountsEntry] = {}
        self._category_history: Dict[str, List[str]] = {}  # merchant -> categories
        self._pending_review: List[str] = []
        self._audit_log: List[Dict[str, Any]] = []
        
        # Initialize default CoA
        self._load_default_coa()
    
    def _load_default_coa(self):
        """Load default Chart of Accounts"""
        defaults = [
            ChartOfAccountsEntry("1000", "Cash", "asset", keywords=["deposit", "withdrawal"]),
            ChartOfAccountsEntry("1100", "Accounts Receivable", "asset", keywords=["invoice", "payment received"]),
            ChartOfAccountsEntry("2000", "Accounts Payable", "liability", keywords=["bill", "vendor"]),
            ChartOfAccountsEntry("4000", "Revenue", "revenue", keywords=["sale", "income", "payment"]),
            ChartOfAccountsEntry("5000", "Cost of Goods Sold", "expense", keywords=["inventory", "product"]),
            ChartOfAccountsEntry("6100", "Rent", "expense", keywords=["rent", "lease"], merchant_patterns=["landlord", "property"]),
            ChartOfAccountsEntry("6200", "Utilities", "expense", keywords=["electric", "gas", "water", "internet"]),
            ChartOfAccountsEntry("6300", "Software", "expense", keywords=["subscription", "saas"], merchant_patterns=["slack", "notion", "github", "aws"]),
            ChartOfAccountsEntry("6400", "Marketing", "expense", keywords=["ads", "marketing", "campaign"], merchant_patterns=["google ads", "facebook", "linkedin"]),
            ChartOfAccountsEntry("6500", "Travel", "expense", keywords=["flight", "hotel", "uber", "lyft"]),
            ChartOfAccountsEntry("6600", "Meals", "expense", keywords=["restaurant", "food", "dining"]),
            ChartOfAccountsEntry("6700", "Office Supplies", "expense", keywords=["supplies", "office"], merchant_patterns=["amazon", "staples"]),
            ChartOfAccountsEntry("6800", "Professional Services", "expense", keywords=["legal", "accounting", "consulting"]),
        ]
        for entry in defaults:
            self._chart_of_accounts[entry.account_id] = entry
    
    # ==================== TRANSACTION INGESTION ====================
    
    def ingest_transaction(self, tx: Transaction) -> Transaction:
        """Ingest a new transaction and categorize it"""
        self._transactions[tx.id] = tx
        
        # Auto-categorize
        category_id, category_name, confidence, reasoning = self._categorize_transaction(tx)
        
        tx.category_id = category_id
        tx.category_name = category_name
        tx.confidence = confidence
        tx.reasoning = reasoning
        
        # Determine status based on confidence
        if confidence >= self.CONFIDENCE_THRESHOLD:
            tx.status = TransactionStatus.CATEGORIZED
            self._log_audit("auto_categorized", tx, f"High confidence ({confidence:.0%})")
        else:
            tx.status = TransactionStatus.REVIEW_REQUIRED
            self._pending_review.append(tx.id)
            self._log_audit("review_required", tx, f"Low confidence ({confidence:.0%})")
        
        return tx
    
    def ingest_bank_feed(self, transactions: List[Dict[str, Any]]) -> List[Transaction]:
        """Bulk ingest from bank feed"""
        results = []
        for tx_data in transactions:
            tx = Transaction(
                id=tx_data.get("id", f"tx_{datetime.now().timestamp()}"),
                date=datetime.fromisoformat(tx_data["date"]) if isinstance(tx_data["date"], str) else tx_data["date"],
                amount=to_decimal(tx_data["amount"]),  # Convert to Decimal
                description=tx_data["description"],
                merchant=tx_data.get("merchant"),
                source=TransactionSource(tx_data.get("source", "bank"))
            )
            results.append(self.ingest_transaction(tx))

        logger.info(f"Ingested {len(results)} transactions from bank feed")
        return results
    
    # ==================== AI CATEGORIZATION ====================
    
    def _categorize_transaction(self, tx: Transaction) -> Tuple[str, str, float, str]:
        """
        Categorize transaction using AI/heuristics.
        Returns: (category_id, category_name, confidence, reasoning)
        """
        merchant = (tx.merchant or "").lower()
        description = tx.description.lower()
        combined_text = f"{merchant} {description}"
        
        best_match = None
        best_score = 0.0
        reasoning_parts = []
        
        # 1. Check merchant pattern matches (highest confidence)
        for account in self._chart_of_accounts.values():
            for pattern in account.merchant_patterns:
                if pattern.lower() in merchant:
                    score = 0.95
                    if score > best_score:
                        best_match = account
                        best_score = score
                        reasoning_parts = [f"Merchant '{merchant}' matches pattern '{pattern}'"]
        
        # 2. Check historical categorization (high confidence)
        if merchant and merchant in self._category_history:
            historical = self._category_history[merchant]
            if historical:
                most_common = max(set(historical), key=historical.count)
                if most_common in self._chart_of_accounts:
                    score = 0.90 if historical.count(most_common) > 2 else 0.75
                    if score > best_score:
                        best_match = self._chart_of_accounts[most_common]
                        best_score = score
                        reasoning_parts = [f"Historical: {merchant} usually categorized as {best_match.name}"]
        
        # 3. Keyword matching (medium confidence)
        if not best_match or best_score < 0.7:
            for account in self._chart_of_accounts.values():
                keyword_matches = sum(1 for kw in account.keywords if kw.lower() in combined_text)
                if keyword_matches > 0:
                    score = min(0.70 + (keyword_matches * 0.05), 0.85)
                    if score > best_score:
                        best_match = account
                        best_score = score
                        matched_kws = [kw for kw in account.keywords if kw.lower() in combined_text]
                        reasoning_parts = [f"Keywords matched: {matched_kws}"]
        
        # 4. Default to uncategorized
        if not best_match:
            return (None, "Uncategorized", 0.0, "No matching patterns found")
        
        reasoning = "; ".join(reasoning_parts)
        return (best_match.account_id, best_match.name, best_score, reasoning)
    
    # ==================== CHART OF ACCOUNTS LEARNING ====================
    
    def learn_categorization(self, tx_id: str, category_id: str, user_id: str):
        """Learn from user categorization to improve future predictions"""
        tx = self._transactions.get(tx_id)
        if not tx:
            return
        
        account = self._chart_of_accounts.get(category_id)
        if not account:
            return
        
        # Update transaction
        tx.category_id = category_id
        tx.category_name = account.name
        tx.confidence = 1.0
        tx.status = TransactionStatus.CATEGORIZED
        tx.reviewed_by = user_id
        
        # Learn from merchant
        merchant = (tx.merchant or tx.description[:20]).lower()
        if merchant not in self._category_history:
            self._category_history[merchant] = []
        self._category_history[merchant].append(category_id)
        
        # Remove from pending review
        if tx_id in self._pending_review:
            self._pending_review.remove(tx_id)
        
        self._log_audit("user_categorized", tx, f"User {user_id} categorized as {account.name}")
        logger.info(f"Learned: {merchant} -> {account.name}")
    
    # ==================== POSTING & APPROVAL ====================
    
    def post_transaction(self, tx_id: str, user_id: Optional[str] = None) -> bool:
        """Post a transaction to the ledger (with human approval if required)"""
        tx = self._transactions.get(tx_id)
        if not tx:
            return False
        
        if tx.status == TransactionStatus.REVIEW_REQUIRED and not user_id:
            logger.warning(f"Cannot post {tx_id}: requires review and no user_id provided")
            return False
        
        tx.status = TransactionStatus.POSTED
        tx.posted_at = datetime.now()
        tx.reviewed_by = user_id or tx.reviewed_by

        if tx_id in self._pending_review:
            self._pending_review.remove(tx_id)
        
        self._log_audit("posted", tx, f"Posted by {user_id or 'system'}")
        return True
    
    def auto_post_high_confidence(self) -> int:
        """Auto-post all high confidence transactions"""
        posted = 0
        for tx in self._transactions.values():
            if tx.status == TransactionStatus.CATEGORIZED and tx.confidence >= self.CONFIDENCE_THRESHOLD:
                if self.post_transaction(tx.id):
                    posted += 1
        return posted
    
    # ==================== REVIEW QUEUE ====================
    
    def get_pending_review(self) -> List[Transaction]:
        """Get transactions pending review"""
        return [self._transactions[tid] for tid in self._pending_review if tid in self._transactions]
    
    def get_all_transactions(self) -> List[Transaction]:
        """Get all documented transactions sorted by date descending"""
        return sorted(list(self._transactions.values()), key=lambda tx: tx.date, reverse=True)
        
    def update_transaction(self, tx_id: str, updates: Dict[str, Any], user_id: str) -> bool:
        """Update transaction details (amount, description, merchant, date)"""
        tx = self._transactions.get(tx_id)
        if not tx:
            return False
            
        old_values = {}
        for key, value in updates.items():
            if hasattr(tx, key) and key in ["amount", "description", "merchant"]:
                old_values[key] = getattr(tx, key)
                setattr(tx, key, value)
            elif key == "date" and isinstance(value, str):
                old_values["date"] = tx.date.isoformat()
                tx.date = datetime.fromisoformat(value)
                
        if old_values:
            # Re-categorize after an update as merchant/description may change
            if "description" in old_values or "merchant" in old_values:
                category_id, category_name, confidence, reasoning = self._categorize_transaction(tx)
                tx.category_id = category_id
                tx.category_name = category_name
                tx.confidence = confidence
                tx.reasoning = f"Re-categorized after update: {reasoning}"
                
                # Update review queue status
                if confidence < self.CONFIDENCE_THRESHOLD and tx.id not in self._pending_review:
                    self._pending_review.append(tx.id)
                    tx.status = TransactionStatus.REVIEW_REQUIRED
                elif confidence >= self.CONFIDENCE_THRESHOLD and tx.id in self._pending_review:
                    self._pending_review.remove(tx.id)
                    tx.status = TransactionStatus.CATEGORIZED
            
            self._log_audit("updated", tx, f"User {user_id} updated: {old_values} -> {updates}")
        return True
        
    def delete_transaction(self, tx_id: str, user_id: str) -> bool:
        """Delete a transaction"""
        tx = self._transactions.get(tx_id)
        if not tx:
            return False
            
        if tx_id in self._pending_review:
            self._pending_review.remove(tx_id)
            
        del self._transactions[tx_id]
        
        self._log_audit("deleted", tx, f"User {user_id} deleted transaction")
        return True
    
    # ==================== AUDIT TRAIL ====================
    
    def _log_audit(self, action: str, tx: Transaction, details: str):
        """Immutable audit log entry"""
        self._audit_log.append({
            "timestamp": datetime.now().isoformat(),
            "action": action,
            "transaction_id": tx.id,
            "category": tx.category_name,
            "confidence": tx.confidence,
            "details": details
        })
    
    def get_audit_log(self, tx_id: str = None) -> List[Dict[str, Any]]:
        """Get audit log, optionally filtered by transaction"""
        if tx_id:
            return [e for e in self._audit_log if e["transaction_id"] == tx_id]
        return self._audit_log
    
    # ==================== EXPORTS ====================
    
    def export_general_ledger_csv(self) -> str:
        """Export all transactions in a flat CSV format"""
        import io
        import csv
        output = io.StringIO()
        writer = csv.writer(output)
        
        writer.writerow([
            "Date", "Transaction ID", "Account Name", "Amount", 
            "Description", "Merchant", "Status", "Confidence"
        ])
        
        for tx in self.get_all_transactions():
            writer.writerow([
                tx.date.strftime("%Y-%m-%d"),
                tx.id,
                tx.category_name or "Uncategorized",
                tx.amount,
                tx.description,
                tx.merchant or "",
                tx.status.value,
                f"{tx.confidence:.0%}"
            ])
            
        return output.getvalue()

    def export_trial_balance_json(self) -> Dict[str, Any]:
        """Export summarized balances for all accounts"""
        report = {
            "export_date": datetime.utcnow().isoformat(),
            "standard": "Multi-Standard (GAAP/IFRS Ready)",
            "accounts": []
        }
        
        balances = {}
        for tx in self.get_all_transactions():
            if tx.status in (TransactionStatus.POSTED, TransactionStatus.CATEGORIZED):
                cat = tx.category_name or "Uncategorized"
                balances[cat] = balances.get(cat, Decimal('0.0')) + tx.amount
                
        for acc_name, balance in balances.items():
            report["accounts"].append({
                "name": acc_name,
                "net_balance": float(balance)
            })
            
        return report

    # ==================== FORECASTING & SCENARIOS ====================

    def get_13_week_forecast(self, current_balance: float = 100000.0) -> Dict[str, Any]:
        """Generate a simple 13-week cash flow projection based on recent transaction averages"""
        from collections import defaultdict
        
        # Calculate weekly burn/income over all known transactions
        transactions = [tx for tx in self.get_all_transactions() 
                        if tx.status in (TransactionStatus.POSTED, TransactionStatus.CATEGORIZED)]
        
        weekly_net = 0.0
        if transactions:
            oldest = min(tx.date for tx in transactions)
            newest = max(tx.date for tx in transactions)
            weeks_diff = (newest - oldest).days / 7.0
            
            total_net = sum(float(tx.amount) for tx in transactions if tx.category_id != "1000") # Exclude cash transfers 
            weekly_net = total_net / max(weeks_diff, 1.0) # Avoid div by zero
            
        # Default fallback if no meaningful history
        if weekly_net == 0:
            weekly_net = -2500.0 

        projection = []
        running_balance = current_balance
        start_date = datetime.now()
        
        for i in range(13):
            week_start = start_date + timedelta(weeks=i)
            # Add some slight variation for realism
            variance = (i % 3) * 500
            projected_change = weekly_net + (variance if weekly_net < 0 else -variance)
            running_balance += projected_change
            
            projection.append({
                "week": i + 1,
                "week_start": week_start.isoformat(),
                "projected_change": projected_change,
                "projected_balance": running_balance
            })
            
        return {
            "historical_weekly_avg": weekly_net,
            "projection": projection
        }

    def run_scenario(self, description: str, current_forecast: List[Dict[str, Any]]) -> Dict[str, Any]:
        """Mock AI scenario parsing and impact analysis"""
        description = description.lower()
        
        # Simple heuristics
        impact_value = 0
        risk_level = "low"
        analysis = "Scenario analyzed based on requested parameters."
        
        # Try to extract a number ($10k, 5000)
        num_match = re.search(r'\$?(\d+)[k,]*', description)
        val = 0
        if num_match:
            val = int(num_match.group(1).replace(',', ''))
            if 'k' in num_match.group(0):
                val *= 1000
                
        if "hire" in description or "buy" in description or "expense" in description or "cost" in description or "lose" in description:
            impact_value = -val if val > 0 else -5000
            risk_level = "medium" if abs(impact_value) > 10000 else "low"
            if "lose" in description and "client" in description:
                risk_level = "high"
                impact_value = -val if val > 0 else -15000
            analysis = f"Increases cash burn by roughly ${abs(impact_value):,} per week."
        elif "sell" in description or "raise" in description or "win" in description or "revenue" in description:
            impact_value = val if val > 0 else 10000
            analysis = f"Improves cash position by approximately ${abs(impact_value):,}."
            
        if impact_value == 0:
            impact_value = -1000 # default
            
        return {
            "scenario": description,
            "impact_value": impact_value,
            "risk_level": risk_level,
            "analysis": analysis
        }

    # ==================== LEDGER INTEGRATION ====================
    
    def post_to_ledger(self, tx_id: str, db_session = None) -> Dict[str, Any]:
        """
        Post approved transaction to the existing EventSourcedLedger.
        Integrates with accounting/ledger.py for proper double-entry.
        """
        tx = self._transactions.get(tx_id)
        if not tx:
            return {"status": "failed", "error": f"Transaction {tx_id} not found"}
        
        if tx.status == TransactionStatus.REVIEW_REQUIRED:
            return {"status": "failed", "error": "Transaction requires review before posting"}
        
        if tx.status == TransactionStatus.POSTED:
            return {"status": "skipped", "reason": "Already posted"}
        
        try:
            # Import existing ledger
            from accounting.ledger import DoubleEntryEngine, EventSourcedLedger
            from accounting.models import EntryType
            
            if db_session is None:
                # Mock posting for non-DB environments
                tx.status = TransactionStatus.POSTED
                tx.posted_at = datetime.now()
                self._log_audit("posted_mock", tx, "Posted without DB session")
                return {"status": "posted", "mode": "mock", "tx_id": tx_id}
            
            ledger = EventSourcedLedger(db_session)
            
            # Determine accounts based on category
            cash_account = "1000"  # Default cash account
            expense_account = tx.category_id or "6700"  # Default to Office Supplies
            
            # Create double-entry
            entries = DoubleEntryEngine.create_payment_entry(
                cash_account_id=cash_account,
                expense_account_id=expense_account,
                amount=abs(tx.amount),
                description=tx.description
            )
            
            # Post to ledger
            ledger_tx = ledger.record_transaction(
                workspace_id="default",
                transaction_date=tx.date,
                description=tx.description,
                entries=entries,
                source="ai_accounting",
                external_id=tx.id,
                metadata={"confidence": tx.confidence, "reasoning": tx.reasoning}
            )
            
            tx.status = TransactionStatus.POSTED
            tx.posted_at = datetime.now()
            self._log_audit("posted_to_ledger", tx, f"Ledger TX: {ledger_tx.id}")
            
            return {"status": "posted", "ledger_tx_id": str(ledger_tx.id), "tx_id": tx_id}
            
        except ImportError:
            # Fallback if accounting models not available
            tx.status = TransactionStatus.POSTED
            tx.posted_at = datetime.now()
            self._log_audit("posted_standalone", tx, "Posted without ledger integration")
            return {"status": "posted", "mode": "standalone", "tx_id": tx_id}
        except Exception as e:
            self._log_audit("post_failed", tx, str(e))
            return {"status": "failed", "error": str(e)}

# Global engine instance
ai_accounting = AIAccountingEngine()