import re
from typing import Dict, Any, List
from ..data.schema import DocumentRecord, CLASS_FIELD_SCHEMAS

INV_NUM_REGEX = re.compile(r'(?i)(?:invoice\s*(?:no|num|number|#)?[:.\s]*|inv\s*#?\s*|no[:.\s]+)([a-z0-9\-_/]+)')
DATE_REGEX = re.compile(r'(\d{1,4}[/\-.]\d{1,2}[/\-.]\d{1,4})')
AMOUNT_REGEX = re.compile(r'(?i)(?:total|amount|due|balance|subtotal)[\s:$]*([0-9]+[.,][0-9]{2,3})')

class RuleBasedExtractor:
    """Baseline A: Rule-based regular expression and keyword heuristic extractor."""
    def __init__(self):
        pass

    def extract(self, record: DocumentRecord, doc_class: str = None) -> Dict[str, Any]:
        cls_name = doc_class or record.doc_class
        tokens = [w.text for w in record.ocr_words]
        full_text = " ".join(tokens)
        schema_fields = CLASS_FIELD_SCHEMAS.get(cls_name, [])
        extracted = {f: None for f in schema_fields}

        if cls_name == "invoice":
            # 1. Invoice Number
            m_num = INV_NUM_REGEX.search(full_text)
            if m_num:
                extracted["invoice_number"] = m_num.group(1).strip()

            # 2. Invoice Date
            m_date = DATE_REGEX.search(full_text)
            if m_date:
                extracted["invoice_date"] = m_date.group(1).strip()

            # 3. Total Amount
            m_amt = AMOUNT_REGEX.search(full_text)
            if m_amt:
                extracted["total_amount"] = m_amt.group(1).strip()

            # 4. Buyer / Seller heuristics
            m_to = re.search(r'(?i)(?:bill\s*to|sold\s*to|to)[:\s]+([A-Za-z0-9\s]{3,30})', full_text)
            if m_to:
                extracted["buyer_name"] = m_to.group(1).strip()

        elif cls_name == "receipt":
            # 1. Company Name: top token block
            if tokens:
                extracted["company_name"] = " ".join(tokens[:min(3, len(tokens))])

            # 2. Date
            m_date = DATE_REGEX.search(full_text)
            if m_date:
                extracted["date"] = m_date.group(1).strip()

            # 3. Total Amount
            m_amt = AMOUNT_REGEX.search(full_text)
            if m_amt:
                extracted["total_amount"] = m_amt.group(1).strip()

        elif cls_name == "application_form":
            # 1. Applicant Name
            m_name = re.search(r'(?i)(?:name|applicant|signature)[:\s]+([A-Za-z\s]{3,25})', full_text)
            if m_name:
                extracted["applicant_name"] = m_name.group(1).strip()

            # 2. Date
            m_date = DATE_REGEX.search(full_text)
            if m_date:
                extracted["date"] = m_date.group(1).strip()

            # 3. Address
            m_addr = re.search(r'(?i)(?:address|location|city)[:\s]+([A-Za-z0-9\s,]{5,40})', full_text)
            if m_addr:
                extracted["address"] = m_addr.group(1).strip()

        return extracted