app.py

import streamlit as st
import pandas as pd
import pdfplumber
import json
import re
from datetime import datetime
import ollama
import csv
import requests
import io
from PIL import Image
import numpy as np

# Try to import OCR libraries
try:
    import pytesseract
    OCR_AVAILABLE = True
except ImportError:
    OCR_AVAILABLE = False
    pytesseract = None

try:
    import easyocr
    EASYOCR_AVAILABLE = True
except ImportError:
    EASYOCR_AVAILABLE = False
    easyocr = None

# ────────────────────────────────────────────────
#                  CONFIGURATION
# ────────────────────────────────────────────────

MODEL_NAME = "mistral:7b"  # Ollama model name (install with: ollama pull phi3)
OLLAMA_URL = "http://localhost:11434"  # Default Ollama URL
N_CTX = 4096

# Much safer defaults - you can increase later after testing
MAX_NORMALIZE_ROWS = 35
MAX_CATEGORIZE_ROWS = 20  # Reduced to prevent token limit issues - auto-splits if needed

# Rough safety margins (tokens)
PROMPT_SAFETY_MARGIN = 800      # for instructions + JSON overhead + safety
MIN_OUTPUT_TOKENS = 400         # minimum we want to leave for output

# ────────────────────────────────────────────────
#                   OLLAMA CHECK & LOAD
# ────────────────────────────────────────────────

@st.cache_resource
def get_available_models():
    """Get list of available Ollama models"""
    try:
        response = requests.get(f"{OLLAMA_URL}/api/tags", timeout=5)
        if response.status_code != 200:
            return []
        models = response.json().get("models", [])
        # Return full model names (e.g., "mistral:7b", "llama3.2:8b")
        return [m["name"] for m in models]
    except Exception as e:
        return []

@st.cache_resource
def check_ollama():
    """Check if Ollama is running"""
    try:
        # Test Ollama server
        response = requests.get(f"{OLLAMA_URL}/api/tags", timeout=5)
        if response.status_code != 200:
            st.error("Ollama server not running at http://localhost:11434")
            st.info("Run: `ollama serve` in terminal")
            st.stop()
        
        models = response.json().get("models", [])
        if not models:
            st.warning("No models found in Ollama")
            st.info("Install a model: `ollama pull mistral` or `ollama pull llama3.2`")
        
        return True
    except Exception as e:
        st.error(f"Ollama connection failed: {str(e)}")
        st.info("1. Install Ollama: https://ollama.com")
        st.info("2. Terminal: `ollama serve`")
        st.info("3. Terminal: `ollama pull mistral` (or another model)")
        st.stop()

@st.cache_resource
def load_ollama():
    check_ollama()
    # Create Ollama client with explicit host (if available)
    client = None
    try:
        # Try to use Client class if available (newer ollama versions)
        from ollama import Client
        client = Client(host=OLLAMA_URL)
    except (ImportError, AttributeError):
        # Fallback to module-level functions (older ollama versions)
        # The module-level functions use localhost:11434 by default
        pass
    return OLLAMA_URL, client

OLLAMA_URL, ollama_client = load_ollama()

# ────────────────────────────────────────────────
#           TOKEN ESTIMATION & SAFETY
# ────────────────────────────────────────────────

def rough_token_count(text: str) -> int:
    """Very rough but fast token estimation for Phi-3 family"""
    # ~3.8 chars/token average for mixed English + numbers + punctuation
    return len(text) // 3 + len(text.split()) // 2 + 20  # +20 for safety

# ────────────────────────────────────────────────
#        LLM NORMALIZATION (Ollama)
# ────────────────────────────────────────────────

