import sys
import json
import re
import os
from difflib import SequenceMatcher
from rapidocr_onnxruntime import RapidOCR
from pdf2image import convert_from_bytes, convert_from_path

# Initialize OCR engine once
engine = RapidOCR()

def smart_spacing(ocr_result):
    """
    Handle output RapidOCR (tuple / list)
    Ambil teks OCR apa adanya
    """

    if not ocr_result:
        return ""

    # ✅ HANDLE RapidOCR tuple
    if isinstance(ocr_result, tuple):
        blocks = ocr_result[0]
    else:
        blocks = ocr_result

    if not blocks or not isinstance(blocks, list):
        return ""

    texts = []

    # RapidOCR format: [ [box, text, score], ... ]
    for item in blocks:
        if (
            isinstance(item, (list, tuple)) and
            len(item) >= 2 and
            isinstance(item[1], str)
        ):
            texts.append(item[1].strip())

    return "\n".join(texts)

def clean(v: str) -> str:
    if not v:
        return ""
    return v.lstrip(": ").strip()

def similarity(a, b):
    return SequenceMatcher(None, a, b).ratio()

def normalize(text: str) -> list[str]:
    t = text.upper()
    t = t.replace("：", ":")

    fixes = {
        "NAME": "NAMA",
        "NAMA:": "NAMA",
        "TEMPAL": "TEMPAT",
        "TEMPA": "TEMPAT",
        "TGLLAHIR": "TGL LAHIR",
        "TEMPAT/TGLLAHIR": "TEMPAT/TGL LAHIR",
        "TEMPAT/TGL": "TEMPAT/TGL LAHIR",
        "JENISKELAMIN": "JENIS KELAMIN",
        "JENISFELANUN": "JENIS KELAMIN",
        "JENIS KELAMIN": "JENIS KELAMIN",
        "GOLDARAH": "GOL DARAH",
        "GOL.DARAH": "GOL DARAH",
        "KEI/DESA": "KEL/DESA",
        "KELDESA": "KEL/DESA",
        "BERIAKU": "BERLAKU",
        "BERLAKUHINGGA": "BERLAKU HINGGA",
        "STATUSPERKAWINAN": "STATUS PERKAWINAN",
    }

    for k, v in fixes.items():
        t = t.replace(k, v)

    t = re.sub(r"[ \t]+", " ", t)
    lines = [l.strip() for l in t.split("\n") if l.strip()]

    return lines

