import os
import gradio as gr
from groq import Groq
import base64
import json

# Inisialisasi Groq dengan API key dari secrets
client = Groq(api_key=os.environ.get("GROQ_API_KEY"))

# Fungsi untuk mengonversi gambar ke base64
def encode_image(image_path):
    with open(image_path, "rb") as image_file:
        return base64.b64encode(image_file.read()).decode('utf-8')

# Fungsi untuk mengekstrak informasi dari gambar KTP menggunakan OCR
def extract_ktp_info(image):
    base64_image = encode_image(image)
    
    prompt = (
        "Ekstrak informasi berikut dari gambar KTP ini dan berikan dalam format JSON dengan urutan sebagai berikut: Provinsi, Kota, NIK, Nama, Tempat/Tgl Lahir, Jenis Kelamin, Gol. Darah, Alamat, RT/RW, Kel/Desa, Kecamatan, Agama, Status Perkawinan, Pekerjaan, Kewarganegaraan. Hapus duplikat dan berikan hanya informasi yang relevan."
    )

    try:
        completion = client.chat.completions.create(
            model="llama-3.2-90b-vision-preview",
            messages=[
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "text",
                            "text": prompt
                        },
                        {
                            "type": "image_url",
                            "image_url": {
                                "url": f"data:image/jpeg;base64,{base64_image}",
                            },
                        },
                    ],
                }
            ],
            response_format={"type": "json_object"},
        )
        
        result = completion.choices[0].message.content
        # Pembersihan hasil untuk menghindari duplikasi
        result_json = json.loads(result)
        return json.dumps(result_json, indent=4)

    except Exception as e:
        return f"Error: {str(e)}"

# Fungsi untuk mengekstrak NPWP
def extract_npwp_info(image):
    base64_image = encode_image(image)
    
    try:
        completion = client.chat.completions.create(
            model="llama-3.2-90b-vision-preview",
            messages=[
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "text",
                            "text": "Ekstrak nomor NPWP dari gambar ini dan berikan dalam format JSON."
                        },
                        {
                            "type": "image_url",
                            "image_url": {
                                "url": f"data:image/jpeg;base64,{base64_image}",
                            },
                        },
                    ],
                }
            ],
            response_format={"type": "json_object"},
        )

        result = completion.choices[0].message.content
        return result

    except Exception as e:
        return f"Error: {str(e)}"

def extract_info(image, doc_type):
    if doc_type == "KTP":
        return extract_ktp_info(image)
    elif doc_type == "NPWP":
        return extract_npwp_info(image)
    else:
        return "Dokumen tidak dikenali."

iface = gr.Interface(
    fn=extract_info,
    inputs=[
        gr.Image(type="filepath", label="Upload Gambar"),
        gr.Radio(["KTP", "NPWP"], label="Jenis Dokumen")
    ],
    outputs=gr.Textbox(label="Informasi Ekstraksi"),
    title="Ekstraksi Informasi KTP dan NPWP",
    description="Unggah gambar KTP atau NPWP untuk mengekstrak informasi.",
)

iface.launch(share=True)