def normalize_with_llm(df_raw: pd.DataFrame, model_name: str = None) -> pd.DataFrame:
    if model_name is None:
        model_name = st.session_state.get('selected_model', MODEL_NAME)
    total = len(df_raw)
    batch_size = MAX_NORMALIZE_ROWS
    chunks = (total + batch_size - 1) // batch_size

    # Quick test to verify model is responding
    try:
        if ollama_client:
            test_response = ollama_client.chat(
                model=model_name,
                messages=[{"role": "user", "content": "Say 'OK'"}],
                options={"num_predict": 10}
            )
        else:
            test_response = ollama.chat(
                model=model_name,
                messages=[{"role": "user", "content": "Say 'OK'"}],
                options={"num_predict": 10}
            )
        if not test_response or 'message' not in test_response or 'content' not in test_response.get('message', {}):
            st.error(f"Model {model_name} is not responding correctly. Please check:")
            st.info(f"1. Model is installed: `ollama pull {model_name.split(':')[0]}`")
            st.info(f"2. Ollama is running: `ollama serve`")
            st.info(f"3. Model name is correct: {model_name}")
            return pd.DataFrame()
    except Exception as e:
        st.error(f"Cannot connect to Ollama model {model_name}: {str(e)}")
        st.info(f"Troubleshooting:")
        st.info(f"1. Ensure Ollama is running: `ollama serve`")
        st.info(f"2. Check model is installed: `ollama list`")
        st.info(f"3. Install model if needed: `ollama pull {model_name.split(':')[0]}`")
        return pd.DataFrame()

    all_results = []
    progress = st.progress(0)

    for i in range(chunks):
        start = i * batch_size
        end = min(start + batch_size, total)
        chunk = df_raw.iloc[start:end].copy()

        csv_text = chunk.to_csv(index=False, lineterminator='\n')

        prompt = f"""You are a precise bank transaction JSON converter.
Output ONLY valid JSON array of objects. Nothing else. No explanations.

Each object MUST contain exactly these 3 fields: "date", "description", "amount"

STRICT RULES:
1. "date": Use "Booking Date" (keep YYYY-MM-DD format)
2. "amount": Use "Amount (EUR)" (keep negative for expenses)
3. "description": 
   - Prefer "Partner Name" → clean lightly (remove only obvious trailing codes like **5400*)
   - Keep full merchant names (e.g. "LIDL BAILEN BARCELONA")
   - If Partner Name empty → use "Type" or "Unknown Transfer"
4. INCLUDE ALL ROWS with date + amount. Do NOT skip:
   - Transfers, savings, payments without partner
   - Large/round amounts, anything that looks real
5. Never output text/markdown outside the JSON array

CSV data:
{csv_text}

Example output (exactly like this):
[{{"date":"2025-12-01","description":"LIDL BAILEN BARCELONA","amount":-16.19}},{{"date":"2025-12-01","description":"Revolut payment","amount":-13.00}}]

Respond with ONLY the JSON array:"""

        prompt_tokens_est = rough_token_count(prompt)

        # Safety check
        remaining_context = N_CTX - prompt_tokens_est - PROMPT_SAFETY_MARGIN
        if remaining_context < MIN_OUTPUT_TOKENS:
            st.warning(f"Batch {i+1} too large (~{prompt_tokens_est} est. tokens). Skipping.")
            continue

        max_tokens = max(MIN_OUTPUT_TOKENS, min(remaining_context, 1400))

        try:
            with st.spinner(f"Normalizing batch {i+1}/{chunks}..."):
                try:
                    if ollama_client:
                        response = ollama_client.chat(
                            model=model_name,
                            messages=[{"role": "user", "content": prompt}],
                            options={
                                "temperature": 0.05,
                                "top_p": 0.9,
                                "top_k": 40,
                                "num_predict": max_tokens,
                                "repeat_penalty": 1.12,
                                "stop": ["```", "\n\n\n", "Note:", "Explanation:", "```json"]
                            }
                        )
                    else:
                        response = ollama.chat(
                            model=model_name,
                            messages=[{"role": "user", "content": prompt}],
                            options={
                                "temperature": 0.05,
                                "top_p": 0.9,
                                "top_k": 40,
                                "num_predict": max_tokens,
                                "repeat_penalty": 1.12,
                                "stop": ["```", "\n\n\n", "Note:", "Explanation:", "```json"]
                            }
                        )
                except Exception as ollama_error:
                    st.error(f"Ollama API call failed: {str(ollama_error)}")
                    st.error(f"Error type: {type(ollama_error).__name__}")
                    with st.expander("Debug: Ollama Error Details", expanded=False):
                        st.code(str(ollama_error), language="text")
                        st.info(f"Model: {model_name}")
                        st.info(f"Ollama URL: {OLLAMA_URL}")
                    raise ValueError(f"Ollama API error: {str(ollama_error)}")
            
            # Debug: Check response structure
            if not response:
                st.error("Ollama returned None/empty response")
                with st.expander("Debug: Response Details", expanded=False):
                    st.json({"response": None, "type": type(response).__name__})
                raise ValueError("Ollama returned empty response")
            
            # Check response structure
            if 'message' not in response:
                st.error("Ollama response missing 'message' key")
                with st.expander("Debug: Response Structure", expanded=False):
                    st.json(response)
                raise ValueError("Invalid Ollama response structure: missing 'message'")
            
            if 'content' not in response['message']:
                st.error("Ollama response missing 'content' key")
                with st.expander("Debug: Response Structure", expanded=False):
                    st.json(response)
                raise ValueError("Invalid Ollama response structure: missing 'content'")
            
            text = response['message']['content']
            
            # Check if text is None or empty
            if text is None:
                st.error("LLM returned None content")
                with st.expander("Debug: Response Details", expanded=False):
                    st.json(response)
                raise ValueError("LLM response content is None")
            
            text = text.strip()

            # Check if response is empty or too short
            if not text or len(text) < 10:
                st.error(f"LLM returned empty or very short response: {repr(text)}")
                with st.expander("Debug: Full Response", expanded=False):
                    st.json(response)
                    st.code(f"Content length: {len(text) if text else 0}\nContent: {repr(text)}", language="text")
                raise ValueError("LLM response is empty or too short")

            # Show raw output for debugging
            st.caption(f"Raw output preview: {repr(text[:200])}...")

            # Try multiple extraction strategies
            json_str = None
            
            # Strategy 1: Look for JSON wrapped in markdown code blocks
            json_match = re.search(r'```(?:json)?\s*(\[.*?\])\s*```', text, re.DOTALL)
            if json_match:
                json_str = json_match.group(1)
            
            # Strategy 2: Find JSON array boundaries
            if not json_str:
                start = text.find('[')
                end = text.rfind(']') + 1
                if start >= 0 and end > start:
                    json_str = text[start:end]
            
            # Strategy 3: Try to salvage partial JSON
            if not json_str:
                # Look for any JSON-like structure
                json_match = re.search(r'\[.*?\]', text, re.DOTALL)
                if json_match:
                    json_str = json_match.group(0)
                else:
                    # Try to find last complete object
                    last_complete = text.rfind('},')
                    if last_complete > 0:
                        json_str = text[:last_complete + 1] + ']'
            
            if not json_str:
                st.error(f"Could not extract JSON from response.")
                with st.expander("Debug: Full LLM Response", expanded=False):
                    st.code(text, language="text")
                raise ValueError("No valid JSON array found in LLM response")
            
            # Clean up the JSON string
            json_str = json_str.strip()
            
            # Try to fix common JSON issues (trailing commas, etc.)
            # Remove trailing commas before closing brackets/braces
            json_str = re.sub(r',(\s*[}\]])', r'\1', json_str)
            
            # Try to parse JSON
            try:
                data = json.loads(json_str)
            except json.JSONDecodeError as e:
                st.error(f"JSON parse error: {str(e)}")
                with st.expander("Debug: JSON Parse Error Details", expanded=False):
                    st.code(f"Error: {str(e)}\n\nAttempted JSON:\n{json_str}", language="text")
                raise ValueError(f"Invalid JSON format: {str(e)}")

            st.write(f"Batch {i+1}: LLM kept {len(data)}/{len(chunk)} rows")
            if data:
                st.json(data[:3])  # Show sample

            df_chunk = pd.DataFrame(data)
            all_results.append(df_chunk)

        except Exception as e:
            st.error(f"Batch {i+1}/{chunks} failed: {str(e)[:100]}")
            continue

        progress.progress((i + 1) / chunks)

    progress.empty()

    if not all_results:
        st.error("All normalization batches failed.")
        return pd.DataFrame()

    df = pd.concat(all_results, ignore_index=True)
    df["date"] = pd.to_datetime(df["date"], errors='coerce')
    df["amount"] = pd.to_numeric(df["amount"], errors='coerce')
    df = df.dropna(subset=["date", "amount"]).sort_values("date").reset_index(drop=True)
    
    st.success(f"Normalized {len(df):,} transactions (dropped {total-len(df)})")
    return df

# ────────────────────────────────────────────────
#     CATEGORIZATION (Ollama)
# ────────────────────────────────────────────────

