parser/__init__.py

import gradio as gr
import pdfplumber
import pandas as pd
import json
import re
import pathlib
import numpy as np


def classify_pdf(pdf_path: str) -> str:
    """
    Peeks at the first page of the PDF to determine its origin.
    Returns 'MARG', 'TALLY', or 'UNKNOWN'.
    """
    with pdfplumber.open(pdf_path) as pdf:
        if len(pdf.pages) == 0:
            return "UNKNOWN"
            
        first_page = pdf.pages[0]
        words = first_page.extract_words(keep_blank_chars=False)
        
        if not words:
            return "UNKNOWN"
            
        text_stream = " ".join([w["text"].lower() for w in words])
        
        # --- 1. Check for Explicit Metadata/Branding ---
        if "marg erp" in text_stream or "margcompusoft" in text_stream:
            return "MARG"
        if "tally prime" in text_stream or "tally solutions" in text_stream:
            return "TALLY"
            
        # --- 2. Check Structural Vocabulary (Heuristics) ---
        marg_score = 0
        tally_score = 0
        
        # Marg Keywords (Strictly Marg-specific financial headers)
        # Removed "batch", "exp", and "hsn" as they are shared with Tally
        marg_keywords = ["s.no.", "mrp", "p.rate", "total stock", "op.stock", "cl.stock"]
        for kw in marg_keywords:
            if kw in text_stream:
                marg_score += 1
                
        # Tally Keywords (Hierarchical accounting)
        # Added "godown" which is highly specific to Tally
        tally_keywords = ["particulars", "opening balance", "inwards", "outwards", "closing balance", "grand total", "godown"]
        for kw in tally_keywords:
            if kw in text_stream:
                tally_score += 1
                
        # Look for Marg's distinct dashed lines
        dashed_lines = sum(1 for w in words if "----" in w["text"] or "====" in w["text"])
        if dashed_lines > 3:
            marg_score += 2

        # Check for explicit Tally report titles combined with Particulars
        if "stock summary" in text_stream and "particulars" in text_stream:
            tally_score += 2

        # --- 3. Final Decision ---
        if marg_score > tally_score and marg_score >= 1:
            return "MARG"
        elif tally_score > marg_score and tally_score >= 1:
            return "TALLY"
            
        # Check for explicit drawn tables (Vanilla Grids)
        if len(first_page.find_tables()) > 0:
            return "GENERIC"
            
        return "UNKNOWN"





def handle_parsing(file_obj):
    if file_obj is None:
        return None, {"error": "Please upload a document."}
        
    try:
        file_path = file_obj.name
        ext = pathlib.Path(file_path).suffix.lower()
        
        # ==========================================
        # ROUTE 1: PDF HANDLING
        # ==========================================
        if ext == ".pdf":
            file_type = classify_pdf(file_path)
            
            if file_type == "MARG":
                data = parse_universal_marg_pdf(file_path)
            elif file_type == "TALLY":
                data = parse_tally_pdf(file_path)
            elif file_type == "GENERIC":
                data = parse_generic_pdf(file_path)
            else:
                return None, {"error": "Format could not be identified as Marg, Tally, or a standard lined table."}
                
            state = {"format": file_type, "data": data}
            return state, data

        # ==========================================
        # ROUTE 2: SPREADSHEET HANDLING
        # ==========================================
        elif ext in [".csv", ".xlsx", ".xls"]:
            # Read the file blindly first to classify it
            if ext == ".csv":
                df = pd.read_csv(file_path, header=None)
            else:
                df = pd.read_excel(file_path, header=None)
                
            file_type = classify_dataframe(df)
            
            if file_type == "MARG":
                data = parse_marg_spreadsheet(df)
                state = {"format": "MARG", "data": data}
                return state, data
                
            elif file_type == "TALLY":
                # We will add parse_tally_spreadsheet() here next!
                return None, {"error": "Tally Spreadsheet parser coming soon!"}
                
            else:
                return None, {"error": "Spreadsheet format could not be identified as Marg or Tally."}
                
        else:
            return None, {"error": "Unsupported file type. Please upload a PDF, CSV, or Excel file."}
            
    except Exception as e:  
        return None, {"error": f"Failed to parse file: {str(e)}"}
        


