ds_quickbooks/app/services/ai_matcher.py

import logging
import time
from typing import List, Tuple

import groq

from config import settings
from schemas import Match, Receipt, Transaction

# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class AIMatcher:
    def __init__(self, use_batch_matching=True):
        self.client = groq.Groq(api_key=settings.GROQ_API_KEY)
        self.model = "llama-3.1-8b-instant"
        self.max_retries = 3
        self.retry_delay = 2  # seconds - increased for rate limiting
        self.rate_limit_delay = 1.0  # seconds between API calls
        self.last_api_call = 0
        self.use_batch_matching = (
            use_batch_matching  # Toggle between new and legacy methods
        )

    def match_receipts_to_transactions(
        self, receipts: List[Receipt], transactions: List[Transaction]
    ) -> List[Match]:
        """Match receipts to transactions using AI"""
        logger.info(
            f"Starting AI matching for {len(receipts)} receipts against {len(transactions)} transactions"
        )
        matches = []

        for i, receipt in enumerate(receipts):
            logger.info(
                f"Processing receipt {i + 1}/{len(receipts)}: {receipt.vendor} - ${receipt.amount}"
            )

            # Rate limiting
            self._rate_limit()

            # Get the BEST match for this receipt (highest confidence score)
            best_match = self._find_best_match(receipt, transactions)
            if best_match:
                matches.append(best_match)
                logger.info(
                    f"Found match: {best_match.confidence_score:.3f} - {best_match.match_reason}"
                )
            else:
                logger.warning(
                    f"No match found for receipt: {receipt.vendor} - ${receipt.amount}"
                )

        # Sort by confidence score (highest first)
        matches = sorted(matches, key=lambda x: x.confidence_score, reverse=True)
        logger.info(f"AI matching completed. Found {len(matches)} matches")
        return matches

    def _rate_limit(self):
        """Implement rate limiting to avoid API quota exhaustion"""
        current_time = time.time()
        time_since_last_call = current_time - self.last_api_call

        if time_since_last_call < self.rate_limit_delay:
            sleep_time = self.rate_limit_delay - time_since_last_call
            logger.debug(f"Rate limiting: sleeping for {sleep_time:.2f} seconds")
            time.sleep(sleep_time)

        self.last_api_call = time.time()

    def _find_best_match(
        self, receipt: Receipt, transactions: List[Transaction]
    ) -> Match:
        """Find the BEST match for a receipt using a single AI call for all candidates"""
        candidates = self._filter_candidates(receipt, transactions)
        if not candidates:
            logger.warning(
                f"No candidates found for receipt: {receipt.vendor} - ${receipt.amount}"
            )
            return None

        logger.info(f"Found {len(candidates)} candidates for receipt: {receipt.vendor}")

        # Choose matching method based on configuration
        if self.use_batch_matching:
            # New efficient method: single AI call for all candidates
            best_match = self._find_best_match_single_call(receipt, candidates)
        else:
            # Legacy method: individual AI calls (fallback)
            best_match = self._find_best_match_legacy(receipt, candidates)

        return best_match

    def _find_best_match_single_call(
        self, receipt: Receipt, candidates: List[Transaction]
    ) -> Match:
        """Find the best match using a single AI call to evaluate all candidates"""
        if not candidates:
            return None

        # Limit candidates to avoid token limits (adjust based on your needs)
        max_candidates = 10
        if len(candidates) > max_candidates:
            # Sort by amount similarity and take top candidates
            candidates = sorted(
                candidates, key=lambda t: abs(receipt.amount - abs(t.amount))
            )[:max_candidates]
            logger.info(
                f"Limited candidates to top {max_candidates} by amount similarity"
            )

        # Build comprehensive prompt with all candidates
        candidates_text = ""
        for i, transaction in enumerate(candidates):
            transaction_amount_abs = abs(transaction.amount)
            date_diff = abs((receipt.receipt_date - transaction.transaction_date).days)
            amount_diff = abs(receipt.amount - transaction_amount_abs)
            amount_percent_diff = (
                (amount_diff / receipt.amount) * 100 if receipt.amount > 0 else 0
            )

            candidates_text += f"""
Candidate {i + 1}:
- Vendor: {transaction.vendor}
- Amount: ${transaction.amount} (absolute: ${transaction_amount_abs})
- Date: {transaction.transaction_date.strftime("%Y-%m-%d")} ({date_diff} days difference)
- Notes: {transaction.notes}
- Amount difference: ${amount_diff} ({amount_percent_diff:.1f}%)
"""

        prompt = f"""You are an expert at matching receipts to bank transactions. Analyze the receipt below against ALL the candidate transactions and return the BEST match.

RECEIPT TO MATCH:
- Vendor: {receipt.vendor}
- Amount: ${receipt.amount}
- Date: {receipt.receipt_date.strftime("%Y-%m-%d")}
- Description: {receipt.description}
- Category: {receipt.category}

CANDIDATE TRANSACTIONS:
{candidates_text}

SCORING CRITERIA (Amount is the PRIMARY factor):

Amount Similarity (MOST IMPORTANT - 60% weight):
- Exact match or within 1%: Start at 0.9-1.0
- Within 5%: Start at 0.75-0.89
- Within 10%: Start at 0.5-0.74
- Within 20%: Start at 0.3-0.49
- More than 20% difference: Start at 0.0-0.29

Then adjust UP or DOWN based on:
- Vendor similarity (20% weight): Exact or similar name increases score
- Date proximity (15% weight): Within 7 days increases score, within 30 days moderate increase
- Description/notes match (5% weight): Relevant keywords increase score

EXAMPLES:
- Amount match + vendor match + close date = 0.95-1.0 (Perfect match)
- Amount match + different vendor + close date = 0.85-0.94 (High confidence)
- Amount match + different vendor + far date = 0.70-0.84 (Medium-high confidence)
- Amount similar (5%) + vendor match = 0.75-0.85 (Medium-high confidence)
- Amount similar (10%) + some matches = 0.50-0.69 (Medium confidence)
- Amount very different (>20%) = 0.0-0.29 regardless of other factors

CRITICAL: You MUST return valid JSON only. No explanations, no text before or after.

Return format:
{{"candidate_number": 1, "confidence_score": 0.87, "reason": "Exact amount match with similar vendor"}}

Another example:
{{"candidate_number": 3, "confidence_score": 0.15, "reason": "Poor match but best available"}}

Return ONLY JSON for the best candidate:"""

        logger.info(f"This is the prompt: {prompt}")
        for attempt in range(self.max_retries):
            try:
                result = self._call_groq_api_with_timeout(
                    prompt, timeout=45
                )  # Longer timeout for complex prompt

                # Parse the single result
                candidate_num, score, reason = self._parse_single_match_response(result)

                if candidate_num == -1:  # Parsing error occurred
                    logger.warning(
                        f"Failed to parse AI response for receipt: {receipt.vendor}"
                    )
                    return None

                if 0 <= candidate_num < len(candidates):
                    best_transaction = candidates[candidate_num]
                    logger.info(
                        f"AI selected candidate {candidate_num + 1}: {best_transaction.vendor} (score: {score:.3f})"
                    )
                    return Match(receipt, best_transaction, score, reason)
                else:
                    logger.warning(
                        f"AI returned invalid candidate number: {candidate_num}"
                    )
                    return None

            except Exception as e:
                logger.warning(
                    f"Attempt {attempt + 1} failed for receipt {receipt.id}: {str(e)}"
                )
                if attempt < self.max_retries - 1:
                    sleep_time = self.retry_delay * (2**attempt)
                    logger.info(f"Waiting {sleep_time} seconds before retry...")
                    time.sleep(sleep_time)
                else:
                    logger.error(f"All attempts failed for receipt {receipt.id}")
                    return None

        return None

    def _parse_single_match_response(self, result: str) -> Tuple[int, float, str]:
        """Parse AI response for single best match (JSON format)"""
        import json
        import re
        
        result = result.strip()
        logger.debug(f"Parsing single match response: {result}")

        try:
            # First, try to parse the entire result as JSON
            try:
                data = json.loads(result)
                candidate_num = int(data.get("candidate_number", -1)) - 1
                score = float(data.get("confidence_score", 0.0))
                reason = str(data.get("reason", "No reason provided"))
                score = max(0.0, min(1.0, score))
                logger.debug(f"Parsed JSON: candidate={candidate_num}, score={score}, reason={reason}")
                return candidate_num, score, reason
            except json.JSONDecodeError:
                pass
            
            # Try to extract JSON object from the response using improved regex
            # This handles nested braces better
            json_pattern = r'\{[^{}]*"candidate_number"[^{}]*"confidence_score"[^{}]*"reason"[^{}]*\}'
            json_match = re.search(json_pattern, result)
            
            if json_match:
                json_str = json_match.group()
                data = json.loads(json_str)
                candidate_num = int(data.get("candidate_number", -1)) - 1
                score = float(data.get("confidence_score", 0.0))
                reason = str(data.get("reason", "No reason provided"))
                score = max(0.0, min(1.0, score))
                logger.debug(f"Parsed extracted JSON: candidate={candidate_num}, score={score}, reason={reason}")
                return candidate_num, score, reason
            
            # Try to find any JSON-like structure with the required fields
            candidate_match = re.search(r'"candidate_number"\s*:\s*(\d+)', result)
            score_match = re.search(r'"confidence_score"\s*:\s*([\d.]+)', result)
            reason_match = re.search(r'"reason"\s*:\s*"([^"]*)"', result)
            
            if candidate_match and score_match and reason_match:
                candidate_num = int(candidate_match.group(1)) - 1
                score = float(score_match.group(1))
                reason = reason_match.group(1)
                score = max(0.0, min(1.0, score))
                logger.debug(f"Parsed fields individually: candidate={candidate_num}, score={score}, reason={reason}")
                return candidate_num, score, reason

        except (json.JSONDecodeError, ValueError, KeyError) as e:
            logger.warning(f"Error parsing JSON response: {e}")
            
            # Fallback to old pipe-delimited format for backwards compatibility
            try:
                if "|" in result:
                    parts = result.split("|")
                    if len(parts) >= 3:
                        candidate_str = parts[0].strip()
                        score_str = parts[1].strip()
                        reason = "|".join(parts[2:]).strip()

                        # Extract candidate number
                        candidate_match = re.search(r"\d+", candidate_str)
                        if candidate_match:
                            candidate_num = (
                                int(candidate_match.group()) - 1
                            )  # Convert to 0-based index
                        else:
                            raise ValueError("No candidate number found")

                        # Extract score
                        score_clean = "".join(
                            c for c in score_str if c.isdigit() or c == "."
                        )
                        score = float(score_clean) if score_clean else 0.0

                        # Ensure score is in valid range
                        score = max(0.0, min(1.0, score))

                        logger.debug(
                            f"Parsed (fallback): candidate={candidate_num}, score={score}, reason={reason}"
                        )
                        return candidate_num, score, reason
            except Exception as fallback_error:
                logger.warning(f"Fallback parsing also failed: {fallback_error}")

        # Final fallback
        logger.warning(f"Could not parse single match response: {result}")
        return -1, 0.0, f"Parse error: {result[:50]}..."

    def _filter_candidates(
        self, receipt: Receipt, transactions: List[Transaction]
    ) -> List[Transaction]:
        """Filter transactions to create a reasonable candidate list"""
        candidates = []
        amount_threshold = receipt.amount * 2.0  # 200% threshold - very inclusive

        for transaction in transactions:
            # Use absolute value for transaction amount comparison
            transaction_amount_abs = abs(transaction.amount)

            # Only exclude transactions with obviously different amounts
            if abs(receipt.amount - transaction_amount_abs) <= amount_threshold:
                candidates.append(transaction)

        logger.debug(
            f"Filtered {len(transactions)} transactions to {len(candidates)} candidates"
        )
        return candidates

    def _find_best_match_legacy(
        self, receipt: Receipt, transactions: List[Transaction]
    ) -> Match:
        """Legacy method: Find the best match using individual API calls (kept as fallback)"""
        candidates = self._filter_candidates(receipt, transactions)
        if not candidates:
            return None

        best_match = None
        highest_score = 0

        for transaction in candidates:
            score, reason = self._calculate_match_score(receipt, transaction)
            logger.debug(
                f"Score {score:.3f} for transaction {transaction.vendor}: {reason}"
            )

            if score > highest_score:
                highest_score = score
                best_match = Match(receipt, transaction, score, reason)

        return best_match

    def _calculate_match_score(
        self, receipt: Receipt, transaction: Transaction
    ) -> Tuple[float, str]:
        """Calculate match score using AI"""
        # Calculate differences for the AI to consider
        date_diff = abs((receipt.receipt_date - transaction.transaction_date).days)
        transaction_amount_abs = abs(transaction.amount)
        amount_diff = abs(receipt.amount - transaction_amount_abs)
        amount_percent_diff = (
            (amount_diff / receipt.amount) * 100 if receipt.amount > 0 else 0
        )

        prompt = f"""
        Compare this receipt with this transaction and provide a confidence score (0-1) and brief reason, the reason must be a single sentence without any special formatting.
        
        Receipt: {receipt.vendor}, ${receipt.amount}, {receipt.receipt_date.strftime("%Y-%m-%d")}
        Receipt Description: {receipt.description}
        Receipt Category: {receipt.category}
        Transaction: {transaction.vendor}, ${transaction.amount} (absolute: ${transaction_amount_abs}), {transaction.transaction_date.strftime("%Y-%m-%d")}
        Transaction Notes: {transaction.notes}
        
        Differences:
        - Date difference: {date_diff} days
        - Amount difference: ${amount_diff} ({amount_percent_diff:.1f}%)
        - Vendor comparison: "{receipt.vendor}" vs "{transaction.vendor}"
        - Description/Notes comparison: "{receipt.description}" vs "{transaction.notes}"
        - Category: {receipt.category}
        
        Score this potential match based on how likely it is the correct match:
        
        - Perfect matches (same vendor, amount, date): 0.95-1.0
        - High confidence (minor differences): 0.8-0.94
        - Medium confidence (moderate differences): 0.6-0.79
        - Low confidence (significant differences): 0.4-0.59
        - Very low confidence (major differences): 0.2-0.39
        - Minimal similarity: 0.1-0.19
        - No meaningful similarity: 0.0-0.09
        
        Consider description and category similarity in your scoring.
        
        The most important factor to consider is the Amount for both the transaction and the receipt. The closer the amounts, the higher the score. If the amounts are different or not close return a low score (0-0.1) based on other factors.
        
        IMPORTANT: Return ONLY the score and reason separated by a pipe character.
        Format: [score]|[reason]
        Example: 0.85|Same vendor, same amount, 2 days apart
        """

        for attempt in range(self.max_retries):
            try:
                result = self._call_groq_api_with_timeout(
                    prompt, timeout=30
                )  # Increased timeout

                # Parse the result - handle multiple formats
                score, reason = self._parse_ai_response(result)

                logger.debug(f"AI Response: {result}")
                logger.debug(f"Parsed: score={score}, reason={reason}")

                return score, reason

            except Exception as e:
                logger.warning(
                    f"Attempt {attempt + 1} failed for receipt {receipt.id}: {str(e)}"
                )
                if attempt < self.max_retries - 1:
                    # Exponential backoff for rate limiting
                    sleep_time = self.retry_delay * (2**attempt)
                    logger.info(f"Waiting {sleep_time} seconds before retry...")
                    time.sleep(sleep_time)
                else:
                    logger.error(f"All attempts failed for receipt {receipt.id}")
                    return 0.0, f"AI error after {self.max_retries} attempts: {str(e)}"

    def _parse_ai_response(self, result: str) -> Tuple[float, str]:
        """Parse AI response with robust error handling"""
        result = result.strip()
        logger.debug(f"Parsing AI response: {result}")

        # Try to find score in various formats
        if "|" in result:
            parts = result.split("|")
            logger.debug(f"Split response into {len(parts)} parts: {parts}")

            # Look for a numeric score in any part
            for i, part in enumerate(parts):
                part = part.strip()
                try:
                    # Remove any non-numeric characters except decimal point
                    score_str_clean = "".join(
                        c for c in part if c.isdigit() or c == "."
                    )
                    if score_str_clean:
                        score = float(score_str_clean)
                        if 0 <= score <= 1:  # Valid confidence score
                            # Get reason from other parts
                            reason_parts = [
                                p.strip()
                                for j, p in enumerate(parts)
                                if j != i and p.strip()
                            ]
                            reason = (
                                " | ".join(reason_parts)
                                if reason_parts
                                else "Score extracted"
                            )
                            logger.debug(
                                f"Found score {score} in part {i}, reason: {reason}"
                            )
                            return score, reason
                except ValueError:
                    continue

        # Try to extract just a number from the response
        try:
            import re

            numbers = re.findall(r"\d+\.?\d*", result)
            if numbers:
                for num_str in numbers:
                    score = float(num_str)
                    if 0 <= score <= 1:  # Valid confidence score
                        logger.debug(f"Extracted score {score} from response")
                        return score, f"Extracted from response: {result[:50]}..."
        except (ValueError, IndexError):
            pass

        # Fallback - try to find any number and normalize it
        try:
            import re

            numbers = re.findall(r"\d+\.?\d*", result)
            if numbers:
                score = float(numbers[0])
                # Normalize to 0-1 range if it's a percentage or other scale
                if score > 1:
                    score = score / 100  # Assume percentage
                score = max(0, min(1, score))  # Clamp to 0-1
                logger.debug(f"Normalized score {score} from response")
                return score, f"Normalized from response: {result[:50]}..."
        except (ValueError, IndexError):
            pass

        # Final fallback
        logger.warning(f"Could not parse AI response: {result}")
        return 0.0, f"Unparseable response: {result[:50]}..."

    def _call_groq_api_with_timeout(self, prompt: str, timeout: int = 15) -> str:
        """Make API call with timeout and retry logic"""
        import concurrent.futures

        def api_call():
            try:
                response = self.client.chat.completions.create(
                    model=self.model,
                    messages=[
                        {"role": "system", "content": "You are a JSON-only response assistant. Return only valid JSON, no explanations."},
                        {"role": "user", "content": prompt}
                    ],
                    max_tokens=150,
                    temperature=0.1,
                )
                return response.choices[0].message.content.strip()
            except Exception as e:
                raise e

        try:
            with concurrent.futures.ThreadPoolExecutor() as executor:
                future = executor.submit(api_call)
                return future.result(timeout=timeout)
        except concurrent.futures.TimeoutError:
            raise Exception(f"API call timed out after {timeout} seconds")
        except Exception as e:
            raise e