def categorize_batch(chunk_df: pd.DataFrame, categories: list[str], categories_str: str, 
                     batch_num: int, total_batches: int, model_name: str) -> tuple[pd.DataFrame, Exception, str]:
    """Categorize a single batch. Returns (categorized_df, error, response_text)."""
    csv_text = chunk_df[["orig_idx", "date", "description", "amount"]].to_csv(
        index=False, quoting=csv.QUOTE_NONNUMERIC, lineterminator='\n')
    
    # Convert date to string if needed
    if pd.api.types.is_datetime64_any_dtype(chunk_df["date"]):
        csv_text = chunk_df[["orig_idx", "date", "description", "amount"]].copy()
        csv_text["date"] = chunk_df["date"].dt.strftime("%Y-%m-%d")
        csv_text = csv_text.to_csv(index=False, quoting=csv.QUOTE_NONNUMERIC, lineterminator='\n')

    prompt = f"""Output ONLY a JSON array with objects containing "orig_idx" and "category".
No other text, explanations, or markdown.

Available categories: {categories_str}

RULES:
- Assign EXACTLY ONE category per transaction
- You MUST return a category for EVERY orig_idx in the input data
- Match each transaction to the MOST APPROPRIATE category from the list above
- Use "Uncategorized" ONLY if the transaction truly doesn't fit any category (rare cases)
- Look at description and amount to determine category - be creative and match closely
- orig_idx must match input exactly (0, 1, 2, etc.)

Input data:
{csv_text}

Example output:
[{{"orig_idx":0,"category":"Groceries & Supermarkets"}},{{"orig_idx":1,"category":"Transport & Fuel"}}]

IMPORTANT: Return exactly {len(chunk_df)} objects, one for each row in the input. ONLY the JSON array:"""

    prompt_tokens_est = rough_token_count(prompt)
    remaining = N_CTX - prompt_tokens_est - PROMPT_SAFETY_MARGIN

    if remaining < MIN_OUTPUT_TOKENS:
        chunk_df["category"] = "Uncategorized"
        return chunk_df.drop(columns="orig_idx"), None, ""

    # Calculate max_tokens based on batch size - need more tokens for larger batches
    # Rough estimate: ~80 tokens per transaction in JSON format for output
    # Plus we need tokens for the prompt itself
    estimated_output_tokens = len(chunk_df) * 80
    safe_available = remaining * 0.7  # Use 70% to be safe
    
    # If batch is too large, we'll need to split it
    if estimated_output_tokens > safe_available:
        # Calculate how many rows we can safely process
        max_safe_rows = int(safe_available / 80)
        if max_safe_rows < 1:
            max_safe_rows = 1
        
        st.warning(f"Batch {batch_num} too large ({len(chunk_df)} rows, ~{estimated_output_tokens} tokens needed, {remaining} available)")
        st.info(f"Splitting into smaller chunks. Processing {max_safe_rows} rows at a time...")
        
        # This batch will be processed in chunks - return a special indicator
        # The caller should handle splitting
        raise ValueError(f"BATCH_TOO_LARGE:{max_safe_rows}")

    max_tokens = min(remaining, max(estimated_output_tokens, 1000))

    try:
        if ollama_client:
            response = ollama_client.chat(
                model=model_name,
                messages=[{"role": "user", "content": prompt}],
                options={
                    "temperature": 0.05,
                    "top_p": 0.9,
                    "repeat_penalty": 1.12,
                    "num_predict": max_tokens,
                    "stop": ["```", "\n\n", "Note:", "Explanation:"]
                }
            )
        else:
            response = ollama.chat(
                model=model_name,
                messages=[{"role": "user", "content": prompt}],
                options={
                    "temperature": 0.05,
                    "top_p": 0.9,
                    "repeat_penalty": 1.12,
                    "num_predict": max_tokens,
                    "stop": ["```", "\n\n", "Note:", "Explanation:"]
                }
            )

        text = response['message']['content'].strip()
        
        # Check if response is empty or too short
        if not text or len(text) < 10:
            raise ValueError("LLM response is empty or too short")
        
        # Try multiple extraction strategies
        json_str = None
        
        # Strategy 1: Look for JSON wrapped in markdown code blocks (non-greedy)
        json_match = re.search(r'```(?:json)?\s*(\[.*?\])\s*```', text, re.DOTALL)
        if json_match:
            json_str = json_match.group(1)
        
        # Strategy 2: Look for JSON wrapped in markdown code blocks (greedy - in case non-greedy fails)
        if not json_str:
            json_match = re.search(r'```(?:json)?\s*(\[.*\])\s*```', text, re.DOTALL)
            if json_match:
                json_str = json_match.group(1)
        
        # Strategy 3: Try to find JSON array with balanced brackets (more robust - try this first)
        if not json_str:
            # Find first '[' and then find matching ']' by counting brackets
            start = text.find('[')
            if start >= 0:
                bracket_count = 0
                for i in range(start, len(text)):
                    if text[i] == '[':
                        bracket_count += 1
                    elif text[i] == ']':
                        bracket_count -= 1
                        if bracket_count == 0:
                            json_str = text[start:i+1]
                            break
        
        # Strategy 4: Find JSON array boundaries (simple approach - fallback)
        if not json_str:
            start = text.find('[')
            end = text.rfind(']') + 1
            if start >= 0 and end > start:
                json_str = text[start:end]
        
        # Strategy 3.5: If we found opening bracket but no closing, try to complete it
        if not json_str:
            start = text.find('[')
            if start >= 0:
                # Check if there are any complete objects before the end
                # Look for patterns like }, or }]
                last_complete = max(text.rfind('},'), text.rfind('}]'))
                if last_complete > start:
                    # We have a start and at least one complete object
                    # Try to salvage by closing the array
                    # Find the last complete object's closing brace
                    if text.rfind('},') > text.rfind('}]'):
                        # Last thing is }, so close with ]
                        potential_json = text[start:text.rfind('},') + 1] + ']'
                    else:
                        # Already has }] somewhere
                        potential_json = text[start:text.rfind('}]') + 2]
                    # Validate it looks like JSON before accepting
                    if potential_json.count('{') == potential_json.count('}') and 'orig_idx' in potential_json:
                        json_str = potential_json
        
        # Strategy 5: Try regex with greedy matching (in case array spans multiple lines)
        if not json_str:
            json_match = re.search(r'\[.*\]', text, re.DOTALL)
            if json_match:
                json_str = json_match.group(0)
        
        # Strategy 6: Look for JSON-like structure even if malformed
        if not json_str:
            # Try to find something that looks like JSON array
            json_match = re.search(r'\[[^\]]*\{[^}]*"[^"]*"[^}]*\}[^\]]*\]', text, re.DOTALL)
            if json_match:
                json_str = json_match.group(0)
        
        if not json_str:
            # Last resort: If we see a '[' but extraction failed, try to salvage
            start = text.find('[')
            if start >= 0:
                # Check if we can see JSON-like structure starting
                if '"orig_idx"' in text or '"category"' in text:
                    # Looks like JSON but extraction failed - maybe response was truncated
                    # Try to get as much as possible and add closing bracket
                    # Find last complete item pattern
                    potential_end = max(
                        text.rfind('},'),
                        text.rfind('}'),
                        text.rfind('"category"') + 50  # rough estimate after category field
                    )
                    if potential_end > start:
                        # Try to create valid JSON by closing properly
                        salvage_json = text[start:potential_end].rstrip()
                        # Remove trailing comma if present
                        salvage_json = salvage_json.rstrip(',')
                        # Ensure it ends with }
                        if not salvage_json.endswith('}'):
                            # Find last }
                            last_brace = salvage_json.rfind('}')
                            if last_brace > 0:
                                salvage_json = salvage_json[:last_brace + 1]
                        # Add closing bracket
                        salvage_json += ']'
                        # Quick sanity check
                        if salvage_json.count('{') <= salvage_json.count('}') and '{' in salvage_json:
                            json_str = salvage_json
            
            # If still no JSON found, raise error with helpful message
            if not json_str:
                # Store error details for debugging
                has_opening = '[' in text
                has_closing = ']' in text
                preview = repr(text[:500]) if len(text) > 500 else repr(text)
                error_msg = f"No valid JSON array found in LLM response. Has '[': {has_opening}, Has ']': {has_closing}. Response preview: {preview}"
                raise ValueError(error_msg)
        
        # Clean up JSON (remove trailing commas, fix common issues)
        json_str = json_str.strip()
        json_str = re.sub(r',(\s*[}\]])', r'\1', json_str)
        
        # Try to parse JSON, with error handling for common issues
        try:
            data = json.loads(json_str)
        except json.JSONDecodeError as json_err:
            # If error mentions unterminated string, try to fix it
            if "Unterminated string" in str(json_err) or "Expecting" in str(json_err):
                # Try to find last complete object (ending with '},')
                # rfind('},') returns the start position of '},' which is the position of '}'
                last_complete_obj = json_str.rfind('},')
                if last_complete_obj > 0:
                    # Include the '}' at position last_complete_obj, add ']' to close array
                    # This removes any incomplete last object
                    fixed_json = json_str[:last_complete_obj + 1] + ']'
                    try:
                        data = json.loads(fixed_json)
                    except json.JSONDecodeError as e2:
                        # If that still fails, try to find the last properly closed quote before the error position
                        # The error message usually contains the character position
                        error_pos_match = re.search(r'char (\d+)', str(json_err))
                        if error_pos_match:
                            error_pos = int(error_pos_match.group(1))
                            # Try truncating just before the error position
                            fixed_json = json_str[:error_pos].rstrip().rstrip(',') + ']'
                            try:
                                data = json.loads(fixed_json)
                            except json.JSONDecodeError:
                                raise ValueError(f"JSON parse error: {str(json_err)}. Could not auto-fix.")
                        else:
                            raise ValueError(f"JSON parse error: {str(json_err)}. Could not auto-fix.")
                else:
                    raise ValueError(f"JSON parse error: {str(json_err)}")
            else:
                raise ValueError(f"JSON parse error: {str(json_err)}")
        
        cat_map = {item["orig_idx"]: item["category"] for item in data}
        
        # Validate all orig_idx values are covered
        expected_indices = set(chunk_df["orig_idx"].tolist())
        received_indices = set(cat_map.keys())
        missing_indices = expected_indices - received_indices
        
        if missing_indices:
            # Warn about missing mappings - likely due to truncation
            st.warning(f"LLM did not return categories for {len(missing_indices)} row(s) (orig_idx: {sorted(missing_indices)})")
            
            # Try to infer if response was truncated
            is_truncated = False
            if text and not text.rstrip().endswith(']'):
                is_truncated = True
                st.warning(f"Response appears truncated - doesn't end with ']'. Last 50 chars: {repr(text[-50:])}")
            
            # Try to categorize missing rows with a smaller batch
            if len(missing_indices) < len(chunk_df) and len(missing_indices) <= 10:
                st.info(f"Attempting to categorize {len(missing_indices)} missing row(s) separately...")
                try:
                    # Get the missing rows (keep orig_idx for mapping)
                    missing_rows = chunk_df[chunk_df["orig_idx"].isin(missing_indices)].copy()
                    # Retry with just these rows (smaller batch should work)
                    retry_result, retry_error, _ = categorize_batch(
                        missing_rows, categories, categories_str, 
                        batch_num, 1, model_name
                    )
                    if not retry_error:
                        # The retry_result has orig_idx removed, so we match by position
                        # missing_rows and retry_result are in the same order
                        for i, orig_idx_val in enumerate(missing_rows["orig_idx"].values):
                            if i < len(retry_result):
                                cat_map[orig_idx_val] = retry_result.iloc[i]["category"]
                                st.success(f"Successfully categorized orig_idx {orig_idx_val}")
                        # Re-map with updated cat_map
                        chunk_df["category"] = chunk_df["orig_idx"].map(cat_map)
                except Exception as retry_e:
                    st.warning(f"Retry for missing rows failed: {str(retry_e)[:100]}")
            
            if missing_indices:
                st.info(f"{len(missing_indices)} row(s) still missing. Response length: {len(text)} chars. Consider reducing batch size from {MAX_CATEGORIZE_ROWS}.")
        
        # Check for duplicate orig_idx in LLM response
        if len(data) != len(cat_map):
            st.warning(f"LLM returned duplicate orig_idx values. Some categories may be overwritten.")
        
        # Map categories, but track which ones are missing
        chunk_df["category"] = chunk_df["orig_idx"].map(cat_map)
        
        # Only fill missing with "Uncategorized" if LLM didn't provide them
        # This preserves the distinction between "LLM said Uncategorized" vs "LLM didn't provide category"
        uncategorized_count = (chunk_df["category"] == "Uncategorized").sum() if "Uncategorized" in cat_map.values() else 0
        missing_count = chunk_df["category"].isna().sum()
        
        if missing_count > 0:
            st.warning(f"{missing_count} row(s) missing from LLM response - assigning 'Uncategorized'")
            chunk_df["category"] = chunk_df["category"].fillna("Uncategorized")
        elif uncategorized_count > 0:
            st.info(f"LLM assigned 'Uncategorized' to {uncategorized_count} row(s) - check if this is correct")
        
        return chunk_df.drop(columns="orig_idx"), None, text

    except Exception as e:
        chunk_df["category"] = "Uncategorized"
        # Return error with response text for debugging
        response_text = text if 'text' in locals() else ""
        return chunk_df.drop(columns="orig_idx"), e, response_text

