import re
import json
import tempfile
from datetime import datetime
from typing import List, Dict, Any, Optional

import gradio as gr
import pdfplumber

DATE_PATTERNS = [
    (re.compile(r"\b(\d{1,2})[/-](\d{1,2})[/-](\d{2,4})\b"), "dmy_num"),
    (re.compile(r"\b([A-Za-z]{3})[/-](\d{1,2})[/-](\d{4})\b"), "mon_d_y"),
    (re.compile(r"\b(\d{1,2})[/-]([A-Za-z]{3})[/-](\d{4})\b"), "d_mon_y"),
    (re.compile(r"\b(\d{1,2})-([A-Za-z]{3})-(\d{2,4})\b"), "d_mon_y"),
]

MONTHS = {
    "jan": 1, "feb": 2, "mar": 3, "apr": 4, "may": 5, "jun": 6,
    "jul": 7, "aug": 8, "sep": 9, "oct": 10, "nov": 11, "dec": 12,
}

SKIP_WORDS = [
    "opening balance", "page total", "closing balance", "statement of account",
    "txn date", "value date", "narration", "particulars", "withdrawals", "deposits",
    "balance", "account number", "branch name", "ifsc", "customer id",
]


def clean_text(x: Any) -> str:
    if x is None:
        return ""
    s = str(x).replace("\n", " ").replace("\r", " ")
    s = re.sub(r"\s+", " ", s).strip()
    return s


def fix_number_spaces(s: str) -> str:
    s = clean_text(s)
    # 11962.9 6 -> 11962.96, 675725 6.04 -> 6757256.04
    s = re.sub(r"(?<=\d)\s+(?=\d(?:\.\d+)?\b)", "", s)
    return s


def amount_value(s: Any) -> str:
    s = fix_number_spaces(str(s or ""))
    s = re.sub(r"\b(Cr|Dr)\b", "", s, flags=re.I).strip()
    m = re.search(r"-?\d[\d,]*\.?\d*", s)
    if not m:
        return ""
    try:
        return f"{float(m.group(0).replace(',', '')):.2f}"
    except Exception:
        return m.group(0).replace(',', '')


def balance_value(s: Any) -> str:
    s = fix_number_spaces(str(s or ""))
    sign = ""
    msign = re.search(r"\b(Cr|Dr)\b", s, flags=re.I)
    if msign:
        sign = msign.group(1).title()
    amt = amount_value(s)
    return f"{amt} {sign}".strip() if amt else ""


def parse_date(text: str, default_year: Optional[int] = None) -> str:
    text = clean_text(text)
    for rx, kind in DATE_PATTERNS:
        m = rx.search(text)
        if not m:
            continue
        try:
            if kind == "dmy_num":
                d, mo, y = m.groups()
                y = int(y)
                if y < 100:
                    y += 2000
                return datetime(y, int(mo), int(d)).strftime("%Y-%m-%d")
            if kind == "mon_d_y":
                mo, d, y = m.groups()
                return datetime(int(y), MONTHS[mo.lower()], int(d)).strftime("%Y-%m-%d")
            if kind == "d_mon_y":
                d, mo, y = m.groups()
                y = int(y)
                if y < 100:
                    y += 2000
                return datetime(y, MONTHS[mo.lower()], int(d)).strftime("%Y-%m-%d")
        except Exception:
            pass
    return ""


def is_skip_row(text: str) -> bool:
    t = text.lower()
    return any(w in t for w in SKIP_WORDS)


def normalize_headers(row: List[str]) -> List[str]:
    return [clean_text(c).lower().replace(" ", "_").replace("/", "_") for c in row]


def parse_known_table(headers: List[str], rows: List[List[str]]) -> List[Dict[str, str]]:
    out = []
    h = headers

    def idx(*names):
        for n in names:
            for i, col in enumerate(h):
                if n in col:
                    return i
        return -1

    date_i = idx("txn_date", "date")
    narr_i = idx("narration", "particular", "description")
    debit_i = idx("debit", "withdrawal")
    credit_i = idx("credit", "deposit")
    bal_i = idx("balance")

    if date_i < 0 or narr_i < 0 or bal_i < 0 or (debit_i < 0 and credit_i < 0):
        return []

    pending = None
    for row in rows:
        row = [clean_text(c) for c in row]
        txt = " ".join(row)
        if not txt or is_skip_row(txt):
            continue
        date = parse_date(row[date_i] if date_i < len(row) else "")
        narr = row[narr_i] if narr_i < len(row) else ""
        debit = amount_value(row[debit_i]) if debit_i >= 0 and debit_i < len(row) else ""
        credit = amount_value(row[credit_i]) if credit_i >= 0 and credit_i < len(row) else ""
        bal = balance_value(row[bal_i]) if bal_i < len(row) else ""

        if date and (debit or credit):
            if pending:
                out.append(pending)
            pending = {
                "date": date,
                "narration": narr,
                "voucher_type": "Payment" if debit else "Receipt",
                "amount": debit or credit,
                "closing_balance": bal,
            }
        elif pending and narr:
            pending["narration"] = clean_text(pending["narration"] + " " + narr)
            if bal and not pending.get("closing_balance"):
                pending["closing_balance"] = bal
    if pending:
        out.append(pending)
    return out