def classify_dataframe(df: pd.DataFrame) -> str:
    """Scans the first 15 rows of a dataframe for Tally/Marg signatures."""
    # Convert the top portion to a single lowercase string for easy keyword matching
    top_rows_text = df.head(15).to_string().lower()
    
    if "marg erp" in top_rows_text or "margcompusoft" in top_rows_text:
        return "MARG"
    if "tally prime" in top_rows_text or "tally solutions" in top_rows_text:
        return "TALLY"
        
    # Heuristics
    if "particulars" in top_rows_text and "godown" in top_rows_text:
        return "TALLY"
    if "s.no" in top_rows_text and "mrp" in top_rows_text and "p.rate" in top_rows_text:
        return "MARG"
        
    return "GENERIC"





import pandas as pd
import numpy as np

def parse_marg_spreadsheet(df: pd.DataFrame) -> list:
    """
    Parses a raw Marg Excel/CSV DataFrame into the Universal JSON Tree structure.
    """
    # 1. Identify the actual table header row
    # Marg exports often have 3-5 rows of metadata at the top before the columns start.
    header_row_index = -1
    target_headers = {'description', 'particulars', 'item', 'product', 'qty', 'rate', 'mrp'}
    
    # Scan the first 20 rows to find where the actual table begins
    for idx, row in df.head(20).iterrows():
        # Convert row values to lowercase strings to check against our targets
        row_str_set = set(str(val).lower().strip() for val in row.values if pd.notna(val))
        # If we find at least 2 common Marg headers, this is our header row
        if len(target_headers.intersection(row_str_set)) >= 2: 
            header_row_index = idx
            break
            
    if header_row_index == -1:
        return [] # Could not find valid table headers
        
    # 2. Reshape the DataFrame
    # Set the found row as the column headers and drop everything above it
    df.columns = df.iloc[header_row_index].astype(str).str.strip()
    df = df.iloc[header_row_index + 1:].reset_index(drop=True)
    
    # Drop completely empty columns and "Unnamed" columns
    df = df.loc[:, df.columns.notna() & (df.columns != '') & (~df.columns.str.contains('Unnamed', case=False, na=False))]
    
    # Identify the primary description column
    desc_col = next((col for col in df.columns if col.lower() in ['description', 'particulars', 'item', 'product']), None)
            
    if not desc_col:
        return []

    # 3. Parse Data into the JSON Tree Structure
    data_tree = []
    current_block_name = "Default View"
    current_items = []
    
    for _, row in df.iterrows():
        desc_val = row[desc_col]
        
        # Skip completely empty rows
        if pd.isna(desc_val) or str(desc_val).strip() == '':
            continue
            
        desc_str = str(desc_val).strip()
        desc_lower = desc_str.lower()
        
        # Skip summary lines (Totals)
        if "total" in desc_lower or "page" in desc_lower or "grand" in desc_lower:
            continue
            
        # Extract row data and check for financials
        has_financials = False
        row_data = {}
        for col in df.columns:
            if col != desc_col:
                val = row[col]
                # If the cell is not empty/NaN
                if pd.notna(val) and str(val).strip() != '':
                    row_data[col] = str(val).strip()
                    # If any secondary column contains a digit, we consider it a financial/item row
                    if any(char.isdigit() for char in str(val)):
                        has_financials = True
                        
        # 4. Routing Logic: Is it a Group Header or an Item?
        # In Marg, if a row has a Description but NO numbers in other columns, it's a Group Header
        is_block_header = (not has_financials and len(row_data) == 0) or (not has_financials and desc_str.isupper())
        
        if is_block_header:
            # Commit the previous block if it has items
            if current_items:
                data_tree.append({
                    "name": current_block_name,
                    "type": "view",
                    "items": current_items
                })
                current_items = [] # Reset for next block
            
            # Start a new block, cleaning up Marg's "=" or "-" formatting
            current_block_name = desc_str.replace("=", "").replace("-", "").strip()
        else:
            # It's an item row, add it to the current block
            current_items.append({
                "particulars": desc_str,
                "data": row_data
            })
            
    # Append the final block after the loop finishes
    if current_items:
        data_tree.append({
            "name": current_block_name,
            "type": "view",
            "items": current_items
        })
        
    return data_tree


        