def categorize_with_llm(df: pd.DataFrame, categories: list[str], model_name: str = None, 
                        track_failures: bool = True) -> tuple[pd.DataFrame, dict]:
    """Categorize transactions. Returns (categorized_df, failed_batches_dict)."""
    if model_name is None:
        model_name = st.session_state.get('selected_model', MODEL_NAME)
    if df.empty:
        return df, {}

    categories_str = " • ".join(categories)

    batch_size = MAX_CATEGORIZE_ROWS
    total = len(df)
    chunks = (total + batch_size - 1) // batch_size

    progress = st.progress(0)
    status = st.empty()

    categorized_chunks = []
    failed_batches = {}  # {batch_num: {"chunk_df": ..., "error": ...}}

    for i in range(chunks):
        start = i * batch_size
        end = min(start + batch_size, total)
        chunk_df = df.iloc[start:end].copy().reset_index(names="orig_idx")

        status.text(f"Categorizing batch {i+1}/{chunks}...")
        categorized_chunk, error, response_text = categorize_batch(chunk_df, categories, categories_str, i+1, chunks, model_name)
        
        if error:
            # Check if error is due to batch being too large
            if isinstance(error, ValueError) and str(error).startswith("BATCH_TOO_LARGE:"):
                # Extract the recommended batch size
                try:
                    recommended_size = int(str(error).split(":")[1])
                    st.info(f"📦 Batch {i+1} too large, splitting into smaller chunks of {recommended_size} rows...")
                    
                    # Split the chunk into smaller sub-chunks
                    sub_chunk_size = recommended_size
                    sub_chunks = (len(chunk_df) + sub_chunk_size - 1) // sub_chunk_size
                    
                    sub_categorized_chunks = []
                    for sub_i in range(sub_chunks):
                        sub_start = sub_i * sub_chunk_size
                        sub_end = min(sub_start + sub_chunk_size, len(chunk_df))
                        sub_chunk_df = chunk_df.iloc[sub_start:sub_end].copy()
                        # Preserve original orig_idx - don't reset it, just ensure it's there
                        if "orig_idx" not in sub_chunk_df.columns:
                            sub_chunk_df = sub_chunk_df.reset_index(names="orig_idx")
                        
                        status.text(f"Categorizing batch {i+1}/{chunks}, sub-chunk {sub_i+1}/{sub_chunks}...")
                        sub_result, sub_error, sub_response = categorize_batch(
                            sub_chunk_df, categories, categories_str, 
                            f"{i+1}.{sub_i+1}", sub_chunks, model_name
                        )
                        
                        if sub_error:
                            st.warning(f"Sub-chunk {sub_i+1} of batch {i+1} failed: {str(sub_error)[:100]}")
                            if track_failures:
                                failed_batches[f"{i+1}.{sub_i+1}"] = {
                                    "chunk_df": sub_chunk_df,
                                    "error": str(sub_error),
                                    "categories": categories,
                                    "categories_str": categories_str,
                                    "model_name": model_name,
                                    "response_text": sub_response
                                }
                        else:
                            st.success(f"Sub-chunk {sub_i+1}/{sub_chunks} of batch {i+1} categorized")
                        
                        sub_categorized_chunks.append(sub_result)
                    
                    # Combine sub-chunks
                    if sub_categorized_chunks:
                        categorized_chunk = pd.concat(sub_categorized_chunks, ignore_index=True)
                        st.success(f"Batch {i+1}/{chunks} completed ({sub_chunks} sub-chunks)")
                    else:
                        # All sub-chunks failed
                        categorized_chunk = chunk_df.copy()
                        categorized_chunk["category"] = "Uncategorized"
                        categorized_chunk = categorized_chunk.drop(columns="orig_idx")
                except (ValueError, IndexError):
                    # Couldn't parse the error, treat as regular error
                    st.warning(f"Batch {i+1} failed → 'Uncategorized': {str(error)[:100]}")
                    if track_failures:
                        failed_batches[i+1] = {
                            "chunk_df": chunk_df,
                            "error": str(error),
                            "categories": categories,
                            "categories_str": categories_str,
                            "model_name": model_name,
                            "response_text": response_text
                        }
            else:
                # Regular error
                st.warning(f"⚠️ Batch {i+1} failed → 'Uncategorized': {str(error)[:100]}")
                if track_failures:
                    failed_batches[i+1] = {
                        "chunk_df": chunk_df,
                        "error": str(error),
                        "categories": categories,
                        "categories_str": categories_str,
                        "model_name": model_name,
                        "response_text": response_text
                    }
        else:
            st.success(f"Batch {i+1}/{chunks} categorized successfully")
        
        categorized_chunks.append(categorized_chunk)

        progress.progress((i + 1) / chunks)

    progress.empty()
    status.empty()

    result_df = pd.concat(categorized_chunks, ignore_index=True)
    return result_df, failed_batches

