# ============================================================================
# Tax Document Intelligence System — HuggingFace Spaces
# ============================================================================

import os
import re
import json
import time
import warnings
from pathlib import Path
from datetime import datetime
from typing import List

warnings.filterwarnings('ignore')

# Computer Vision & OCR
import easyocr
import cv2
import numpy as np
from PIL import Image, ImageDraw, ImageFont

# Deep Learning
import torch
from transformers import pipeline

# Database
from sqlalchemy import create_engine, Column, Integer, String, Float, Text, DateTime
from sqlalchemy.orm import declarative_base, sessionmaker

# Interface
import gradio as gr

# ============================================================================
# DATABASE SETUP
# ============================================================================

Base = declarative_base()

class Document(Base):
    __tablename__ = "documents"
    id              = Column(Integer, primary_key=True, autoincrement=True)
    filename        = Column(String(255))
    document_type   = Column(String(100))
    vendor          = Column(String(255))
    invoice_number  = Column(String(100))
    date            = Column(String(50))
    total_amount    = Column(String(50))
    subtotal        = Column(String(50))
    tax_amount      = Column(String(50))
    raw_text        = Column(Text)
    items           = Column(Text)
    confidence      = Column(Float)
    created_at      = Column(DateTime, default=datetime.utcnow)

engine  = create_engine("sqlite:///tax_documents.db", echo=False)
Base.metadata.create_all(engine)
Session = sessionmaker(bind=engine)

def save_to_db(data: dict) -> int:
    session = Session()
    try:
        doc = Document(
            filename       = data.get("filename", "uploaded_image"),
            document_type  = data.get("document_type", "unknown"),
            vendor         = data.get("vendor", ""),
            invoice_number = data.get("invoice_number", ""),
            date           = data.get("date", ""),
            total_amount   = data.get("total_amount", ""),
            subtotal       = data.get("subtotal", ""),
            tax_amount     = data.get("tax_amount", ""),
            raw_text       = data.get("raw_text", ""),
            items          = json.dumps(data.get("items", [])),
            confidence     = data.get("confidence", 0.0),
            created_at     = datetime.utcnow()
        )
        session.add(doc)
        session.commit()
        return doc.id
    except Exception as e:
        session.rollback()
        print(f"DB save error: {e}")
        return -1
    finally:
        session.close()

def get_all_records() -> List[dict]:
    session = Session()
    try:
        docs = session.query(Document).order_by(Document.created_at.desc()).all()
        return [
            {
                "ID"        : d.id,
                "File"      : d.filename,
                "Type"      : d.document_type,
                "Vendor"    : d.vendor,
                "Invoice #" : d.invoice_number,
                "Date"      : d.date,
                "Total"     : d.total_amount,
                "Saved At"  : str(d.created_at)[:19] if d.created_at else ""
            }
            for d in docs
        ]
    finally:
        session.close()

# ============================================================================
# OCR ENGINE
# ============================================================================

print("🔍 Initializing EasyOCR...")
use_gpu    = torch.cuda.is_available()
ocr_reader = easyocr.Reader(['en'], gpu=use_gpu, verbose=False)
print("✅ EasyOCR ready!")

def extract_text_from_image(image_path: str) -> str:
    try:
        results = ocr_reader.readtext(image_path, detail=0, paragraph=True)
        return "\n".join(results).strip()
    except Exception as e:
        return f"OCR Error: {e}"

def preprocess_image(image_path: str) -> str:
    try:
        img = cv2.imread(image_path)
        if img is None:
            return image_path
        gray     = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        denoised = cv2.fastNlMeansDenoising(gray, h=10)
        clahe    = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
        enhanced = clahe.apply(denoised)
        out_path = "/tmp/preprocessed.png"
        cv2.imwrite(out_path, enhanced)
        return out_path
    except:
        return image_path

# ============================================================================
# DOCUMENT CLASSIFIER
# ============================================================================