def midpoint(w: dict) -> float:
    return (w["x0"] + w["x1"]) / 2.0


# ==========================================
# 2. TALLY PARSER ENGINE
# ==========================================
KEYWORDS   = ("HSN/SAC:", "Batch:", "Godown:")
TALLY_STOP_WORDS = ("Grand Total", "Closing Stock", "Opening Stock")
NUM_RE     = re.compile(r"^-?[\d,]+(\.\d+)?$")
VALID_HEADERS = {"quantity", "qty", "rate", "value", "amount"}
SUPER_HEADERS = {"opening", "inwards", "outwards", "closing"}
CLUSTER_TOL = 3.0

def is_number(t: str) -> bool:
    return bool(NUM_RE.match(t.replace(" ", "")))

def line_text(words: list) -> str:
    return " ".join(w["text"] for w in words)

def has_financial(words: list, dividers: list) -> bool:
    if not dividers: return False
    for w in words:
        if midpoint(w) >= dividers[0]:
            if is_number(w["text"]):
                return True
    return False

def is_keyword_line(words: list) -> bool:
    t = line_text(words)
    return any(k in t for k in KEYWORDS)

def tally_is_stop_line(words: list) -> bool:
    t = line_text(words)
    return any(s in t for s in TALLY_STOP_WORDS)

def build_indent_map(blocks: list) -> dict:
    raw_x0s = [b["x0"] for b in blocks]
    if not raw_x0s:
        return {}
    unique = sorted(set(raw_x0s))
    clusters = []
    for x in unique:
        if clusters and x - clusters[-1] <= CLUSTER_TOL:
            continue
        clusters.append(x)
    indent_map = {}
    for x in raw_x0s:
        level = min(range(len(clusters)), key=lambda i: abs(clusters[i] - x))
        indent_map[x] = level
    return indent_map

def detect_columns(words: list, header_top: float):
    anchor_words = [w for w in words if abs(w["top"] - header_top) < 15 and w["text"].lower() in VALID_HEADERS]
    anchor_words.sort(key=lambda w: midpoint(w))
    super_words = [w for w in words if header_top - 25 < w["top"] < header_top + 5 and w["text"].lower() in SUPER_HEADERS]
    # super_words = [w for w in words if header_top - 25 < w["top"] < header_top - 2 and w["text"].lower() not in VALID_HEADERS and w["text"].lower() != "particulars"]
    super_words.sort(key=lambda w: w["x0"]) 

    dynamic_cols = []
    for aw in anchor_words:
        prefix = ""
        for i, sw in enumerate(super_words):
            if midpoint(aw) >= sw["x0"] - 20:
                if i + 1 < len(super_words):
                    if midpoint(aw) < super_words[i+1]["x0"] - 20:
                        prefix = sw["text"] + " "
                        break 
                else:
                    prefix = sw["text"] + " "
                    break
        label = (prefix + aw["text"]).strip().title()
        dynamic_cols.append({"label": label, "center": midpoint(aw)})

    if not dynamic_cols:
        dynamic_cols = [
            {"label": "Quantity", "center": 300},
            {"label": "Rate", "center": 380},
            {"label": "Value", "center": 460}
        ]

    dividers = [dynamic_cols[0]["center"] - 30]
    for i in range(len(dynamic_cols) - 1):
        dividers.append((dynamic_cols[i]["center"] + dynamic_cols[i+1]["center"]) / 2.0)
    return dynamic_cols, dividers