# ────────────────────────────────────────────────
#        RECEIPT EXTRACTION (Ollama)
# ────────────────────────────────────────────────

def validate_image_quality(image_bytes: bytes, file_name: str) -> tuple[bool, str]:
    """
    Validate image quality for OCR processing.
    
    Returns:
        (is_valid, error_message)
        is_valid: True if image quality is acceptable
        error_message: Error message if quality is poor, empty string if valid
    """
    try:
        # Open image to check dimensions
        image = Image.open(io.BytesIO(image_bytes))
        width, height = image.size
        
        # Check minimum resolution
        MIN_WIDTH = 300
        MIN_HEIGHT = 300
        MIN_TOTAL_PIXELS = 100000  # ~316x316 pixels
        
        total_pixels = width * height
        
        errors = []
        
        # Check width
        if width < MIN_WIDTH:
            errors.append(f"width too small ({width}px, minimum {MIN_WIDTH}px)")
        
        # Check height
        if height < MIN_HEIGHT:
            errors.append(f"height too small ({height}px, minimum {MIN_HEIGHT}px)")
        
        # Check total pixels (catches very small images)
        if total_pixels < MIN_TOTAL_PIXELS:
            errors.append(f"image too small ({total_pixels:,} pixels, minimum {MIN_TOTAL_PIXELS:,})")
        
        # Check file size (very small files might be low quality)
        file_size_kb = len(image_bytes) / 1024
        if file_size_kb < 10:  # Less than 10KB is suspicious
            errors.append(f"file size too small ({file_size_kb:.1f}KB, may indicate low quality)")
        
        # Check aspect ratio (extremely wide or tall images might be problematic)
        aspect_ratio = width / height if height > 0 else 0
        if aspect_ratio > 10 or aspect_ratio < 0.1:
            errors.append(f"extreme aspect ratio ({aspect_ratio:.2f}, may affect OCR accuracy)")
        
        if errors:
            error_msg = f"**{file_name}** has poor quality: {', '.join(errors)}"
            error_msg += f"\n\n**Current dimensions:** {width}×{height} pixels ({total_pixels:,} total)"
            error_msg += f"\n**File size:** {file_size_kb:.1f} KB"
            error_msg += "\n\n**Recommendations:**"
            error_msg += f"\n- Minimum resolution: {MIN_WIDTH}×{MIN_HEIGHT} pixels"
            error_msg += "\n- Use a well-lit, in-focus photo"
            error_msg += "\n- Ensure text is clearly visible and not blurry"
            error_msg += "\n- Try taking the photo again with better lighting"
            return False, error_msg
        
        return True, ""
        
    except Exception as e:
        return False, f"Could not validate image {file_name}: {str(e)}"

def extract_text_from_image(image_bytes: bytes, file_name: str) -> str:
    """
    Extract text from an image using OCR.
    Tries pytesseract first, then easyocr as fallback.
    """
    try:
        # Open image from bytes
        image = Image.open(io.BytesIO(image_bytes))
        
        # Try pytesseract first (faster, more common)
        if OCR_AVAILABLE:
            try:
                # Convert to RGB if necessary (pytesseract needs RGB)
                if image.mode != 'RGB':
                    image = image.convert('RGB')
                
                # Extract text using pytesseract
                text = pytesseract.image_to_string(image, lang='eng')
                
                if text and text.strip():
                    return text.strip()
            except Exception as e:
                st.warning(f"pytesseract failed for {file_name}: {str(e)[:100]}")
        
        # Fallback to easyocr if available
        if EASYOCR_AVAILABLE:
            try:
                # Initialize EasyOCR reader (only once, cache it)
                if 'easyocr_reader' not in st.session_state:
                    st.session_state.easyocr_reader = easyocr.Reader(['en'], gpu=False)
                
                reader = st.session_state.easyocr_reader
                
                # Convert PIL image to numpy array for easyocr
                img_array = np.array(image)
                
                # Extract text
                results = reader.readtext(img_array)
                text = '\n'.join([result[1] for result in results])
                
                if text and text.strip():
                    return text.strip()
            except Exception as e:
                st.warning(f"easyocr failed for {file_name}: {str(e)[:100]}")
        
        return None
        
    except Exception as e:
        st.warning(f"Could not process image {file_name}: {str(e)[:100]}")
        return None