print("🤖 Loading Document Classifier...")
classifier = pipeline(
    "zero-shot-classification",
    model="typeform/distilbert-base-uncased-mnli",
    device=-1
)

DOCUMENT_LABELS = ["invoice", "receipt", "tax form", "bank statement", "purchase order"]

def classify_document(text: str) -> dict:
    if not text or len(text.strip()) < 10:
        return {"label": "unknown", "confidence": 0.0}
    try:
        result = classifier(text[:512], DOCUMENT_LABELS, multi_label=False)
        return {
            "label"     : result["labels"][0],
            "confidence": round(result["scores"][0], 3)
        }
    except Exception as e:
        return {"label": "document", "confidence": 0.5}

print("✅ Classifier ready!")

# ============================================================================
# ENTITY EXTRACTOR
# ============================================================================

def extract_entities(text: str) -> dict:
    entities = {
        "vendor": "", "invoice_number": "", "date": "",
        "total_amount": "", "subtotal": "", "tax_amount": "", "items": []
    }
    lines = text.split("\n")

    vendor_keywords = ["from:", "vendor:", "supplier:", "company:", "store:", "shop:"]
    for line in lines[:5]:
        line = line.strip()
        if len(line) > 3 and not re.match(r'^[\d\s\.\-\/]+$', line):
            entities["vendor"] = line
            break
    for line in lines:
        for kw in vendor_keywords:
            if kw in line.lower():
                entities["vendor"] = re.sub(rf'(?i){kw}\s*', '', line).strip()

    for pattern in [
        r'(?i)invoice\s*(?:no|number|#|num)[:\s]*([A-Z0-9\-/]+)',
        r'(?i)inv\s*(?:no|#)[:\s]*([A-Z0-9\-/]+)',
        r'(?i)bill\s*(?:no|#)[:\s]*([A-Z0-9\-/]+)',
        r'#([A-Z0-9]{4,15})\b',
    ]:
        m = re.search(pattern, text)
        if m:
            entities["invoice_number"] = m.group(1).strip()
            break

    for pattern in [
        r'\b(\d{1,2}[\-\/\.\s]\d{1,2}[\-\/\.\s]\d{2,4})\b',
        r'\b(\d{4}[\-\/\.]\d{1,2}[\-\/\.]\d{1,2})\b',
        r'(?i)(\d{1,2}\s+(?:jan|feb|mar|apr|may|jun|jul|aug|sep|oct|nov|dec)[a-z]*\s+\d{2,4})',
        r'(?i)((?:jan|feb|mar|apr|may|jun|jul|aug|sep|oct|nov|dec)[a-z]*\s+\d{1,2},?\s+\d{2,4})',
    ]:
        m = re.search(pattern, text)
        if m:
            entities["date"] = m.group(1).strip()
            break

    for pattern in [
        r'(?i)total\s*(?:amount)?\s*[:\$£€Rs\.]*\s*([\d,\.]+)',
        r'(?i)grand\s*total\s*[:\$£€Rs\.]*\s*([\d,\.]+)',
        r'(?i)amount\s*due\s*[:\$£€Rs\.]*\s*([\d,\.]+)',
        r'(?i)balance\s*due\s*[:\$£€Rs\.]*\s*([\d,\.]+)',
    ]:
        m = re.search(pattern, text)
        if m:
            entities["total_amount"] = m.group(1).strip().replace(",", "")
            break

    for pattern in [r'(?i)sub\s*total\s*[:\$£€Rs\.]*\s*([\d,\.]+)']:
        m = re.search(pattern, text)
        if m:
            entities["subtotal"] = m.group(1).strip().replace(",", "")
            break

    for pattern in [r'(?i)(?:vat|gst|tax|hst|pst)\s*[:\(\d%]*\s*[:\$£€Rs\.]*\s*([\d,\.]+)']:
        m = re.search(pattern, text)
        if m:
            entities["tax_amount"] = m.group(1).strip().replace(",", "")
            break

    items = re.findall(r'([A-Za-z][A-Za-z\s]{3,40})\s+[\$£€Rs\.]*\s*(\d+[\.,]\d{2})', text)
    entities["items"] = [f"{i[0].strip()} — {i[1]}" for i in items[:8]]

    return entities