def parse_tally_pdf(pdf_path: str) -> list:
    full_tree = []
    with pdfplumber.open(pdf_path) as pdf:
        for page in pdf.pages:
            words = page.extract_words(keep_blank_chars=False)
            if not words: continue

            header_keywords = {"quantity", "qty", "rate", "value", "amount", "particulars"}
            header_candidates = [w["top"] for w in words if w["text"].lower() in header_keywords]
            header_top = min(header_candidates) if header_candidates else 80.0

            columns, dividers = detect_columns(words, header_top)
            lines = {}
            for w in words:
                if w["top"] < header_top + 2: continue
                y = round(w["top"] / 3) * 3
                lines.setdefault(y, []).append(w)

            blocks = []
            current = None

            def commit(blk):
                if blk: blocks.append(blk)

            for y in sorted(lines.keys()):
                lw = sorted(lines[y], key=lambda w: w["x0"])
                if not lw: continue
                if tally_is_stop_line(lw):
                    commit(current)
                    current = None
                    break

                fin = has_financial(lw, dividers)
                kw  = is_keyword_line(lw)
                first_x0 = lw[0]["x0"]
                p = []
                col_bins = [[] for _ in columns]
                
                for w in lw:
                    cx = midpoint(w)
                    if cx < dividers[0]:
                        p.append(w["text"])
                    else:
                        placed = False
                        for i in range(len(dividers) - 1):
                            if cx < dividers[i+1]:
                                col_bins[i].append(w["text"])
                                placed = True
                                break
                        if not placed:
                            col_bins[-1].append(w["text"])

                if fin or kw:
                    commit(current)
                    current = {"x0": first_x0, "particulars": p, "financials": col_bins}
                else:
                    current_is_anchor = (
                        current is not None and
                        (any(len(bin) > 0 for bin in current["financials"]) 
                         or is_keyword_line([{"text": t} for t in current["particulars"]]))
                    )
                    same_indent = (current is not None and abs(first_x0 - current["x0"]) <= CLUSTER_TOL)

                    if current_is_anchor and same_indent:
                        current["particulars"].extend([w["text"] for w in lw])
                    else:
                        commit(current)
                        current = {"x0": first_x0, "particulars": p, "financials": [[] for _ in columns]}

            commit(current)
            if not blocks: continue

            indent_map = build_indent_map(blocks)
            page_nodes = []
            hierarchy = []

            for blk in blocks:
                name = " ".join(blk["particulars"]).strip()
                if not name: continue
                level = indent_map.get(blk["x0"], 0)
                node = {"name": name, "level": level, "type": "group", "children": []}

                stock_data = {}
                has_stock = False
                for i, col in enumerate(columns):
                    val = " ".join(blk["financials"][i]).strip()
                    if val:
                        stock_data[col["label"]] = val
                        has_stock = True

                if has_stock: node["stock"] = stock_data

                if "HSN/SAC:" in name: node["type"] = "hsn_detail"
                elif name.startswith("Batch:"): node["type"] = "batch_detail"
                elif name.startswith("Godown:"): node["type"] = "godown_detail"
                elif has_stock: node["type"] = "item"

                while len(hierarchy) > level: hierarchy.pop()

                if not hierarchy:
                    page_nodes.append(node)
                else:
                    hierarchy[-1].setdefault("children", []).append(node)

                if node["type"] in ("group", "item"):
                    hierarchy.append(node)

            full_tree.extend(page_nodes)
    return full_tree


# ==========================================
# 3. MARG PARSER ENGINE
# ==========================================
MARG_HEADERS = {
    "s.no.", "s.no", "description", "product", "item", "particulars", 
    "packing", "pack", "batch", "btch", "exp", "expiry", 
    "hsn", "sac", "sku", "code", "unit", "gst%", "tax", 
    "op.stock", "opening", "purchase", "receipt", 
    "sale", "outwards", "cl.stock", "closing", "total", "stock", 
    "qty", "free", "rate", "p.rate", "mrp", "value", "amount"
}

def marg_is_stop_line(words: list) -> bool:
    text = " ".join(w["text"] for w in words).lower()
    return "grand total" in text or "page total" in text or text.startswith("total")