def parse_ktp(text: str) -> dict:
    lines = normalize(text)

    data = {
        "provinsi": "",
        "kota": "",
        "nik": "",
        "nama": "",
        "tempat_lahir": "",
        "tanggal_lahir": "",
        "jenis_kelamin": "",
        "gol_darah": "",
        "alamat": "",
        "rt_rw": "",
        "kel_desa": "",
        "kecamatan": "",
        "agama": "",
        "status_perkawinan": "",
        "pekerjaan": "",
        "kewarganegaraan": "",
        "berlaku_hingga": "",
    }

    i = 0
    while i < len(lines):
        line = lines[i]

        # ================= HEADER =================
        if line.startswith("PROVINSI"):
            data["provinsi"] = clean(line.replace("PROVINSI", ""))

        elif line.startswith(("KOTA", "KABUPATEN")):
            data["kota"] = clean(
                line.replace("KOTA", "").replace("KABUPATEN", "")
            )

        # ================= NIK =================
        elif line == "NIK" and i + 1 < len(lines):
            # NIK ada di baris berikutnya
            next_line = lines[i + 1]
            # Extract 16 digit dari baris berikutnya
            nik_match = re.search(r"\d{16}", next_line)
            if nik_match:
                data["nik"] = nik_match.group(0)
                # Setelah NIK, biasanya langsung NAMA
                # Check if ada nama di baris setelah NIK
                if i + 2 < len(lines):
                    potential_nama = lines[i + 2]
                    # Jika bukan field label dan tidak ada angka banyak, kemungkinan nama
                    if not any(keyword in potential_nama for keyword in ["TEMPAT", "TGL", "JENIS", "ALAMAT", "AGAMA", "PEKERJAAN", "STATUS", "KEWARGANEGARAAN", "BERLAKU"]):
                        if not re.search(r"\d{3,}", potential_nama):  # Tidak ada 3+ digit berurutan
                            if not data["nama"]:  # Belum terisi
                                data["nama"] = clean(potential_nama)
            i += 1  # Skip next line karena sudah diproses

        elif re.fullmatch(r":?\d{16}", line):
            data["nik"] = clean(line)

        # ================= NAMA =================
        elif line == "NAMA" and i + 1 < len(lines):
            data["nama"] = clean(lines[i + 1])

        elif line.startswith("NAMA"):
            nama_value = clean(line.split(":", 1)[-1])
            if nama_value:
                data["nama"] = nama_value

        # ================= TTL (TEMPAT/TGL LAHIR) =================
        elif "TEMPAT" in line and "TGL" in line:
            # Format bisa: "TEMPAT/TGL LAHIR" atau "TEMPAL/Tgi Lahir" dst
            value = clean(line.split(":", 1)[-1])
            if not value and i + 1 < len(lines):
                # Value di baris berikutnya
                value = lines[i + 1]

            # Pattern: NAMATEMPAT, DD-MM-YYYY atau NAMATEMPAT.DD-MM-YYYY
            # Contoh: RUMBAIJAYA.18-07-1997 atau RUMBAI JAYA, 18-07-1997
            m = re.search(r"([A-Z][A-Z\s]+?)[\s\.,]+(\d{2}-\d{2}-\d{4})", value)
            if m:
                tempat = m.group(1).strip()
                # Clean up tempat lahir - remove trailing numbers/codes
                tempat = re.sub(r"\d+$", "", tempat).strip()
                data["tempat_lahir"] = tempat
                data["tanggal_lahir"] = m.group(2)

        # ================= JENIS KELAMIN =================
        elif "JENIS" in line and "KELAMIN" in line:
            jk_value = clean(line.replace("JENIS KELAMIN", "").replace("JENIS", "").replace("KELAMIN", ""))
            if jk_value and jk_value not in [":", ""]:
                data["jenis_kelamin"] = jk_value
            elif i + 1 < len(lines):
                potential_jk = lines[i + 1]
                if potential_jk in ["LAKI-LAKI", "PEREMPUAN", "LAKILAKI"]:
                    data["jenis_kelamin"] = potential_jk

        # ================= GOL DARAH =================
        elif "GOL" in line and "DARAH" in line:
            v = clean(line.split(":", 1)[-1])
            if v in {"A", "B", "AB", "O", "-"}:
                data["gol_darah"] = v if v != "-" else ""

        # ================= ALAMAT =================
        elif line == "ALAMAT" and i + 1 < len(lines):
            data["alamat"] = clean(lines[i + 1])

        elif line.startswith("ALAMAT"):
            addr = clean(line.replace("ALAMAT", ""))
            if addr:
                data["alamat"] = addr

        # ================= RT/RW =================
        elif line == "RT/RW" and i + 1 < len(lines):
            data["rt_rw"] = clean(lines[i + 1])

        elif line.startswith("RT/RW") or line.startswith("RTRW"):
            rtrw = clean(line.replace("RT/RW", "").replace("RTRW", ""))
            if rtrw:
                data["rt_rw"] = rtrw
        
        # Pattern RT/RW tanpa label: 001/002 atau 1/2
        elif re.fullmatch(r"\d{1,3}/\d{1,3}", line):
            if not data["rt_rw"]:  # Belum terisi
                data["rt_rw"] = line

        # ================= KEL/DESA =================
        elif "KEL" in line and "DESA" in line:
            kd = clean(line.replace("KEL/DESA", "").replace("KELDESA", ""))
            if kd:
                data["kel_desa"] = kd
            elif i + 1 < len(lines):
                data["kel_desa"] = clean(lines[i + 1])

        # ================= KECAMATAN =================
        elif line.startswith("KECAMATAN") or line == "KECAMATAN":
            if i + 1 < len(lines):
                data["kecamatan"] = clean(lines[i + 1])
            else:
                kec = clean(line.replace("KECAMATAN", ""))
                if kec:
                    data["kecamatan"] = kec

        # ================= AGAMA =================
        elif line == "AGAMA" and i + 1 < len(lines):
            data["agama"] = clean(lines[i + 1])
        
        elif line.startswith("AGAMA"):
            agama = clean(line.replace("AGAMA", ""))
            if agama:
                data["agama"] = agama

        # ================= STATUS PERKAWINAN =================
        elif "STATUS" in line and "PERKAWINAN" in line:
            # Ambil dari baris berikutnya jika ada
            if i + 1 < len(lines):
                status_value = clean(lines[i + 1])
                # Validasi jika bukan field label
                if status_value not in ["PEKERJAAN", "KEWARGANEGARAAN", "BERLAKU", "AGAMA"]:
                    data["status_perkawinan"] = status_value
            # Atau bisa inline
            status_inline = clean(line.replace("STATUS PERKAWINAN", "").replace("STATUSPERKAWINAN", ""))
            if status_inline and status_inline != ":":
                data["status_perkawinan"] = status_inline

        # ================= PEKERJAAN =================
        elif line == "PEKERJAAN" and i + 1 < len(lines):
            potential_pekerjaan = clean(lines[i + 1])
            # Pastikan bukan field label
            if potential_pekerjaan not in ["KEWARGANEGARAAN", "STATUS", "BERLAKU", "WNI", "INDONESIA"]:
                data["pekerjaan"] = potential_pekerjaan

        elif line.startswith("PEKERJAAN"):
            pek = clean(line.split(":", 1)[-1])
            if pek and pek not in ["KEWARGANEGARAAN"]:
                data["pekerjaan"] = pek

        # ================= KEWARGANEGARAAN =================
        elif "KEWARGANEGARAAN" in line:
            # Value di baris yang sama atau berikutnya
            kw_inline = clean(line.replace("KEWARGANEGARAAN", ""))
            if kw_inline and kw_inline != ":":
                data["kewarganegaraan"] = kw_inline
            elif i + 1 < len(lines):
                kw_next = clean(lines[i + 1])
                if kw_next not in ["BERLAKU", "HINGGA"]:
                    data["kewarganegaraan"] = kw_next
        
        # Fallback: WNI standalone
        elif line == "WNI":
            if not data["kewarganegaraan"]:
                data["kewarganegaraan"] = "WNI"

        # ================= BERLAKU HINGGA =================
        elif "BERLAKU" in line:
            if "HINGGA" in line:
                berlaku_val = clean(line.replace("BERLAKU HINGGA", "").replace("BERLAKUHINGGA", ""))
                if berlaku_val:
                    data["berlaku_hingga"] = berlaku_val
                elif i + 1 < len(lines):
                    data["berlaku_hingga"] = clean(lines[i + 1])

        i += 1
        
