import re
from typing import Dict, Any, List, Tuple, Optional
import numpy as np

from ..data.schema import DocumentRecord, OCRWord, CLASS_FIELD_SCHEMAS

INV_NUM_PATTERNS = [
    re.compile(r'(?i)\b(?:invoice\s*(?:number|num|no|#)?|inv(?:\s*#|\s*no|\.)?)\s*[:.\s#]*\s*([a-z0-9\-_/]+)'),
    re.compile(r'(?i)\b(?:order|po|bill)\s*(?:no|number|#)?\s*[:.\s#]*\s*([a-z0-9\-_/]+)')
]
DATE_PATTERN = re.compile(r'(\d{1,4}[/\-.]\d{1,2}[/\-.]\d{1,4})')
AMOUNT_PATTERNS = [
    re.compile(r'(?i)(?:total|amount\s*due|net\s*amount|balance|grand\s*total)[\s:$]*([0-9]+[.,][0-9]{2,3})'),
    re.compile(r'(?i)(?:total|amount|due)[\s:$]*([0-9]+[.,]?[0-9]*)'),
    re.compile(r'[$€£¥]\s*([0-9]+[.,][0-9]{2})')
]

class HybridFieldExtractor:
    """
    Intelligent Hybrid Field Extractor:
    Combines spatial bounding-box proximity heuristics, regex pattern matching,
    and token-level confidence scoring to extract structured fields and bounding boxes.
    """
    def __init__(self, doc_class: str):
        self.doc_class = doc_class
        self.schema_fields = CLASS_FIELD_SCHEMAS.get(doc_class, [])

    def _find_words_near(self, key_word: OCRWord, all_words: List[OCRWord], max_dx=350, max_dy=70) -> List[OCRWord]:
        kx0, ky0, kx1, ky1 = key_word.bbox
        candidates = []
        for w in all_words:
            if w == key_word:
                continue
            wx0, wy0, wx1, wy1 = w.bbox

            # 1. Directly to the right (same line)
            is_right = (wx0 >= kx0) and (wx0 - kx1 <= max_dx) and (abs(wy0 - ky0) <= 30)

            # 2. Directly below (within column)
            is_below = (wy0 >= ky1) and (wy0 - ky1 <= max_dy) and (abs(wx0 - kx0) <= 120)

            if is_right or is_below:
                dist = np.sqrt((wx0 - kx1)**2 + (wy0 - ky0)**2)
                candidates.append((dist, w))

        candidates.sort(key=lambda x: x[0])
        return [c[1] for c in candidates]

    def _union_bboxes(self, words: List[OCRWord]) -> Optional[List[int]]:
        if not words:
            return None
        x0 = min(w.bbox[0] for w in words)
        y0 = min(w.bbox[1] for w in words)
        x1 = max(w.bbox[2] for w in words)
        y1 = max(w.bbox[3] for w in words)
        return [int(x0), int(y0), int(x1), int(y1)]

    def extract_with_details(self, record: DocumentRecord) -> Tuple[Dict[str, Any], Dict[str, List[int]], Dict[str, float]]:
        words = record.ocr_words
        full_text = " ".join([w.text for w in words])
        extracted = {f: None for f in self.schema_fields}
        bboxes = {f: None for f in self.schema_fields}
        confidences = {f: 0.0 for f in self.schema_fields}

        if self.doc_class == "invoice":
            # 1. Invoice Number
            for pat in INV_NUM_PATTERNS:
                m = pat.search(full_text)
                if m:
                    cand = m.group(1).strip()
                    if cand.lower() not in ["number", "num", "no", "date", "due", "total", "ber", "page", "order"]:
                        extracted["invoice_number"] = cand
                        confidences["invoice_number"] = 0.92
                        break
            
            # Spatial fallback for invoice number
            if not extracted["invoice_number"]:
                for w in words:
                    if any(t in w.text.lower() for t in ["invoice", "inv", "number", "no.", "no:"]):
                        near = self._find_words_near(w, words, max_dx=250, max_dy=50)
                        val_words = [nw for nw in near if any(c.isalnum() for c in nw.text) and nw.text.lower() not in ["date", "tax", "total", "to", "number"]]
                        if val_words:
                            extracted["invoice_number"] = val_words[0].text
                            bboxes["invoice_number"] = val_words[0].bbox
                            confidences["invoice_number"] = 0.85
                            break

            # 2. Invoice Date
            for i, w in enumerate(words):
                m_date = DATE_PATTERN.search(w.text)
                if m_date:
                    extracted["invoice_date"] = m_date.group(1)
                    bboxes["invoice_date"] = w.bbox
                    confidences["invoice_date"] = 0.95
                    break
            
            if not extracted["invoice_date"]:
                # Check for "date" trigger
                for w in words:
                    if "date" in w.text.lower():
                        near = self._find_words_near(w, words)
                        for nw in near:
                            m_d = DATE_PATTERN.search(nw.text)
                            if m_d:
                                extracted["invoice_date"] = m_d.group(1)
                                bboxes["invoice_date"] = nw.bbox
                                confidences["invoice_date"] = 0.90
                                break
                        if extracted["invoice_date"]:
                            break

            # 3. Total Amount
            for w in words:
                if any(t in w.text.lower() for t in ["total", "due", "balance", "amount"]):
                    near = self._find_words_near(w, words, max_dx=350, max_dy=60)
                    for nw in near:
                        cleaned = re.sub(r'[^\d.-]', '', nw.text.replace(',', ''))
                        if cleaned and any(c.isdigit() for c in cleaned):
                            try:
                                val = float(cleaned)
                                if val > 0:
                                    extracted["total_amount"] = nw.text
                                    bboxes["total_amount"] = nw.bbox
                                    confidences["total_amount"] = 0.93
                                    break
                            except ValueError:
                                pass
                    if extracted["total_amount"]:
                        break

            if not extracted["total_amount"]:
                m_amt = AMOUNT_PATTERNS[0].search(full_text) or AMOUNT_PATTERNS[1].search(full_text)
                if m_amt:
                    extracted["total_amount"] = m_amt.group(1).strip()
                    confidences["total_amount"] = 0.88

            # 4. Seller / Vendor Name (typically top of document)
            top_words = [w for w in words if w.bbox[1] < 250 and len(w.text) > 2]
            if top_words:
                seller_tokens = [w for w in top_words[:4] if w.text.lower() not in ["invoice", "date", "number", "page", "to", "from", "bill"]]
                if seller_tokens:
                    extracted["seller_name"] = " ".join([w.text for w in seller_tokens])
                    bboxes["seller_name"] = self._union_bboxes(seller_tokens)
                    confidences["seller_name"] = 0.80

            # 5. Buyer Name
            for w in words:
                if any(t in w.text.lower() for t in ["bill to", "sold to", "customer", "attn", "client"]):
                    near = self._find_words_near(w, words, max_dx=300, max_dy=50)
                    if near:
                        buyer_tokens = [nw for nw in near[:3] if nw.text.lower() not in [":", "to", "name"]]
                        if buyer_tokens:
                            extracted["buyer_name"] = " ".join([nw.text for nw in buyer_tokens])
                            bboxes["buyer_name"] = self._union_bboxes(buyer_tokens)
                            confidences["buyer_name"] = 0.82
                            break

        elif self.doc_class == "receipt":
            # 1. Company Name (top header)
            top_words = sorted([w for w in words if w.bbox[1] < 200], key=lambda w: (w.bbox[1], w.bbox[0]))
            if top_words:
                company_tokens = [w for w in top_words[:3] if w.text.lower() not in ["receipt", "welcome", "tel", "phone", "date", "time"]]
                if company_tokens:
                    extracted["company_name"] = " ".join([w.text for w in company_tokens])
                    bboxes["company_name"] = self._union_bboxes(company_tokens)
                    confidences["company_name"] = 0.88

            # 2. Date
            for w in words:
                m_d = DATE_PATTERN.search(w.text)
                if m_d:
                    extracted["date"] = m_d.group(1)
                    bboxes["date"] = w.bbox
                    confidences["date"] = 0.94
                    break

            # 3. Address
            for w in words:
                if any(t in w.text.lower() for t in ["street", "st.", "ave", "road", "rd", "blvd", "suite", "city", "mall", "plaza"]):
                    near = self._find_words_near(w, words, max_dx=200, max_dy=30)
                    addr_tokens = [w] + near[:4]
                    extracted["address"] = " ".join([nw.text for nw in addr_tokens])
                    bboxes["address"] = self._union_bboxes(addr_tokens)
                    confidences["address"] = 0.80
                    break

            # 4. Total Amount
            for w in words:
                if any(t in w.text.lower() for t in ["total", "due", "cash", "amount", "paid", "subtotal", "charge"]):
                    near = self._find_words_near(w, words, max_dx=350, max_dy=50)
                    for nw in near:
                        cleaned = re.sub(r'[^\d.-]', '', nw.text.replace(',', ''))
                        if cleaned and any(c.isdigit() for c in cleaned):
                            try:
                                val = float(cleaned)
                                if val > 0:
                                    extracted["total_amount"] = nw.text
                                    bboxes["total_amount"] = nw.bbox
                                    confidences["total_amount"] = 0.95
                                    break
                            except ValueError:
                                pass
                    if extracted["total_amount"]:
                        break

            if not extracted["total_amount"]:
                # Try finding biggest currency/numeric number near bottom
                bottom_words = [w for w in words if w.bbox[1] > 300]
                for bw in reversed(bottom_words):
                    cleaned = re.sub(r'[^\d.-]', '', bw.text.replace(',', ''))
                    if cleaned and any(c.isdigit() for c in cleaned):
                        try:
                            val = float(cleaned)
                            if val > 0:
                                extracted["total_amount"] = bw.text
                                bboxes["total_amount"] = bw.bbox
                                confidences["total_amount"] = 0.85
                                break
                        except ValueError:
                            pass

        elif self.doc_class == "application_form":
            # 1. Applicant Name
            for w in words:
                if any(t in w.text.lower() for t in ["applicant", "employee", "name", "signature", "client"]):
                    near = self._find_words_near(w, words, max_dx=300, max_dy=40)
                    cand_tokens = [nw for nw in near[:3] if nw.text.lower() not in [":", "name", "date", "is", "of", "title"]]
                    if cand_tokens:
                        extracted["applicant_name"] = " ".join([nw.text for nw in cand_tokens])
                        bboxes["applicant_name"] = self._union_bboxes(cand_tokens)
                        confidences["applicant_name"] = 0.86
                        break

            # 2. Date
            for w in words:
                m_d = DATE_PATTERN.search(w.text)
                if m_d:
                    extracted["date"] = m_d.group(1)
                    bboxes["date"] = w.bbox
                    confidences["date"] = 0.95
                    break

            # 3. Address
            for w in words:
                if any(t in w.text.lower() for t in ["address", "location", "residence", "city", "state", "street"]):
                    near = self._find_words_near(w, words, max_dx=350, max_dy=50)
                    cand_tokens = [nw for nw in near[:4] if nw.text.lower() not in [":", "address", "city", "phone"]]
                    if cand_tokens:
                        extracted["address"] = " ".join([nw.text for nw in cand_tokens])
                        bboxes["address"] = self._union_bboxes(cand_tokens)
                        confidences["address"] = 0.83
                        break

            # 4. Other key value pairs
            kv_pairs = []
            for w in words:
                if ":" in w.text and len(w.text) > 2:
                    key_part = w.text.replace(":", "").strip()
                    near = self._find_words_near(w, words, max_dx=200, max_dy=30)
                    if near:
                        val_part = " ".join([nw.text for nw in near[:3]])
                        kv_pairs.append(f"{key_part}: {val_part}")
            if kv_pairs:
                extracted["other_key_value_pairs"] = "; ".join(kv_pairs[:3])
                confidences["other_key_value_pairs"] = 0.78

        # Fallback bounding box search if text was extracted via regex without explicit bbox
        for f, val in extracted.items():
            if val and not bboxes[f]:
                val_first = str(val).split()[0]
                for w in words:
                    if val_first.lower() in w.text.lower():
                        bboxes[f] = w.bbox
                        break

        return extracted, bboxes, confidences

    def extract(self, record: DocumentRecord) -> Dict[str, Any]:
        fields, _, _ = self.extract_with_details(record)
        return fields