def parse_universal_marg_pdf(pdf_path: str) -> list:
    data_tree = []
    with pdfplumber.open(pdf_path) as pdf:
        for page_num, page in enumerate(pdf.pages):
            words = page.extract_words(keep_blank_chars=False)
            if not words: continue

            y_buckets = {}
            for w in words:
                y_b = round(w["top"] / 4) * 4
                if w["text"].lower() in MARG_HEADERS:
                    y_buckets.setdefault(y_b, []).append(w)

            if not y_buckets: continue
            best_y = max(y_buckets.keys(), key=lambda y: len(y_buckets[y]))
            if len(y_buckets[best_y]) < 2: continue 

            header_top = min([w["top"] for w in y_buckets[best_y]])
            anchor_words = [w for w in words if abs(w["top"] - header_top) < 15 and w["text"].lower() in MARG_HEADERS]
            anchor_words.sort(key=lambda w: midpoint(w))

            columns = []
            current_label = anchor_words[0]
            for w in anchor_words[1:]:
                if w["x0"] - current_label["x1"] < 15:
                    current_label["text"] += " " + w["text"]
                    current_label["x1"] = w["x1"]
                else:
                    columns.append({"label": current_label["text"].title(), "center": midpoint(current_label), "x0": current_label["x0"]})
                    current_label = w
            columns.append({"label": current_label["text"].title(), "center": midpoint(current_label), "x0": current_label["x0"]})

            desc_index = 0
            for i, col in enumerate(columns):
                if col["label"].lower() in {"description", "product", "item", "particulars"}:
                    desc_index = i
                    break
            
            dividers = []
            for i in range(desc_index): dividers.append((columns[i]["center"] + columns[i+1]["center"]) / 2.0)
            if desc_index + 1 < len(columns): dividers.append(columns[desc_index + 1]["x0"] - 15)
            for i in range(desc_index + 1, len(columns) - 1): dividers.append((columns[i]["center"] + columns[i+1]["center"]) / 2.0)
            if not dividers: dividers = [9999]

            lines = {}
            for w in words:
                if w["top"] < header_top + 10: continue
                if "-" * 5 in w["text"] or "=" * 5 in w["text"]:
                    if len(w["text"].replace("=", "").replace("-", "").strip()) == 0: continue 
                y = round(w["top"] / 3) * 3
                lines.setdefault(y, []).append(w)

            current_block_name = f"Report View {page_num + 1}"
            current_item = None

            for y in sorted(lines.keys()):
                lw = sorted(lines[y], key=lambda w: w["x0"])
                if not lw: continue
                if marg_is_stop_line(lw): break

                main_divider = dividers[desc_index] if len(dividers) > desc_index else dividers[0]
                fin_words = [w["text"] for w in lw if midpoint(w) >= main_divider]
                has_grid_data = any(any(char.isdigit() or char == '-' for char in text) for text in fin_words)

                col_bins = [[] for _ in columns] 
                for w in lw:
                    cx = midpoint(w)
                    placed = False
                    for i in range(len(dividers)):
                        if cx < dividers[i]:
                            col_bins[i].append(w["text"])
                            placed = True
                            break
                    if not placed:
                        col_bins[-1].append(w["text"])

                text_str = " ".join(col_bins[desc_index]).strip()
                pre_desc_text = ""
                for i in range(desc_index): pre_desc_text += " ".join(col_bins[i]).strip()

                is_block_header = text_str.isupper() and ("=" in text_str or "LIMITED" in text_str or "PHARMA" in text_str)

                if is_block_header:
                    if current_item:
                        block_node = next((n for n in data_tree if n["name"] == current_block_name), None)
                        if not block_node:
                            block_node = {"name": current_block_name, "type": "view", "items": []}
                            data_tree.append(block_node)
                        block_node["items"].append(current_item)
                        current_item = None
                    current_block_name = text_str.replace("=", "").replace("-", "").strip()
                    continue

                grid_dict = {}
                for i, col in enumerate(columns):
                    if i == desc_index: continue
                    val = " ".join(col_bins[i]).strip()
                    if val: grid_dict[col["label"]] = val

                starts_new_item = False
                if pre_desc_text: starts_new_item = True
                elif has_grid_data and current_item and any(k in current_item["data"] for k in grid_dict.keys() if k != "Unit"):
                    starts_new_item = True
                elif not current_item and (text_str or has_grid_data):
                    starts_new_item = True

                if starts_new_item:
                    if current_item:
                        block_node = next((n for n in data_tree if n["name"] == current_block_name), None)
                        if not block_node:
                            block_node = {"name": current_block_name, "type": "view", "items": []}
                            data_tree.append(block_node)
                        block_node["items"].append(current_item)
                    current_item = {"particulars": text_str, "data": grid_dict}
                else:
                    if current_item:
                        if text_str:
                            current_item["particulars"] = (current_item["particulars"] + " " + text_str) if current_item["particulars"] else text_str
                        if grid_dict: current_item["data"].update(grid_dict)

            if current_item:
                block_node = next((n for n in data_tree if n["name"] == current_block_name), None)
                if not block_node:
                    block_node = {"name": current_block_name, "type": "view", "items": []}
                    data_tree.append(block_node)
                block_node["items"].append(current_item)

    return data_tree