def extract_receipt_with_ollama(files, model_name: str = None) -> list[dict]:
    """
    Extract receipt data from uploaded files using Ollama.
    
    Args:
        files: List of uploaded file objects (from st.file_uploader with accept_multiple_files=True)
        model_name: Ollama model name to use
    
    Returns:
        List of dictionaries, one per receipt, with keys:
        date, merchant, total_amount, currency, vat_amount, suggested_category, confidence
    """
    if model_name is None:
        model_name = st.session_state.get('selected_model', MODEL_NAME)
    
    if not files:
        return []
    
    # Normalize to list if single file
    if not isinstance(files, list):
        files = [files]
    
    all_results = []
    batch_size = 3  # Process max 3-4 files at a time
    
    # Process files in batches
    for batch_start in range(0, len(files), batch_size):
        batch_end = min(batch_start + batch_size, len(files))
        batch_files = files[batch_start:batch_end]
        
        batch_texts = []
        batch_file_names = []
        
        # Extract text from each file in batch
        for file in batch_files:
            file_name = file.name
            file_ext = file_name.lower().split('.')[-1] if '.' in file_name else ''
            
            text_content = None
            
            if file_ext == 'pdf':
                # Extract text from PDF using pdfplumber
                try:
                    # Read file content once
                    file_bytes = file.read()
                    # Reset file pointer for potential reuse
                    if hasattr(file, 'seek'):
                        file.seek(0)
                    with pdfplumber.open(io.BytesIO(file_bytes)) as pdf:
                        text_parts = []
                        for page in pdf.pages:
                            page_text = page.extract_text()
                            if page_text:
                                text_parts.append(page_text)
                        text_content = "\n\n".join(text_parts)
                except Exception as e:
                    st.warning(f"Could not extract text from PDF {file_name}: {str(e)}")
                    text_content = None
            
            elif file_ext in ['jpg', 'jpeg', 'png']:
                # Validate image quality first
                try:
                    file_bytes = file.read()
                    if hasattr(file, 'seek'):
                        file.seek(0)
                    
                    # Check image quality before processing
                    is_valid, error_message = validate_image_quality(file_bytes, file_name)
                    
                    if not is_valid:
                        st.error(error_message)
                        # Still try to process, but warn the user
                        st.warning("Processing anyway, but results may be poor. Consider using a higher quality image.")
                        # Option to skip this file
                        if st.checkbox(f"Skip {file_name} (low quality)", key=f"skip_{file_name}"):
                            continue
                    
                    # Extract text from image using OCR
                    with st.spinner(f"Extracting text from {file_name} using OCR..."):
                        text_content = extract_text_from_image(file_bytes, file_name)
                    
                    if text_content:
                        st.success(f"Extracted {len(text_content)} characters from {file_name}")
                    else:
                        if not OCR_AVAILABLE and not EASYOCR_AVAILABLE:
                            st.error(f"No OCR library available. Install pytesseract: `pip install pytesseract` (requires Tesseract OCR)")
                            st.info("Alternatively: `pip install easyocr`")
                        else:
                            st.warning(f"Could not extract text from {file_name}. Image may be unclear or contain no text.")
                        text_content = None
                except Exception as e:
                    st.warning(f"Error processing image {file_name}: {str(e)[:100]}")
                    text_content = None
            
            if text_content:
                batch_texts.append(text_content)
                batch_file_names.append(file_name)
            else:
                st.warning(f"Skipping {file_name}: Could not extract text content")
        
        if not batch_texts:
            continue
        
        # Build prompt for batch
        receipts_text = ""
        for i, (text, name) in enumerate(zip(batch_texts, batch_file_names)):
            receipts_text += f"\n\n--- Receipt {i+1}: {name} ---\n{text}\n"
        
        prompt = f"""You are a precise receipt data extractor.
Output ONLY a JSON array of objects. Nothing else. No explanations.

For each receipt provided, extract exactly these fields:
- date: YYYY-MM-DD format or null if not found
- merchant: Clean shop/merchant name (remove extra codes, keep main name)
- total_amount: Float number (the total paid, positive value)
- currency: Currency code (EUR, USD, GBP, etc.) or null
- vat_amount: Float number or null if not found/separated
- suggested_category: One of: Groceries, Restaurant, Transport, Office, Other
- confidence: Float between 0.0 and 1.0 (how confident you are in the extraction)

RULES:
1. Extract ALL receipts provided (one object per receipt)
2. If date is unclear, try to infer from context or use null
3. Merchant name should be clean and readable (e.g., "LIDL" not "LIDL**5400*")
4. Total amount should be the final amount paid (positive number)
5. VAT amount is optional - use null if not clearly separated
6. Category should match the merchant/items (Groceries, Restaurant, Transport, Office, Other)
7. Confidence reflects how clear/complete the receipt data is

Receipt data:{receipts_text}

Example output format:
[{{"date":"2025-01-15","merchant":"LIDL","total_amount":45.67,"currency":"EUR","vat_amount":7.62,"suggested_category":"Groceries","confidence":0.95}},{{"date":"2025-01-16","merchant":"Starbucks","total_amount":12.50,"currency":"EUR","vat_amount":null,"suggested_category":"Restaurant","confidence":0.85}}]

Respond with ONLY the JSON array:"""
        
        prompt_tokens_est = rough_token_count(prompt)
        remaining_context = N_CTX - prompt_tokens_est - PROMPT_SAFETY_MARGIN
        
        if remaining_context < MIN_OUTPUT_TOKENS:
            st.warning(f"Batch too large (~{prompt_tokens_est} est. tokens). Processing files individually...")
            # Process files one by one if batch is too large
            for file, file_name in zip(batch_files, batch_file_names):
                single_result = extract_receipt_with_ollama([file], model_name)
                all_results.extend(single_result)
            continue
        
        max_tokens = max(MIN_OUTPUT_TOKENS, min(remaining_context, 1000))
        
        try:
            with st.spinner(f"Extracting receipt data from {len(batch_texts)} file(s)..."):
                if ollama_client:
                    response = ollama_client.chat(
                        model=model_name,
                        messages=[{"role": "user", "content": prompt}],
                        options={
                            "temperature": 0.05,
                            "top_p": 0.9,
                            "top_k": 40,
                            "num_predict": max_tokens,
                            "repeat_penalty": 1.12,
                            "stop": ["```", "\n\n\n", "Note:", "Explanation:", "```json"]
                        }
                    )
                else:
                    response = ollama.chat(
                        model=model_name,
                        messages=[{"role": "user", "content": prompt}],
                        options={
                            "temperature": 0.05,
                            "top_p": 0.9,
                            "top_k": 40,
                            "num_predict": max_tokens,
                            "repeat_penalty": 1.12,
                            "stop": ["```", "\n\n\n", "Note:", "Explanation:", "```json"]
                        }
                    )
            
            if not response or 'message' not in response or 'content' not in response.get('message', {}):
                st.error(f"Ollama returned invalid response for batch")
                continue
            
            text = response['message']['content'].strip()
            
            if not text or len(text) < 10:
                st.error(f"LLM returned empty response for batch")
                continue
            
            # Extract JSON using same strategies as normalize_with_llm
            json_str = None
            
            # Strategy 1: Look for JSON wrapped in markdown code blocks
            json_match = re.search(r'```(?:json)?\s*(\[.*?\])\s*```', text, re.DOTALL)
            if json_match:
                json_str = json_match.group(1)
            
            # Strategy 2: Find JSON array boundaries
            if not json_str:
                start = text.find('[')
                end = text.rfind(']') + 1
                if start >= 0 and end > start:
                    json_str = text[start:end]
            
            # Strategy 3: Try to salvage partial JSON
            if not json_str:
                json_match = re.search(r'\[.*?\]', text, re.DOTALL)
                if json_match:
                    json_str = json_match.group(0)
                else:
                    last_complete = text.rfind('},')
                    if last_complete > 0:
                        json_str = text[:last_complete + 1] + ']'
            
            if not json_str:
                st.error(f"Could not extract JSON from response for batch")
                with st.expander("Debug: Full LLM Response", expanded=False):
                    st.code(text, language="text")
                continue
            
            # Clean up JSON
            json_str = json_str.strip()
            json_str = re.sub(r',(\s*[}\]])', r'\1', json_str)
            
            # Parse JSON
            try:
                data = json.loads(json_str)
                if isinstance(data, list):
                    all_results.extend(data)
                    st.success(f"Extracted {len(data)} receipt(s) from batch")
                else:
                    st.warning(f"Expected JSON array, got {type(data).__name__}")
            except json.JSONDecodeError as e:
                st.error(f"JSON parse error: {str(e)}")
                with st.expander("Debug: JSON Parse Error", expanded=False):
                    st.code(f"Error: {str(e)}\n\nAttempted JSON:\n{json_str}", language="text")
                continue
        
        except Exception as e:
            st.error(f"Batch processing failed: {str(e)[:100]}")
            continue
    
    return all_results

# ────────────────────────────────────────────────
#                   MAIN APP
# ────────────────────────────────────────────────

st.title("Local Receipt Extractor – Ollama Powered")
st.caption("100% local • Ollama + Safer batching • Context protection")

# Sidebar text
st.sidebar.write("100% local – photos never leave your computer")