# ================= FALLBACK HEURISTIC =================

    # NAMA fallback - cari setelah NIK
    if not data["nama"] and data["nik"]:
        for idx, l in enumerate(lines):
            if data["nik"] in l and idx + 1 < len(lines):
                potential_nama = lines[idx + 1]
                # Pastikan bukan field label
                if not any(keyword in potential_nama for keyword in ["TEMPAT", "TGL", "JENIS", "ALAMAT", "AGAMA", "PEKERJAAN"]):
                    if not re.search(r"\d{3,}", potential_nama):
                        data["nama"] = clean(potential_nama)
                        break

    # TTL fallback (baris mana saja)
    if not data["tanggal_lahir"]:
        for l in lines:
            # Pattern lengkap dengan nama tempat
            m = re.search(r"([A-Z][A-Z\s]+?)[\s\.,]+(\d{2}-\d{2}-\d{4})", l)
            if m:
                tempat = m.group(1).strip()
                tempat = re.sub(r"\d+$", "", tempat).strip()
                # Ambil kata terakhir saja jika ada multiple words
                # Contoh: "RUMBAI JAYA" lebih baik dari "TEMPAL TGL LAHIR RUMBAI JAYA"
                if len(tempat.split()) > 3:
                    tempat = " ".join(tempat.split()[-2:])
                data["tempat_lahir"] = tempat
                data["tanggal_lahir"] = m.group(2)
                break

    # Jenis kelamin fallback
    if not data["jenis_kelamin"]:
        for l in lines:
            if l in {"LAKI-LAKI", "PEREMPUAN", "LAKILAKI"}:
                data["jenis_kelamin"] = l
                break

    # Gol darah fallback - scan semua baris
    if not data["gol_darah"]:
        for l in lines:
            # Standalone A, B, AB, O
            if l in {"A", "B", "AB", "O"}:
                data["gol_darah"] = l
                break
            # Atau setelah "GOL"
            if "GOL" in l:
                for gd in ["AB", "A", "B", "O"]:  # AB dulu karena lebih panjang
                    if gd in l:
                        data["gol_darah"] = gd
                        break

    # STATUS PERKAWINAN fallback
    if not data["status_perkawinan"] or data["status_perkawinan"] == "STATUS PERKAWINAN":
        for l in lines:
            if l in ["BELUM KAWIN", "KAWIN", "CERAI HIDUP", "CERAI MATI"]:
                data["status_perkawinan"] = l
                break

    # KEWARGANEGARAAN fallback
    if not data["kewarganegaraan"] or data["kewarganegaraan"] == "KEWARGANEGARAAN":
        for l in lines:
            if l in ["WNI", "WNA", "INDONESIA"]:
                data["kewarganegaraan"] = l
                break

    # Post fix
    if data["jenis_kelamin"] == "LAKILAKI":
        data["jenis_kelamin"] = "LAKI-LAKI"

    # ================= MERGE ADDRESS =================
    # Gabungkan alamat, rt_rw, kel_desa, kecamatan
    address_parts = []
    if data["alamat"]:
        address_parts.append(data["alamat"])
    if data["rt_rw"]:
        address_parts.append(f"RT/RW {data['rt_rw']}")
    if data["kel_desa"]:
        address_parts.append(f"Kel. {data['kel_desa']}")
    if data["kecamatan"]:
        address_parts.append(f"Kec. {data['kecamatan']}")
    
    if data["kota"]:
        address_parts.append(data["kota"])
    if data["provinsi"]:
        address_parts.append(data["provinsi"])

    if address_parts:
        data["alamat"] = ", ".join(address_parts)

    return data