# ==========================================
# NEW: GENERIC LINED TABLE PARSER
# ==========================================
def parse_generic_pdf(pdf_path: str) -> list:
    """Extracts standard tables with explicit drawn borders."""
    data_tree = []
    with pdfplumber.open(pdf_path) as pdf:
        for page_num, page in enumerate(pdf.pages):
            # Native extraction using drawn lines
            tables = page.extract_tables()
            
            for table_idx, table in enumerate(tables):
                if not table or len(table) < 2: 
                    continue
                    
                # Clean up None types and newlines
                cleaned_table = [
                    [str(cell).replace('\n', ' ').strip() if cell else "" for cell in row]
                    for row in table
                ]
                
                headers = cleaned_table[0]
                items = []
                
                # Build the JSON structure to match our standard format
                for row in cleaned_table[1:]:
                    row_data = {}
                    for col_idx, cell in enumerate(row):
                        header_name = headers[col_idx] if col_idx < len(headers) and headers[col_idx] else f"Column_{col_idx+1}"
                        row_data[header_name] = cell
                        
                    items.append({
                        "particulars": row[0] if row else "", 
                        "data": row_data
                    })
                    
                data_tree.append({
                    "name": f"Page {page_num + 1} - Lined Table {table_idx + 1}",
                    "type": "generic_table",
                    "items": items,
                    "raw_table": cleaned_table # We pass this for easy Pandas rendering
                })
                
    return data_tree

# ==========================================
# 4. GRADIO APP & UI LOGIC
# ==========================================

# def handle_parsing(pdf_file):
#     if pdf_file is None:
#         return None, {"error": "Please upload a PDF document."}
#     try:
#         # 1. Route the document
#         file_type = classify_pdf(pdf_file.name)
        
#         # 2. Extract based on format
#         if file_type == "MARG":
#             data = parse_universal_marg_pdf(pdf_file.name)
#             state = {"format": "MARG", "data": data}
            
#         elif file_type == "TALLY":
#             data = parse_tally_pdf(pdf_file.name)
#             state = {"format": "TALLY", "data": data}
            
#         elif file_type == "GENERIC":
#             data = parse_generic_pdf(pdf_file.name)
#             state = {"format": "GENERIC", "data": data}
            
#         else:
#             return None, {"error": "Format could not be identified as Marg, Tally, or a standard lined table."}

#         return state, data 
        
#     except Exception as e:  
#         return None, {"error": f"Failed to parse PDF: {str(e)}"}  
        




