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
033 - Other/Undetected""".split("\n")


model = SentenceTransformer('sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2')

def respond(
    message,
    history: list[tuple[str, str]],
):
    # messages = [{"role": "system", "content": system_message}]
    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]{3})\b'

    matches = re.findall(pattern, message)

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

    request_code = codes[scores.argmax()]

    return "Request code number: " + request_code[:3] + "\nRequest detail: " + request_code[6:] + "\n Plate 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,
)


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