def process_file(file_path):
    if file_path.lower().endswith('.pdf'):
        # Convert first page of PDF to image
        # verify poppler availability or just try
        try:
             images = convert_from_path(file_path)
             if not images:
                 return {"error": "Empty PDF"}
             # Process first page
             # Save to temp or convert to bytes. RapidOCR accepts numpy or path or bytes.
             # Easier to save temp for consistency with current flow or pass bytes if RapidOCR supports it (it accepts bytes)
             # But let's just save temp to be safe
             import tempfile
             with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as tmp:
                 images[0].save(tmp.name, "JPEG")
                 tmp_path = tmp.name
             
             result, _ = engine(tmp_path)
             os.remove(tmp_path)
        except Exception as e:
            return {"error": f"PDF Error: {str(e)}"}
    else:
        # Image
        result, _ = engine(file_path)

    if not result:
        return {"error": "No text detected"}

    text = smart_spacing(result)
    parsed = parse_ktp(text)
    
    return {
        "text": text,
        "ktp": parsed,
        "raw": result
    }

if __name__ == "__main__":
    if len(sys.argv) < 2:
        print(json.dumps({"error": "File path required"}))
        sys.exit(1)

    file_path = sys.argv[1]
    
    try:
        if not os.path.exists(file_path):
             print(json.dumps({"error": "File not found"}))
             sys.exit(1)

        data = process_file(file_path)
        print(json.dumps(data))
    except Exception as e:
        print(json.dumps({"error": str(e)}))