import io
import os
import re
from concurrent.futures import ThreadPoolExecutor
from decimal import Decimal

import fitz
import pytesseract
from PIL import Image, ImageFilter, ImageOps

OCR_TARGET_WIDTH = 2500

ALL_FIELDS = (
    "box1_beneficiary_name", "box2_ssn",
    "box3_benefits_paid", "box4_benefits_repaid",
    "box5_net_benefits", "box6_fed_tax_withheld",
    "box7_address", "box8_claim_number",
    "box3_description", "box4_description"
)

DIGIT_FIX = str.maketrans({
    "O": "0", "o": "0", "I": "1", "l": "1", "|": "1",
    "]": "1", "[": "1", "B": "8", "a": "8", "S": "5",
})

SSN_RE = re.compile(r"(?<!\d)(\d{3})[\s-]?(\d{2})[\s-]?(\d{4})(?!\d)")

def clean(value):
    if value is None:
        return None
    value = re.sub(r"\s+", " ", str(value)).strip()
    return value or None

def _looks_scrambled(text):
    if not text or len(text) < 80:
        return False
    stripped = re.sub(r"\s+", "", text)
    if not stripped:
        return False
    unusual = sum(1 for c in stripped if c in "*()&#|\\")
    unusual_ratio = unusual / len(stripped)
    dollar_28 = len(re.findall(r"\b28\b", text))
    if unusual_ratio > 0.06 and dollar_28 >= 2:
        return True
    if unusual_ratio > 0.12:
        return True
    return False

def page_image(page, document):
    images = page.get_images(full=True)
    if len(images) == 1:
        try:
            info = document.extract_image(images[0][0])
            image = Image.open(io.BytesIO(info["image"])).convert("L")
            if image.width >= 600:
                return image
        except Exception:
            pass
    pix = page.get_pixmap(matrix=fitz.Matrix(300 / 72, 300 / 72), alpha=False)
    return Image.open(io.BytesIO(pix.tobytes("png"))).convert("L")

def prepare_for_ocr(image, sharpen):
    scale = OCR_TARGET_WIDTH / image.width
    if scale > 1.05:
        image = image.resize((round(image.width * scale), round(image.height * scale)), Image.Resampling.LANCZOS)
    image = ImageOps.autocontrast(image)
    if sharpen:
        image = image.filter(ImageFilter.UnsharpMask(2, 150, 3))
    return image

def ocr_image(image):
    text = pytesseract.image_to_string(image, lang="eng", config="--psm 6")
    if len(text.strip()) < 10:
        text = pytesseract.image_to_string(image, lang="eng", config="--psm 3")
    return text

def extract_text(pdf_bytes):
    pages = []
    ocr_jobs = []
    document = fitz.open(stream=pdf_bytes, filetype="pdf")

    try:
        for index, page in enumerate(document):
            text = page.get_text("text", sort=True) or ""
            scrambled = _looks_scrambled(text)
            
            try:
                base = page_image(page, document)
            except Exception:
                base = None

            item = {
                "page": index + 1,
                "text": text,
                "text_alt": "",
                "width": round(page.rect.width, 2),
                "height": round(page.rect.height, 2),
                "source": "pdf_text",
                "image": base,
            }

            needs_ocr = (len(re.sub(r"\s+", "", text)) < 40 or scrambled)

            if needs_ocr:
                if scrambled:
                    item["pdf_text_raw"] = text
                    item["text"] = ""
                tesseract_cmd = pytesseract.pytesseract.tesseract_cmd
                if not os.path.isfile(tesseract_cmd):
                    item["ocr_error"] = "Tesseract executable was not found."
                else:
                    try:
                        item["source"] = "ocr"
                        for sharpen in (False, True):
                            ocr_jobs.append((index, sharpen, prepare_for_ocr(base, sharpen)))
                    except Exception as exc:
                        item["ocr_error"] = f"OCR failed on page {index + 1}: {exc}"

            pages.append(item)
    finally:
        document.close()

    if ocr_jobs:
        def run(job):
            index, sharpen, image = job
            try:
                return index, sharpen, ocr_image(image), None
            except Exception as exc:
                return index, sharpen, "", exc

        workers = max(2, min(4, os.cpu_count() or 2))
        with ThreadPoolExecutor(max_workers=workers) as pool:
            for index, sharpen, text, exc in pool.map(run, ocr_jobs):
                if exc is not None:
                    pages[index]["ocr_error"] = f"OCR failed on page {index + 1}: {exc}"
                elif sharpen:
                    pages[index]["text_alt"] = text
                else:
                    pages[index]["text"] = text

    return pages