# ============================================================================
# SAMPLE INVOICE GENERATOR
# ============================================================================

def create_sample_invoice():
    img  = Image.new('RGB', (700, 950), color='white')
    draw = ImageDraw.Draw(img)
    try:
        font_bold  = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", 22)
        font_title = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", 28)
        font_reg   = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", 17)
        font_small = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", 14)
    except:
        font_bold = font_title = font_reg = font_small = ImageFont.load_default()

    draw.rectangle([(0, 0), (700, 100)], fill='#1a1a2e')
    draw.text((30, 20),  "TECHMART SOLUTIONS PVT LTD",          fill='white',    font=font_title)
    draw.text((30, 60),  "www.techmart.com | info@techmart.com", fill='#aaaaaa', font=font_small)
    draw.text((530, 20), "INVOICE",                              fill='#f0a500', font=font_title)

    y = 120
    draw.text((30,  y), "Invoice No : INV-2024-00892",    fill='#333', font=font_bold)
    draw.text((400, y), "Date : 15 March 2024",           fill='#333', font=font_bold)
    y += 35
    draw.text((30,  y), "Bill To  : Muhammad Attiq",      fill='#333', font=font_reg)
    draw.text((400, y), "Due Date : 30 March 2024",       fill='#333', font=font_reg)
    y += 25
    draw.text((30,  y), "Address  : Islamabad, Pakistan", fill='#555', font=font_small)

    y += 40
    draw.rectangle([(30, y), (670, y+2)], fill='#1a1a2e')
    y += 15
    draw.rectangle([(30, y), (670, y+35)], fill='#f0a500')
    draw.text((40,  y+8), "Description", fill='white', font=font_bold)
    draw.text((370, y+8), "Qty",         fill='white', font=font_bold)
    draw.text((450, y+8), "Unit Price",  fill='white', font=font_bold)
    draw.text((570, y+8), "Amount",      fill='white', font=font_bold)
    y += 45

    items = [
        ("Laptop - Dell Inspiron 15",    1, 85000),
        ("Wireless Mouse Logitech M185", 2,  2500),
        ("USB-C Hub (7-port)",           1,  4500),
        ("HDMI Cable 2m",                3,   800),
    ]
    for i, (desc, qty, price) in enumerate(items):
        draw.rectangle([(30, y), (670, y+32)], fill='#f9f9f9' if i % 2 == 0 else 'white')
        draw.text((40,  y+7), desc,                  fill='#222', font=font_reg)
        draw.text((385, y+7), str(qty),               fill='#222', font=font_reg)
        draw.text((450, y+7), f"Rs. {price:,}",      fill='#222', font=font_reg)
        draw.text((555, y+7), f"Rs. {qty*price:,}",  fill='#222', font=font_reg)
        y += 32

    y += 20
    draw.rectangle([(30, y), (670, y+2)], fill='#cccccc')
    y += 15
    subtotal = sum(q*p for _,q,p in items)
    tax      = int(subtotal * 0.17)
    total    = subtotal + tax

    draw.text((450, y),   "Subtotal :",       fill='#333',    font=font_bold)
    draw.text((570, y),   f"Rs. {subtotal:,}", fill='#333',   font=font_reg)
    y += 30
    draw.text((450, y),   "GST (17%) :",      fill='#333',    font=font_bold)
    draw.text((570, y),   f"Rs. {tax:,}",     fill='#333',    font=font_reg)
    y += 30
    draw.rectangle([(440, y-5), (670, y+38)], fill='#1a1a2e')
    draw.text((450, y+8), "TOTAL DUE :",      fill='white',   font=font_bold)
    draw.text((555, y+8), f"Rs. {total:,}",   fill='#f0a500', font=font_bold)

    y += 80
    draw.text((30, y),    "Payment Terms: Net 15 Days",          fill='#888',    font=font_small)
    draw.text((30, y+20), "Bank: HBL | Account: 1234-5678-9012", fill='#888',    font=font_small)
    draw.text((30, y+40), "Thank you for your business!",         fill='#1a1a2e', font=font_bold)

    path = "/tmp/sample_invoice.png"
    img.save(path)
    return path