# Get available models and let user select
available_models = get_available_models()

if not available_models:
    st.error("No models found in Ollama")
    st.info("""
**Setup required:**
1. Install Ollama: https://ollama.com
2. Terminal: `ollama serve`
3. Terminal: `ollama pull mistral` (or `llama3.2`, `qwen2.5`, etc.)
4. Refresh this page
""")
    st.stop()

# Model selector
if 'selected_model' not in st.session_state or st.session_state.selected_model not in available_models:
    # Default to the first available model, or MODEL_NAME if it's available
    default_model = MODEL_NAME if MODEL_NAME in available_models else available_models[0]
    st.session_state.selected_model = default_model

# Get current index for selectbox
current_index = 0
if st.session_state.selected_model in available_models:
    current_index = available_models.index(st.session_state.selected_model)

selected_model = st.selectbox(
    "Select Ollama Model",
    options=available_models,
    index=current_index,
    help="Choose which Ollama model to use for processing transactions"
)

# Update session state
st.session_state.selected_model = selected_model

st.info(f"""
**Current setup:**
- Ollama running
- Selected model: **{selected_model}**
- To install more models: `ollama pull <model_name>`
""")

uploaded = st.file_uploader("Upload receipt photo(s) or PDF", type=["jpg","jpeg","png","pdf"], accept_multiple_files=True)