def _parse_amount_or_none(val_str):
    if not val_str:
        return None
    val_str = val_str.strip()
    
    if re.search(r"\bN[O0]NE\b", val_str, re.I):
        return "NONE"
    
    val_str = re.sub(r"[^\d.,-]", "", val_str)
    if not val_str:
        return None
    val_str = val_str.replace(",", "")
    try:
        return round(float(val_str), 2)
    except ValueError:
        return None

def _extract_boxes(text):
    res = {"box3": None, "box4": None, "box5": None}
    
    top_text = re.split(r"DESCRIPTION\s+OF\s+AMOUNT", text, flags=re.I)[0]
    token_pattern = re.compile(r"\bN[O0]NE\b|(?:\$\s*)?\d{1,3}(?:,\d{3})*\.\d{2}", re.I)
    tokens = token_pattern.findall(top_text)
    
    if len(tokens) >= 1:
        res["box3"] = _parse_amount_or_none(tokens[0])
    if len(tokens) >= 2:
        res["box4"] = _parse_amount_or_none(tokens[1])
    if len(tokens) >= 3:
        res["box5"] = _parse_amount_or_none(tokens[2])
        
    return res

def _extract_box6(text):
    m = re.search(r"Box\s*6[^\n]*(?:Voluntary[^\n]*)?\n+(.*?)(?=\n\s*Box|\n\s*Form|$)", text, re.I | re.S)
    if m:
        val = clean(m.group(1))
        if val:
            if re.search(r"\bN[O0]NE\b", val, re.I):
                return "NONE"
            token = re.search(r"(?:\$\s*)?\d{1,3}(?:,\d{3})*\.\d{2}", val, re.I)
            if token:
                return _parse_amount_or_none(token.group(0))
    return None

def _extract_box7(text):
    m = re.search(r"Box\s*7[^\n]*\n(.*?)(?=Box\s*8|$)", text, re.I | re.S)
    if m:
        lines = [clean(l) for l in m.group(1).splitlines() if clean(l)]
        filtered = []
        for line in lines:
            line = re.sub(r"[<>\{\}\[\]\|\`\~].*$", "", line).strip()
            if not line: continue
            if re.search(r"(form|ssa|1099|noise)", line, re.I): continue
            filtered.append(line)
        return " ".join(filtered) if filtered else None
    return None

def _extract_box8(text, image=None):
    # 1. Target Image Crop for Box 8 area (lower-middle section of SSA-1099 forms)
    if image is not None:
        try:
            width, height = image.size
            # Crop specifically around where Box 8 usually resides
            crop = image.crop((int(width * 0.05), int(height * 0.50), int(width * 0.95), int(height * 0.85)))
            crop = crop.resize((crop.width * 3, crop.height * 3), Image.Resampling.LANCZOS)
            crop = ImageOps.autocontrast(crop)
            
            raw = pytesseract.image_to_string(crop, config="--psm 6")
            # Look for 9 digits followed by letters/numbers (e.g. 073-84-8705D)
            m = re.search(r"(\d{3})[\s-]?(\d{2})[\s-]?(\d{4})([A-Z0-9]+)", raw)
            if m:
                return f"{m.group(1)}-{m.group(2)}-{m.group(3)}{m.group(4)}"
        except Exception:
            pass

    # 2. Fallback text scan
    for pattern in [
        r"(?:Box\s*8|Claim\s*Number)[\s\S]{0,100}?(\d{3})[\s-]?(\d{2})[\s-]?(\d{4})([A-Z0-9]+)",
        r"\b(\d{3})[\s-]?(\d{2})[\s-]?(\d{4})([A-Z0-9]+)\b"
    ]:
        m = re.search(pattern, text, re.I)
        if m:
            return f"{m.group(1)}-{m.group(2)}-{m.group(3)}{m.group(4)}"

    return None

