"""
BillStructAI demo — paste noisy OCR text from an invoice/receipt and get
back structured JSON, using a Qwen2.5-1.5B-Instruct base model fine-tuned
with QLoRA for this task.

Deploy: push this file + requirements.txt to a Hugging Face Space
(Gradio SDK). Set ADAPTER_REPO below to your uploaded LoRA adapter repo.
"""

import json
import re

import spaces  # add this import

import gradio as gr
import torch
from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

BASE_MODEL_NAME = "Qwen/Qwen2.5-1.5B-Instruct"
ADAPTER_REPO = "trishpurkait/billstructai-qwen-lora"  # <-- update after uploading adapter

SYSTEM_PROMPT = """
You are a precise invoice information extraction assistant.
Your task is to extract structured invoice data from noisy OCR text.
Return only valid JSON.
Do not explain.
Do not add markdown.
Use null for missing fields.
Do not hallucinate values that are not present in the OCR text.
""".strip()

EXAMPLE_OCR = """HAZLETON
REISSUE
LABORATORIES AMERICA, INC.
INVOICE
5510 NICHOLSON LANE, SUITE 400, KENSINGTON, MARYLAND 20095
PLEASE SEND PAYMENT TO:
INVOICE NO .: 812516
HAZLETON LABORATORIES AMERICA, INC.
USE THIS NUMBER
PO BOX 25065
FOR CHECK REMITTANCE
RICHMOND, VA
23260
USA
AND CORRESPONDENCE
BILLED TO:
DATE :
02-06-89
DR. J. D. HECK
LORILLARD RESEARCH CENTER
420 ENGLISH STREET
P.O. BOX 21688
GREENSBORO, NC 27420
SPONSOR REFERENCE : :selected:
IDENTIFICATION :
PURCHASE ORDER NO. 344A
HLA REFERENCE :
PROJECT AUTHORIZATION NO. :
1480
COST CENTER :
6130
PROJECT NUMBER :
20988
SERVICES RENDERED:
AMES SALMONELLA/MICROSOME REVERSE MUTATION ASSAY
IN TRIPLICATE
HLA REFERENCE
FINAL REPORT
NUMBER
TEST MATERIAL
MAILED
VALUE
10532 0 401
A105 BF7-142-3 BF9-135-2
11-08-88
1,300.00
10534 0 401
A167 OR117-89
11-15-88
1,300.00
TOTAL AMOUNT DUE UPON RECEIPT : :selected:
$2,600.00
OK-Ag Minnenyer
2-10-89
87147782
Dept. 8700
Acct. 4111
GENETICS AND IN VITRO TOXICOLOGY DEPARTMENT (301) 230-0001"""

_model = None
_tokenizer = None


def load_model():
    """Lazily load base model + LoRA adapter (4-bit) on first request."""
    global _model, _tokenizer
    if _model is not None:
        return _model, _tokenizer

    bnb_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=torch.float16,
        bnb_4bit_use_double_quant=True,
    )

    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_NAME, trust_remote_code=True)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    base_model = AutoModelForCausalLM.from_pretrained(
        BASE_MODEL_NAME,
        quantization_config=bnb_config,
        device_map="auto",
        trust_remote_code=True,
    )
    model = PeftModel.from_pretrained(base_model, ADAPTER_REPO)
    model.eval()

    _model, _tokenizer = model, tokenizer
    return model, tokenizer


def extract_json_from_text(text: str):
    """Best-effort JSON extraction from model output (handles markdown fences)."""
    if not text:
        return None
    text = text.strip()
    try:
        return json.loads(text)
    except Exception:
        pass
    text = text.replace("```json", "").replace("```", "").strip()
    try:
        return json.loads(text)
    except Exception:
        pass
    match = re.search(r"\{.*\}", text, re.DOTALL)
    if match:
        try:
            return json.loads(match.group(0))
        except Exception:
            return None
    return None

@spaces.GPU  # add this decorator
def run_extraction(ocr_text: str) -> str:
    if not ocr_text or not ocr_text.strip():
        return "Paste some OCR text first."

    model, tokenizer = load_model()

    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {
            "role": "user",
            "content": f"Extract invoice information from the OCR text below and "
            f"return only valid JSON.\n\nOCR_TEXT:\n{ocr_text.strip()}",
        },
    ]
    prompt = tokenizer.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

    with torch.no_grad():
        output_ids = model.generate(
            **inputs,
            max_new_tokens=768,
            do_sample=False,
            pad_token_id=tokenizer.pad_token_id,
        )

    generated = tokenizer.decode(
        output_ids[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True
    )

    parsed = extract_json_from_text(generated)
    if parsed is not None:
        return json.dumps(parsed, indent=2, ensure_ascii=False)
    return generated  # fall back to raw output if parsing fails


with gr.Blocks(title="BillStructAI — Invoice OCR to JSON") as demo:
    gr.Markdown(
        "# BillStructAI\n"
        "Paste noisy OCR text from an invoice or receipt below. "
        "A Qwen2.5-1.5B-Instruct model fine-tuned with QLoRA extracts it into "
        "structured JSON. [See the full writeup and evaluation results](https://github.com/trishpurkait/billstructai)."
    )
    with gr.Row():
        with gr.Column():
            ocr_input = gr.Textbox(
                label="OCR text",
                lines=12,
                value=EXAMPLE_OCR,
                placeholder="Paste raw OCR text here...",
            )
            submit_btn = gr.Button("Extract", variant="primary")
        with gr.Column():
            json_output = gr.Code(label="Extracted JSON", language="json", lines=20)

    submit_btn.click(fn=run_extraction, inputs=ocr_input, outputs=json_output)

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