sample_path = create_sample_invoice()
print(f"✅ Sample invoice created: {sample_path}")

# ============================================================================
# PROCESSING PIPELINE
# ============================================================================

def process_document(image_input, filename: str = "uploaded_image") -> dict:
    result = {
        "filename": filename, "document_type": "unknown", "confidence": 0.0,
        "vendor": "", "invoice_number": "", "date": "", "total_amount": "",
        "subtotal": "", "tax_amount": "", "items": [], "raw_text": "",
        "db_id": -1, "processing_time": 0.0, "error": None
    }
    start = time.time()
    try:
        temp_path = "/tmp/processing_input.png"
        if isinstance(image_input, np.ndarray):
            Image.fromarray(image_input).save(temp_path)
        elif isinstance(image_input, str):
            temp_path = image_input
        else:
            image_input.save(temp_path)

        preprocessed = preprocess_image(temp_path)
        raw_text      = extract_text_from_image(preprocessed)
        result["raw_text"] = raw_text

        if not raw_text or len(raw_text.split()) < 3:
            result["error"] = "Could not extract text. Check image quality."
            return result

        cls = classify_document(raw_text)
        result["document_type"] = cls["label"]
        result["confidence"]    = cls["confidence"]

        entities = extract_entities(raw_text)
        result.update(entities)

        result["processing_time"] = round(time.time() - start, 2)
        result["db_id"]           = save_to_db(result)

    except Exception as e:
        result["error"]           = str(e)
        result["processing_time"] = round(time.time() - start, 2)

    return result


def format_result(result: dict) -> str:
    if result.get("error"):
        return f"❌ Error: {result['error']}"
    items_text = ""
    if result.get("items"):
        items_text = "\n" + "\n".join(f"     • {item}" for item in result["items"])
    return f"""✅ DOCUMENT PROCESSED SUCCESSFULLY
{'='*55}

📄 Document Info:
   Type          : {result['document_type'].upper()}
   Confidence    : {result['confidence']:.0%}
   Database ID   : #{result['db_id']}
   Processing    : {result['processing_time']}s

🏢 Extracted Data:
   Vendor        : {result['vendor'] or '—'}
   Invoice #     : {result['invoice_number'] or '—'}
   Date          : {result['date'] or '—'}

💰 Financial Data:
   Subtotal      : {result['subtotal'] or '—'}
   Tax           : {result['tax_amount'] or '—'}
   Total Amount  : {result['total_amount'] or '—'}

🛒 Line Items:{items_text if items_text else ' —'}

📝 Raw Text Preview:
{result['raw_text'][:300]}{'...' if len(result['raw_text']) > 300 else ''}
{'='*55}"""

# ============================================================================
# GRADIO INTERFACE
# ============================================================================

custom_css = """
.result-box { font-family: monospace; font-size: 13px; }
.gr-button-primary { background: #f0a500 !important; border: none !important; }
"""

last_result = {}

def gradio_process(image):
    global last_result
    if image is None:
        return "⚠️ Please upload an image first.", "", ""
    result      = process_document(image, filename="user_upload.png")
    last_result = result
    formatted   = format_result(result)
    summary = f"""
📊 Quick Summary
────────────────
Type     : {result['document_type'].upper()}
Vendor   : {result['vendor'] or 'Not detected'}
Total    : {result['total_amount'] or 'Not detected'}
Date     : {result['date'] or 'Not detected'}
DB ID    : #{result['db_id']}
Time     : {result['processing_time']}s
"""
    raw = result.get("raw_text", "No text extracted")
    return formatted, summary, raw