def _extract_descriptions(text):
    box3_desc = "NONE"
    box4_desc = "NONE"
    
    # Box 3 Description
    m3 = re.search(
        r"DESCRIPTION\s+OF\s+AMOUNT\s+IN\s+BOX\s+3[^\n]*\n(.*?)(?=DESCRIPTION\s+OF\s+AMOUNT\s+IN\s+BOX\s+4|\bBox\s*6\b|$)",
        text,
        re.I | re.S
    )
    if m3:
        lines = [clean(l) for l in m3.group(1).splitlines() if clean(l)]
        filtered = []
        for line in lines:
            line = re.sub(r"\bNOISE\b.*$", "", line, flags=re.I).strip()
            line = re.sub(r"[<>\{\}\[\]\|\`\~].*$", "", line).strip()
            if not line or re.fullmatch(r"N[O0]NE\.?", line, re.I): continue
            if re.search(r"Box\s*\d", line, re.I) or re.search(r"DESCRIPTION\s+OF", line, re.I): continue
            if re.search(r"(paid|benefit|check|deposit|medicare|tax|withheld|premium|addition|supplemental|retroactive)", line, re.I):
                filtered.append(line)
        if filtered:
            box3_desc = " | ".join(filtered)

    # Box 4 Description
    m4 = re.search(
        r"DESCRIPTION\s+OF\s+AMOUNT\s+IN\s+BOX\s+4[^\n]*\n(.*?)(?=\bBox\s*6\b|\bBox\s*7\b|$)",
        text,
        re.I | re.S
    )
    if m4:
        lines = [clean(l) for l in m4.group(1).splitlines() if clean(l)]
        filtered = []
        for line in lines:
            line = re.sub(r"\bNOISE\b.*$", "", line, flags=re.I).strip()
            line = re.sub(r"[<>\{\}\[\]\|\`\~].*$", "", line).strip()
            if not line: continue
            if re.search(r"Box\s*\d", line, re.I) or re.search(r"DESCRIPTION\s+OF", line, re.I): continue
            if re.search(r"(paid by check|direct deposit|benefits for|paid in)", line, re.I):
                continue
            filtered.append(line)
        if filtered:
            box4_desc = " | ".join(filtered)
            
    return box3_desc, box4_desc

def _text_after(text, patterns, min_len=3):
    for pattern in patterns:
        for match in re.finditer(pattern, text, re.I | re.M):
            tail = text[match.end():match.end() + 200]
            for line in tail.splitlines():
                candidate = clean(line)
                if not candidate or len(candidate) < min_len:
                    continue
                if re.fullmatch(r"[\d.,\s$]+", candidate):
                    continue
                if re.search(r"\bbox\s*\d", candidate, re.I):
                    continue
                return candidate
    return None