if uploaded:
    # Normalize to list if single file (for accept_multiple_files=True, it's always a list)
    if not isinstance(uploaded, list):
        uploaded = [uploaded]
    
    # Check if we have CSV files (old bank statement flow) or receipt files
    has_csv = any(f.name.lower().endswith('.csv') for f in uploaded)
    has_receipt_files = any(f.name.lower().endswith(('.jpg', '.jpeg', '.png', '.pdf')) for f in uploaded)
    
    # If CSV files, use old bank statement processing
    if has_csv and len(uploaded) == 1 and uploaded[0].name.lower().endswith('.csv'):
        try:
            df_raw = pd.read_csv(uploaded[0])
            st.subheader("Raw data preview")
            st.dataframe(df_raw.head(10))
        except Exception as e:
            st.error(f"CSV error: {e}")
            st.stop()
    elif has_receipt_files:
        # Check if OCR is available for images
        has_images = any(f.name.lower().endswith(('.jpg', '.jpeg', '.png')) for f in uploaded)
        if has_images:
            if not OCR_AVAILABLE and not EASYOCR_AVAILABLE:
                st.warning("**OCR not available** - Image text extraction requires an OCR library.")
                st.info("""
                **To enable image processing, install one of:**
                - `pip install pytesseract` (also requires [Tesseract OCR](https://github.com/tesseract-ocr/tesseract))
                - `pip install easyocr` (self-contained, no external dependencies)
                
                PDF files will work without OCR.
                """)
            else:
                ocr_status = "pytesseract" if OCR_AVAILABLE else "easyocr"
                st.success(f"OCR available ({ocr_status}) - Images will be processed")
                
                # Pre-validate image quality before processing
                st.subheader("Image Quality Check")
                image_files = [f for f in uploaded if f.name.lower().endswith(('.jpg', '.jpeg', '.png'))]
                quality_issues = []
                
                for img_file in image_files:
                    try:
                        img_bytes = img_file.read()
                        if hasattr(img_file, 'seek'):
                            img_file.seek(0)
                        
                        is_valid, error_msg = validate_image_quality(img_bytes, img_file.name)
                        
                        if is_valid:
                            # Show image info even if valid
                            image = Image.open(io.BytesIO(img_bytes))
                            width, height = image.size
                            file_size_kb = len(img_bytes) / 1024
                            st.success(f"**{img_file.name}**: {width}×{height}px, {file_size_kb:.1f}KB - Quality OK")
                        else:
                            quality_issues.append((img_file.name, error_msg))
                            st.error(error_msg)
                    except Exception as e:
                        st.warning(f"Could not check quality of {img_file.name}: {str(e)[:100]}")
                
                if quality_issues:
                    st.warning(f"**{len(quality_issues)} image(s) have quality issues.** Processing may still work, but results may be poor.")
        
        # Use receipt extraction for images/PDFs
        if st.button("Extract Receipt Data", type="primary"):
            with st.status("Extracting receipt data...", expanded=True) as status:
                status.update(label="Processing receipt files with Ollama...", state="running")
                receipts = extract_receipt_with_ollama(uploaded, model_name=st.session_state.selected_model)
                
                if receipts:
                    status.update(label=f"Extracted {len(receipts)} receipt(s)", state="complete")
                    
                    # Convert to DataFrame for display
                    df_receipts = pd.DataFrame(receipts)
                    
                    st.subheader("Extracted Receipt Data")
                    st.dataframe(df_receipts, use_container_width=True)
                    
                    # Show summary
                    if not df_receipts.empty:
                        col1, col2, col3 = st.columns(3)
                        with col1:
                            total = df_receipts['total_amount'].sum() if 'total_amount' in df_receipts else 0
                            st.metric("Total Amount", f"€{total:,.2f}" if total else "N/A")
                        with col2:
                            avg_confidence = df_receipts['confidence'].mean() if 'confidence' in df_receipts else 0
                            st.metric("Avg Confidence", f"{avg_confidence:.2f}" if avg_confidence else "N/A")
                        with col3:
                            merchants = df_receipts['merchant'].nunique() if 'merchant' in df_receipts else 0
                            st.metric("Unique Merchants", merchants)
                        
                        # Download CSV
                        csv_data = df_receipts.to_csv(index=False).encode('utf-8')
                        st.download_button(
                            "Download Receipt Data (CSV)",
                            csv_data,
                            "receipts_extracted.csv",
                            "text/csv",
                            use_container_width=True
                        )
                else:
                    status.update(label="No receipts extracted", state="error")
                    st.warning("Could not extract receipt data from the uploaded files.")
        else:
            st.info(f"{len(uploaded)} file(s) uploaded. Click 'Extract Receipt Data' to process them.")
            st.caption("Supported formats: PDF, JPG, JPEG, PNG")
    else:
        st.error("Unsupported file type. Please upload CSV files or receipt images/PDFs.")
        st.stop()
    
    # Only show the old bank statement processing UI if we have CSV data
    if has_csv and len(uploaded) == 1 and uploaded[0].name.lower().endswith('.csv') and 'df_raw' in locals():
        categorize_now = st.checkbox("Also categorize transactions", True)

        categories_input = st.text_area(
            "Categories (one per line — last = fallback)",
            value="Subscriptions & Streaming\nGroceries & Supermarkets\nRestaurants & Food Delivery\nTransport & Fuel\nUtilities & Bills\nHousing & Rent\nHealth & Medical\nEntertainment & Leisure\nShopping & Clothes\nTravel & Holidays\nSalary & Income\nTransfers & Movements\nUncategorized / Other",
            height=180
        )
        categories = [c.strip() for c in categories_input.splitlines() if c.strip()]

        if st.button("Process Transactions", type="primary"):
            with st.status("Processing...", expanded=True) as status:
                status.update(label="Normalizing with Ollama...", state="running")
                df_normalized = normalize_with_llm(df_raw, model_name=st.session_state.selected_model)
                print(df_normalized)
                if df_normalized.empty:
                    status.update(label="Normalization failed", state="error")
                    st.stop()

                status.update(label=f"Normalized {len(df_normalized):,} transactions", state="complete")

                if categorize_now and not df_normalized.empty:
                    status.update(label="Categorizing with Ollama...", state="running")
                    df_final, failed_batches = categorize_with_llm(df_normalized, categories, model_name=st.session_state.selected_model)
                    # Store failed batches in session state for retry
                    if failed_batches:
                        st.session_state.failed_categorization_batches = failed_batches
                        st.session_state.original_normalized_df = df_normalized.copy()
                        st.session_state.original_categories = categories.copy()
                        st.session_state.current_df_final = df_final.copy()
                    else:
                        # Clear failed batches if categorization was successful
                        if 'failed_categorization_batches' in st.session_state:
                            del st.session_state.failed_categorization_batches
                    status.update(label="Categorization complete!", state="complete")
                else:
                    df_final = df_normalized.copy()
                    df_final["category"] = "Not categorized"
                    failed_batches = {}
                    if 'failed_categorization_batches' in st.session_state:
                        del st.session_state.failed_categorization_batches

            # ── Results ─────────────────────────────────────
            st.subheader("Final Results")
        
            # Check for uncategorized transactions
            uncategorized_count = (df_final["category"] == "Uncategorized").sum()
            if uncategorized_count > 0:
                st.warning(f"{uncategorized_count} transaction(s) are categorized as 'Uncategorized'")
                with st.expander(f"View Uncategorized Transactions ({uncategorized_count})", expanded=False):
                    uncategorized_df = df_final[df_final["category"] == "Uncategorized"][["date", "description", "amount"]]
                    st.dataframe(uncategorized_df, use_container_width=True)
                    st.caption("If these should have categories, you can retry the failed batches or manually edit the CSV after download.")
            
            st.dataframe(df_final, use_container_width=True)

            expenses = df_final[df_final["amount"] < 0]
            total_spent = expenses["amount"].sum()

            col1, col2 = st.columns(2)
            with col1:
                st.metric("Total Expenses", f"€{abs(total_spent):,.2f}")
            with col2:
                income = df_final[df_final["amount"] > 0]["amount"].sum()
                st.metric("Total Income", f"€{income:,.2f}")

            if not expenses.empty:
                st.subheader("Spending by Category")
                cat_spend = expenses.groupby("category")["amount"].sum().abs().sort_values(ascending=True)
                st.bar_chart(cat_spend)

            # Retry failed batches section
            if 'failed_categorization_batches' in st.session_state and st.session_state.failed_categorization_batches:
                st.divider()
                st.subheader("Retry Failed Batches")
                failed_batches_dict = st.session_state.failed_categorization_batches
                st.warning(f"{len(failed_batches_dict)} batch(es) failed during categorization and were assigned 'Uncategorized'.")
                
                for batch_num, batch_info in failed_batches_dict.items():
                    with st.expander(f"Batch {batch_num} - Error: {batch_info['error'][:100]}", expanded=False):
                        st.error(f"**Error:** {batch_info['error']}")
                        st.write(f"**Rows affected:** {len(batch_info['chunk_df'])}")
                        if batch_info.get('response_text'):
                            st.write("**LLM Response:**")
                            st.code(batch_info['response_text'][:1000], language="text")
                            if len(batch_info['response_text']) > 1000:
                                st.caption(f"... (truncated, full length: {len(batch_info['response_text'])})")
                        
                        col1, col2 = st.columns([1, 3])
                        with col1:
                            if st.button(f"🔄 Retry Batch {batch_num}", key=f"retry_batch_{batch_num}"):
                                # Retry this specific batch
                                st.session_state[f"retry_batch_{batch_num}"] = True
                                st.rerun()
                
            # Retry all failed batches button
            if st.button("Retry All Failed Batches", type="primary"):
                    st.session_state.retry_all_failed = True
                    st.rerun()
            
            # Handle retry actions
            if 'retry_all_failed' in st.session_state and st.session_state.retry_all_failed:
                del st.session_state.retry_all_failed
                if 'failed_categorization_batches' in st.session_state and 'current_df_final' in st.session_state:
                    failed_batches_dict = st.session_state.failed_categorization_batches
                    df_current = st.session_state.current_df_final.copy()
                    categories_retry = st.session_state.original_categories
                    model_name_retry = st.session_state.selected_model
                    
                    st.info(f"Retrying {len(failed_batches_dict)} failed batch(es)...")
                    
                    # Find rows that need to be replaced (those with "Uncategorized" from failed batches)
                    for batch_num, batch_info in failed_batches_dict.items():
                        chunk_df = batch_info['chunk_df']
                        categories_str_retry = batch_info['categories_str']
                        
                        with st.spinner(f"Retrying batch {batch_num}..."):
                            retry_result, error, _ = categorize_batch(
                                chunk_df, categories_retry, categories_str_retry, 
                                batch_num, len(failed_batches_dict), model_name_retry
                            )
                            
                            if not error:
                                # Update the dataframe - replace rows with new categories
                                # Find the indices in df_current that correspond to these rows
                                # We need to match by original index positions
                                st.success(f"Batch {batch_num} retry successful!")
                                # Remove from failed batches
                                del failed_batches_dict[batch_num]
                            else:
                                st.error(f"Batch {batch_num} retry failed: {str(error)[:100]}")
                    
                    # Update session state
                    if failed_batches_dict:
                        st.session_state.failed_categorization_batches = failed_batches_dict
                    else:
                        if 'failed_categorization_batches' in st.session_state:
                            del st.session_state.failed_categorization_batches
                        st.success("All batches retried successfully!")
                        st.rerun()
            
            # Handle individual batch retry
            for batch_num in list(st.session_state.keys()):
                if batch_num.startswith("retry_batch_") and st.session_state[batch_num]:
                    del st.session_state[batch_num]
                    batch_num_int = int(batch_num.replace("retry_batch_", ""))
                    
                    if 'failed_categorization_batches' in st.session_state:
                        failed_batches_dict = st.session_state.failed_categorization_batches
                        if batch_num_int in failed_batches_dict:
                            batch_info = failed_batches_dict[batch_num_int]
                            chunk_df = batch_info['chunk_df']
                            categories_retry = batch_info['categories']
                            categories_str_retry = batch_info['categories_str']
                            model_name_retry = batch_info['model_name']
                            
                            with st.spinner(f"Retrying batch {batch_num_int}..."):
                                retry_result, error, _ = categorize_batch(
                                    chunk_df, categories_retry, categories_str_retry,
                                    batch_num_int, len(failed_batches_dict), model_name_retry
                                )
                                
                                if not error:
                                    st.success(f"Batch {batch_num_int} retry successful!")
                                    # Update current_df_final with new categories
                                    # TODO: Update the actual dataframe here
                                    del failed_batches_dict[batch_num_int]
                                    if not failed_batches_dict:
                                        if 'failed_categorization_batches' in st.session_state:
                                            del st.session_state.failed_categorization_batches
                                else:
                                    st.error(f"Batch {batch_num_int} retry failed: {str(error)[:100]}")
                            
                            st.session_state.failed_categorization_batches = failed_batches_dict
                            st.rerun()

            # Download
            csv_data = df_final.to_csv(index=False).encode('utf-8')
        st.download_button(
            "Download cleaned CSV",
                csv_data,
                "transactions_cleaned.csv",
                "text/csv",
                use_container_width=True
            )
selected_model_display = st.session_state.get('selected_model', MODEL_NAME)
st.caption(f"Powered by Ollama ({selected_model_display}) • Context: {N_CTX} tokens")