import gradio as gr
import re

from sentence_transformers import SentenceTransformer
from sentence_transformers.util import cos_sim


codes = """001 - Vehicle Registration (New)
002 - Vehicle Registration Renewal
003 - Vehicle Ownership Transfer
004 - Vehicle De-registration
005 - Lost Registration Certificate Replacement
006 - Address Change Update
007 - Vehicle Data Correction
008 - Ownership Name Correction
009 - Vehicle Tax Payment
010 - Late Payment Fee Processing
011 - Vehicle Type/Specification Update
012 - BBNKB (Transfer Fee of Vehicle Ownership)
013 - STNK Issuance (Vehicle Registration Certificate)
014 - STNK Renewal
015 - Motor Vehicle Roadworthiness Inspection
016 - Plate Number Renewal
017 - Lost Plate Replacement
018 - Vehicle Export Registration
019 - Vehicle Import Registration
020 - Fleet Vehicle Registration
021 - Bulk Vehicle Registration Update
022 - Vehicle Insurance Assistance
023 - Vehicle Accident Reporting
024 - Vehicle Usage Change Declaration (e.g., personal to commercial)
025 - Legal Document Verification
026 - Ownership Transfer for Inherited Vehicle
027 - STNK Temporary Suspension
028 - Proof of Ownership Document Update
029 - Vehicle Ownership History Check
030 - Vehicle Tax Recalculation Request
031 - Tax Exemption Application (for special cases)
032 - Deceased Owner’s Vehicle Ownership Transfer""".split("\n")

undetected = "099 - Other/Undetected"


codes = """001 - Pendaftaran Kendaraan (Baru)
002 - Pembaruan Pendaftaran Kendaraan
003 - Alih Kepemilikan Kendaraan
004 - Pembatalan Pendaftaran Kendaraan
005 - Penggantian Sertifikat Pendaftaran Kendaraan yang Hilang
006 - Pembaruan Perubahan Alamat
007 - Koreksi Data Kendaraan
008 - Koreksi Nama Kepemilikan
009 - Pembayaran Pajak Kendaraan
010 - Proses Denda Keterlambatan Pembayaran
011 - Pembaruan Jenis/Spesifikasi Kendaraan
012 - Pembayaran Pajak Kendaraan Melalui E-Samsat
013 - Penerbitan STNK (Sertifikat Pendaftaran Kendaraan)
014 - Pembaruan STNK
015 - Pemeriksaan Kelayakan Jalan Kendaraan Bermotor
016 - Pembaruan Nomor Plat Kendaraan
017 - Penggantian Plat yang Hilang
018 - Pendaftaran Ekspor Kendaraan
019 - Pendaftaran Impor Kendaraan
020 - Pendaftaran Kendaraan Armada
021 - Pembaruan Pendaftaran Kendaraan Massal
022 - Bantuan Asuransi Kendaraan
023 - Pelaporan Kecelakaan Kendaraan
024 - Deklarasi Perubahan Penggunaan Kendaraan (misalnya, pribadi ke komersial)
025 - Verifikasi Dokumen Hukum
026 - Alih Kepemilikan Kendaraan Warisan
027 - Penangguhan Sementara STNK
028 - Pembaruan Dokumen Bukti Kepemilikan
029 - Pemeriksaan Riwayat Kepemilikan Kendaraan
030 - Permintaan Perhitungan Ulang Pajak Kendaraan
031 - Permohonan Pembebasan Pajak (untuk kasus khusus)
032 - Alih Kepemilikan Kendaraan Pemilik yang Meninggal""".split("\n")

undetected = "099 - Lainnya/Tidak Terdeteksi"

# model_id = "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2"
# model_id = "Alibaba-NLP/gte-multilingual-base"
model_id = "BAAI/bge-m3"
model_id = "sentence-transformers/paraphrase-multilingual-mpnet-base-v2"
model_id = "intfloat/multilingual-e5-small"
# model_id = "sentence-transformers/distiluse-base-multilingual-cased-v2"
model = SentenceTransformer(model_id, trust_remote_code=True)

codes_emb = model.encode([x[6:] for x in codes])

def respond(
    message,
    history: list[tuple[str, str]],
    threshold,
    is_multiple
):
    global codes_emb
    global undetected

    undetected_code = undetected[:3]

    if history and history[-1][-1][21:24] == undetected_code: 
        list_his = ""
        for his in history[::-1]:
            if his[-1][21:24] != undetected_code:
                break
            list_his = his[0] + "\n" + list_his

        message += "\n" + list_his
    
    # pattern = r'\b([A-Z]{1,2})\s?(\d{4})\s?([A-Z]{3})\b'
    # pattern = r'\b([A-Z]{1,2})\s?(\d{4})\s?([A-Z]{1,3})\b'
    pattern = r'\b([A-Za-z]{1,2})\s?(\d{4})\s?([A-Za-z]{1,3})\b'

    matches = re.findall(pattern, message)

    plate_numbers = ", ".join([" ".join(x) for i,x in enumerate(matches)]).upper()
    
    text_emb = model.encode(message)
    scores = cos_sim(codes_emb, text_emb)[:,0]

    if is_multiple:
        request_details = []
        request_numbers = []
        request_scores = []
        for i,score in enumerate(scores):
            if score > threshold:
                request_details.append(codes[i][6:])
                request_numbers.append(codes[i][:3])
                request_scores.append(str( round(score.tolist(), 3) ) )

        if not request_details:
            request_details.append(undetected[6:])
            request_numbers.append(undetected_code)

        request_numbers = "\n".join(request_numbers)
        request_details = "\n".join(request_details)
        request_scores = "\n".join(request_scores)

        return "Request code number:\n" + request_numbers + "\n\nRequest detail:\n" + request_details + f"\n\nConfidence score:\n{request_scores}" + "\n\nPlate numbers: " + plate_numbers


    s_max = scores.argmax()

    if scores[s_max] < threshold:
        # request_code = "033 - Other/Undetected"
        request_code = undetected
    else:
        request_code = codes[scores.argmax()]
        # "{:.2f}".format(a)

    return "Request code number: " + request_code[:3] + "\nRequest detail: " + request_code[6:] + f"\nConfidence score: {round(scores[s_max].tolist(),3)}" + "\nPlate numbers: " + plate_numbers

    # for val in history:
    #     if val[0]:
    #         messages.append({"role": "user", "content": val[0]})
    #     if val[1]:
    #         messages.append({"role": "assistant", "content": val[1]})

    # messages.append({"role": "user", "content": message})

    # response = ""

    # for message in client.chat_completion(
    #     messages,
    #     max_tokens=max_tokens,
    #     stream=True,
    #     temperature=temperature,
    #     top_p=top_p,
    # ):
    #     token = message.choices[0].delta.content

    #     response += token
    #     yield response


"""
For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
"""
# demo = gr.ChatInterface(
#     respond,
# )


with gr.Blocks() as demo:
    # Add header title and description
    gr.Markdown("# List of Request Numbers")
    gr.Markdown("<br>".join(codes) + "<br>" + undetected)
    gr.Markdown("# Valid License Plate Number Criteria:")
    gr.Markdown("(1-2 letters) (4 numbers) (1-3 letters)")

    # Add chat interface
    chat_interface = gr.ChatInterface(
        respond, 
        additional_inputs=[
            gr.Number(0.5, label="confidence threshold", show_label=True, minimum=0., maximum=1.0, step=0.1),
            gr.Checkbox(label="multiple", info="Allow multiple request code numbers"),
        ]
    )

if __name__ == "__main__":
    demo.launch()