def parse_ssa1099(pages, filename=None):
    full_text = "\n".join(
        ((p.get("text") or "") + "\n" + (p.get("text_alt") or ""))
        for p in pages if p
    )

    result = {key: None for key in ALL_FIELDS}
    warnings = []

    result["box1_beneficiary_name"] = _text_after(full_text, [
        r"Box\s*1\b[^\n]*", r"\b1\.?\s*Name\b[^\n]*",
        r"Name\s+of\s+(?:the\s+)?beneficiary[:\s]+",
    ])

    box2 = None
    for pattern in [
        r"Box\s*2\b[^\n]{0,120}",
        r"\b2\.?\s*(?:Beneficiary[’']?s?\s+)?(?:Social\s+Security|SSN)[^\n]{0,80}",
        r"Beneficiary[’']?s?\s+(?:Social\s+Security\s+)?Number[^\n]{0,80}",
        r"Social\s+Security\s+Number[^\n]{0,80}",
    ]:
        for match in re.finditer(pattern, full_text, re.I | re.M):
            tail = full_text[match.end():match.end() + 100]
            m = re.search(r"\b(\d{3})[\s-]?(\d{2})[\s-]?(\d{4})\b", tail)
            if m:
                box2 = f"{m.group(1)}-{m.group(2)}-{m.group(3)}"
                break
        if box2: break
    if not box2:
        box2_match = SSN_RE.search(full_text)
        if box2_match:
            box2 = f"{box2_match.group(1)}-{box2_match.group(2)}-{box2_match.group(3)}"
    result["box2_ssn"] = box2

    boxes = _extract_boxes(full_text)
    result["box3_benefits_paid"] = boxes["box3"]
    result["box4_benefits_repaid"] = boxes["box4"]
    result["box5_net_benefits"] = boxes["box5"]
    
    result["box6_fed_tax_withheld"] = _extract_box6(full_text)
    result["box7_address"] = _extract_box7(full_text)
    
    img = pages[0].get("image") if pages else None
    result["box8_claim_number"] = _extract_box8(full_text, img)
    
    b3_desc, b4_desc = _extract_descriptions(full_text)
    result["box3_description"] = b3_desc
    result["box4_description"] = b4_desc

    # Reconciliation
    b3 = result["box3_benefits_paid"]
    b4 = result["box4_benefits_repaid"]
    b5 = result["box5_net_benefits"]

    def _num(v):
        return 0.0 if v == "NONE" else (v if isinstance(v, (int, float)) else None)

    v3, v4, v5 = _num(b3), _num(b4), _num(b5)

    if v3 is not None and v4 is not None and v5 is not None:
        if abs((v3 - v4) - v5) > 0.02:
            warnings.append(f"Box 5 ({v5}) does not equal Box 3 ({v3}) minus Box 4 ({v4}).")
    elif v3 is not None and v4 is None and v5 is None:
        result["box4_benefits_repaid"] = "NONE"
        result["box5_net_benefits"] = v3
    elif v3 is not None and v4 is not None and v5 is None:
        result["box5_net_benefits"] = round(v3 - v4, 2)
    elif v3 is not None and v5 is not None and v4 is None:
        result["box4_benefits_repaid"] = round(v3 - v5, 2)

    if not result["box1_beneficiary_name"]:
        warnings.append("Box 1 name could not be read; enter it manually.")
    if result["box2_ssn"] is None:
        warnings.append("Box 2 SSN could not be extracted.")
    if result["box3_benefits_paid"] is None:
        warnings.append("Box 3 benefits paid could not be read; enter it manually.")
    if result["box5_net_benefits"] is None:
        warnings.append("Box 5 net benefits could not be read; enter it manually.")
    if result["box8_claim_number"] is None:
        warnings.append("Box 8 Claim Number could not be read; enter it manually.")
    if any(p.get("source") == "ocr" for p in pages):
        warnings.append("OCR-derived values require review against the original PDF.")
    if any(p.get("ocr_error") for p in pages):
        warnings.append("One or more pages could not be OCR-processed.")
    if not full_text.strip():
        warnings.append("No readable PDF text was extracted.")

    result["source_page_count"] = len(pages)
    result["source_pages"] = [
        {
            "page": p["page"],
            "source": p.get("source", "unknown"),
            "text": p.get("text", ""),
            "text_alt": p.get("text_alt", ""),
            **({"pdf_text_raw": p["pdf_text_raw"]} if p.get("pdf_text_raw") else {}),
            **({"ocr_error": p["ocr_error"]} if p.get("ocr_error") else {}),
        }
        for p in pages
    ]

    result["parse_warnings"] = warnings
    result["auto_filled"] = []
    result["review_required"] = bool(warnings)
    result["form_type"] = "ssa1099"

    return result