def parse_text_lines(text: str) -> List[Dict[str, str]]:
    out = []
    pending = None
    lines = [clean_text(x) for x in text.splitlines() if clean_text(x)]
    for line in lines:
        low = line.lower()
        if is_skip_row(line):
            continue
        date = parse_date(line)
        nums = re.findall(r"(?:\d[\d,]*\.?\d*)\s*(?:Cr|Dr)?", fix_number_spaces(line), flags=re.I)
        # heuristic: date + at least amount + balance
        if date and len(nums) >= 2:
            if pending:
                out.append(pending)
            bal = balance_value(nums[-1])
            amt = amount_value(nums[-2])
            # if line contains Dr near debit amount, assume Payment; otherwise Receipt for many bank PDFs
            before_bal = line[: line.rfind(nums[-1]) if nums[-1] in line else len(line)]
            vtype = "Payment" if re.search(r"\bDr\b", before_bal, re.I) or " to " in low or low.startswith("to ") else "Receipt"
            narration = line
            for rx, _ in DATE_PATTERNS:
                narration = rx.sub("", narration, count=1)
            narration = clean_text(narration)
            pending = {
                "date": date,
                "narration": narration,
                "voucher_type": vtype,
                "amount": amt,
                "closing_balance": bal,
            }
        elif pending and not date:
            pending["narration"] = clean_text(pending["narration"] + " " + line)
    if pending:
        out.append(pending)
    return out


def extract_pdf(pdf_path: str, max_pages: int = 0) -> Dict[str, Any]:
    transactions = []
    pages_done = 0
    debug = []
    with pdfplumber.open(pdf_path) as pdf:
        pages = pdf.pages[:max_pages] if max_pages and max_pages > 0 else pdf.pages
        for pno, page in enumerate(pages, start=1):
            pages_done += 1
            page_txns = []
            try:
                tables = page.extract_tables() or []
                for table in tables:
                    if not table or len(table) < 2:
                        continue
                    header_row_idx = None
                    for i, r in enumerate(table[:4]):
                        joined = " ".join(clean_text(c) for c in r).lower()
                        if ("date" in joined and ("debit" in joined or "withdraw" in joined) and ("credit" in joined or "deposit" in joined)):
                            header_row_idx = i
                            break
                    if header_row_idx is not None:
                        headers = normalize_headers(table[header_row_idx])
                        rows = table[header_row_idx + 1:]
                        page_txns.extend(parse_known_table(headers, rows))
            except Exception as e:
                debug.append(f"Page {pno} table error: {e}")

            if not page_txns:
                try:
                    text = page.extract_text(x_tolerance=1, y_tolerance=3) or ""
                    page_txns.extend(parse_text_lines(text))
                except Exception as e:
                    debug.append(f"Page {pno} text error: {e}")

            transactions.extend(page_txns)

    # final cleanup + de-duplicate simple duplicates
    final = []
    seen = set()
    for t in transactions:
        if not t.get("date") or not t.get("amount"):
            continue
        t["narration"] = clean_text(t.get("narration", ""))
        key = (t["date"], t["narration"], t["amount"], t.get("closing_balance", ""))
        if key not in seen:
            seen.add(key)
            final.append(t)

    return {
        "success": True,
        "total_transactions": len(final),
        "pages_processed": pages_done,
        "transactions": final,
        "debug": debug[:20],
    }


def gradio_extract(file, max_pages):
    if file is None:
        return {"success": False, "error": "Upload PDF file"}, ""
    result = extract_pdf(file.name, int(max_pages or 0))
    return result, json.dumps(result.get("transactions", []), indent=2, ensure_ascii=False)


def build_ui():
    with gr.Blocks(title="Free Bank Statement PDF to JSON") as demo:
        gr.Markdown("# Free Bank Statement PDF → Transaction JSON\nFast test app for Hugging Face Spaces. Output fields: `date`, `narration`, `voucher_type`, `amount`, `closing_balance`.")
        with gr.Row():
            pdf = gr.File(label="Upload Bank Statement PDF", file_types=[".pdf"])
            max_pages = gr.Number(label="Max pages for testing (0 = all)", value=0, precision=0)
        btn = gr.Button("Extract Transactions", variant="primary")
        output = gr.JSON(label="Full Response")
        tx_json = gr.Code(label="Transactions JSON", language="json")
        btn.click(gradio_extract, inputs=[pdf, max_pages], outputs=[output, tx_json])
    return demo

if __name__ == "__main__":
    build_ui().queue(default_concurrency_limit=2).launch(server_name="0.0.0.0", server_port=7860)