# Recursive flattener for displaying Tally Hierarchies in a DataFrame
def flatten_tally_tree(node, rows_list, all_keys):
    # Use Unicode non-breaking spaces (\u00A0) so the web UI doesn't collapse them
    indent_spacing = "\u00A0\u00A0\u00A0\u00A0\u00A0\u00A0" # 6 spaces per level
    indent = indent_spacing * node.get("level", 0)
    
    prefix = "↳ " if node.get("level", 0) > 0 else ""
    name = indent + prefix + node.get("name", "")
    
    row_dict = {"Particulars": name}
    stock = node.get("stock", {})
    for k, v in stock.items():
        row_dict[k] = v
        if k not in all_keys:
            all_keys.append(k)
            
    rows_list.append(row_dict)
    
    for child in node.get("children", []):
        flatten_tally_tree(child, rows_list, all_keys)


# UI Layout
custom_theme = gr.themes.Soft(primary_hue="blue")

with gr.Blocks(title="Universal Accounting Parser", theme=custom_theme) as app:
    
    parsed_data_state = gr.State(None)

    gr.Markdown(
        """
        # 📊 Universal Tally & Marg Extractor
        Upload a stock report from either **Tally ERP** or **Marg ERP**. 
        Our router will detect the source system automatically, you can download the json.
        """
    )
    
    with gr.Row():
        with gr.Column(scale=1):
            # pdf_input = gr.File(label="Upload PDF Report", file_types=[".pdf"])
            # Change this in your UI block:
            pdf_input = gr.File(label="Upload Report", file_types=[".pdf", ".csv", ".xlsx", ".xls"])
            parse_button = gr.Button("Extract Tables", variant="primary")
            
            # --- ADD THIS EXAMPLES BLOCK ---
            gr.Examples(
                examples=[
                    ["examples/extreme_stress_test.pdf"], # Replace with your exact filename
                    ["examples/marg_enterprise_simulation (1).pdf"], # Replace with your exact filename
                    ["examples/tally_stress_test.pdf"]  # Replace with your exact filename
                ],
                inputs=pdf_input,
                label="Or try one of these examples:"
            )
            # -------------------------------
            
        with gr.Column(scale=3):
            with gr.Tabs():
                
                # TAB 1: DYNAMIC UI
                with gr.Tab("Visual Tables (UI)"):
                    
                    @gr.render(inputs=parsed_data_state)
                    def render_dynamic_tables(state_dict):
                        if not state_dict:
                            gr.Markdown("*(Upload a document to generate tables)*")
                            return
                            
                        doc_format = state_dict.get("format")
                        data_tree = state_dict.get("data", [])
                        
                        gr.Markdown(f"### 🔍 Detected System: **{doc_format}**")
                        gr.HTML("<hr/>")
                        
                        if doc_format == "MARG":
                            for block in data_tree:
                                gr.Markdown(f"#### 📋 {block.get('name', 'Report View')}")
                                if not block.get('items'): continue
                                    
                                all_data_keys = []
                                for item in block['items']:
                                    for k in item.get('data', {}).keys():
                                        if k not in all_data_keys: all_data_keys.append(k)
                                
                                columns = []
                                if "S.No." in all_data_keys:
                                    columns.append("S.No.")
                                    all_data_keys.remove("S.No.")
                                columns.append("Description / Particulars")
                                columns.extend(all_data_keys)
                                
                                rows = []
                                for item in block['items']:
                                    row = []
                                    if "S.No." in columns: row.append(item.get('data', {}).get("S.No.", ""))
                                    row.append(item.get('particulars', ''))
                                    for k in all_data_keys: row.append(item.get('data', {}).get(k, ""))
                                    rows.append(row)
                                    
                                df = pd.DataFrame(rows, columns=columns)
                                gr.Dataframe(value=df, interactive=False, label="")
                                gr.HTML("<br/>")

                        elif doc_format == "GENERIC":
                            for block in data_tree:
                                gr.Markdown(f"#### 📋 {block.get('name')}")
                                raw_table = block.get('raw_table', [])
                                
                                if len(raw_table) > 1:
                                    # Create DataFrame directly from the raw extracted table
                                    df = pd.DataFrame(raw_table[1:], columns=raw_table[0])
                                    gr.Dataframe(value=df, interactive=False, label="")
                                    gr.HTML("<br/>")
                                

                        elif doc_format == "TALLY":
                            all_data_keys = []
                            rows_dict_list = []
                            
                            # Flatten the recursive Tally tree
                            for node in data_tree:
                                flatten_tally_tree(node, rows_dict_list, all_data_keys)
                            
                            if not rows_dict_list:
                                gr.Markdown("*No items found.*")
                                return  # <--- CHANGED FROM continue TO return
                                
                            # Detect if this is a 12-column (Super Header) or 4-column layout
                            super_categories = ["Opening", "Inwards", "Outwards", "Closing"]
                            has_super_headers = any(any(k.startswith(sup) for k in all_data_keys) for sup in super_categories)
                            
                            # Use Gradio CSS variables so it looks perfect in both Light and Dark mode
                            html = "<div style='overflow-x: auto;'><table style='width: 100%; border-collapse: collapse; font-family: sans-serif; font-size: 14px; text-align: left; border: 1px solid var(--border-color-primary);'>"
                            html += "<thead>"
                            
                            if has_super_headers:
                                # Row 1: Super Headers (Colspan)
                                html += "<tr style='background-color: var(--background-fill-secondary);'>"
                                html += "<th rowspan='2' style='padding: 12px; border: 1px solid var(--border-color-primary); min-width: 300px;'>Particulars</th>"
                                for sup in super_categories:
                                    sub_cols = [k for k in all_data_keys if k.startswith(sup)]
                                    if sub_cols:
                                        html += f"<th colspan='{len(sub_cols)}' style='padding: 8px; border: 1px solid var(--border-color-primary); text-align: center; font-weight: bold;'>{sup} Balance</th>"
                                html += "</tr>"
                                
                                # Row 2: Sub Headers
                                html += "<tr style='background-color: var(--background-fill-secondary);'>"
                                for sup in super_categories:
                                    sub_cols = [k for k in all_data_keys if k.startswith(sup)]
                                    for k in sub_cols:
                                        sub_name = k.replace(sup + " ", "")
                                        html += f"<th style='padding: 8px; border: 1px solid var(--border-color-primary);'>{sub_name}</th>"
                                html += "</tr>"
                            else:
                                # Standard Flat Headers (for 4-column reports)
                                html += "<tr style='background-color: var(--background-fill-secondary);'>"
                                html += "<th style='padding: 12px; border: 1px solid var(--border-color-primary);'>Particulars</th>"
                                for k in all_data_keys:
                                    html += f"<th style='padding: 8px; border: 1px solid var(--border-color-primary);'>{k}</th>"
                                html += "</tr>"
                                
                            html += "</thead><tbody>"
                            
                            # Populate Data Rows with Indentations
                            for row in rows_dict_list:
                                html += "<tr style='border-bottom: 1px solid var(--border-color-primary);'>"
                                html += f"<td style='padding: 8px; border-right: 1px solid var(--border-color-primary); white-space: pre;'>{row.get('Particulars', '')}</td>"
                                
                                if has_super_headers:
                                    for sup in super_categories:
                                        sub_cols = [k for k in all_data_keys if k.startswith(sup)]
                                        for k in sub_cols:
                                            html += f"<td style='padding: 8px; border-right: 1px solid var(--border-color-primary);'>{row.get(k, '')}</td>"
                                else:
                                    for k in all_data_keys:  # <--- FIXED TYPO HERE (was all_keys)
                                        html += f"<td style='padding: 8px; border-right: 1px solid var(--border-color-primary);'>{row.get(k, '')}</td>"
                                        
                                html += "</tr>"
                                
                            html += "</tbody></table></div>"
                            
                            # Render the raw HTML block instead of the DataFrame
                            gr.HTML(html)

                            
                
                # TAB 2: RAW JSON
                with gr.Tab("Raw JSON Data"):
                    json_output = gr.JSON(label="Structured JSON Output")
    
    parse_button.click(
        fn=handle_parsing, 
        inputs=pdf_input, 
        outputs=[parsed_data_state, json_output]
    )

if __name__ == "__main__":
    app.launch()