def gradio_load_sample():
    return Image.open(sample_path)

def gradio_get_records():
    records = get_all_records()
    if not records:
        return "📭 No records yet. Process a document to get started!"
    lines = [f"📊 DATABASE RECORDS ({len(records)} total)", "="*70]
    for r in records:
        lines.append(
            f"[#{r['ID']}] {r['Type'].upper():15s} | "
            f"Vendor: {r['Vendor'][:20]:20s} | "
            f"Total: {r['Total']:12s} | "
            f"Date: {r['Date']:12s} | "
            f"{r['Saved At']}"
        )
    lines.append("="*70)
    return "\n".join(lines)


with gr.Blocks(theme=gr.themes.Soft(), title="Tax Document Intelligence", css=custom_css) as demo:

    gr.Markdown("""
    # 🧾 Tax Document Intelligence System
    ### AI-Powered Receipt & Invoice Data Extraction
    *Upload any invoice or receipt → Auto-extract vendor, total, date, line items → Save to database*
    ---
    """)

    with gr.Tabs():

        with gr.Tab("📤 Upload & Extract"):
            gr.Markdown("### Upload your invoice or receipt image")
            with gr.Row():
                with gr.Column(scale=1):
                    image_input = gr.Image(label="Upload Invoice / Receipt", type="numpy", height=400)
                    with gr.Row():
                        process_btn = gr.Button("🚀 Extract Data", variant="primary", size="lg")
                        sample_btn  = gr.Button("📄 Load Sample",  variant="secondary", size="lg")
                    gr.Markdown("""
                    **💡 Tips for best results:**
                    - Use clear, well-lit photos
                    - Ensure all text is readable
                    - Works with JPG, PNG, WEBP
                    - Supports English text
                    """)
                with gr.Column(scale=1):
                    summary_out  = gr.Textbox(label="📊 Quick Summary",      lines=10, interactive=False, elem_classes="result-box")
                    raw_text_out = gr.Textbox(label="📝 Extracted Raw Text", lines=10, interactive=False, elem_classes="result-box")

            full_result_out = gr.Textbox(label="📋 Full Extraction Report", lines=20, interactive=False, elem_classes="result-box")

            process_btn.click(fn=gradio_process,    inputs=[image_input], outputs=[full_result_out, summary_out, raw_text_out])
            sample_btn.click( fn=gradio_load_sample, inputs=[],           outputs=[image_input])

        with gr.Tab("🗄️ Database Records"):
            gr.Markdown("### All processed documents stored in SQLite")
            refresh_btn = gr.Button("🔄 Refresh Records", variant="primary")
            records_out = gr.Textbox(label="Database Records", lines=20, interactive=False,
                                     elem_classes="result-box", value="Click 'Refresh Records' to load data.")
            refresh_btn.click(fn=gradio_get_records, inputs=[], outputs=[records_out])

        with gr.Tab("ℹ️ How It Works"):
            gr.Markdown("""
            ## 🔧 System Architecture
            ```
            Image Upload → Preprocessing → EasyOCR → DistilBERT Classifier → Entity Extractor → SQLite → Gradio UI
            ```

            ## 🏗️ Tech Stack
            | Component        | Technology           |
            |------------------|----------------------|
            | OCR Engine       | EasyOCR              |
            | Classifier       | DistilBERT Zero-Shot |
            | Entity Extraction| Smart Regex          |
            | Database         | SQLite + SQLAlchemy  |
            | Interface        | Gradio               |
            | Image Processing | OpenCV + PIL         |

            ---
            *Built for NLP & MLOps Final Year Project — Bootcamp Cohort 14*
            """)

# ============================================================================
# LAUNCH
# ============================================================================

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)