app.py

import os

import gradio as gr
from core import RAG, parse_cas_input

rag = RAG()
auth_list = []
auth_env = os.getenv("AUTH_USERS", "")
for pair in auth_env.split(","):
    if ":" in pair:
        user, pwd = pair.split(":", 1)
        auth_list.append((user.strip(), pwd.strip()))
EXAMPLE_INPUTS = [
    "7664-39-3",
    "7664-39-3, 128-37-0",
    "50-00-0, 7664-93-9, 67-64-1",
]


def _run_pipeline(cas_numbers):
    result = rag.pipeline(cas_numbers)

    rows = []
    for r in result.results:
        rows.append([r.casNumber, r.chemicalName, r.status, r.reason])

    summary_parts = []
    for r in result.results:
        summary_parts.append(
            f"### {r.casNumber} — {r.chemicalName}\n"
            f"**Yêu cầu pháp lý:** {r.status}\n\n"
            f"**Cơ sở:** {r.reason}\n"
        )
    summary_md = "\n---\n".join(summary_parts)
    return rows, summary_md


def analyse_text(raw_input: str):
    if not raw_input or not raw_input.strip():
        raise gr.Error("Vui lòng nhập ít nhất một mã CAS.")

    cas_numbers = parse_cas_input(raw_input)
    if not cas_numbers:
        raise gr.Error(
            f"Không tìm thấy mã CAS hợp lệ trong '{raw_input}'. "
            "Định dạng đúng: 7664-39-3"
        )
    return _run_pipeline(cas_numbers)


def analyse_image(image_path: str):
    if image_path is None:
        raise gr.Error("Vui lòng upload một ảnh chứa mã CAS.")

    cas_numbers = rag.extract_cas_from_image(image_path)
    if not cas_numbers:
        raise gr.Error("Không tìm thấy mã CAS nào trong ảnh. Vui lòng thử ảnh khác.")

    cas_display = ", ".join(cas_numbers)
    rows, summary_md = _run_pipeline(cas_numbers)
    return cas_display, rows, summary_md


with gr.Blocks(
    title="Chemical & Precursor Declaration Checker",
    theme=gr.themes.Soft(primary_hue="blue", secondary_hue="sky"),
    css="""
        .main-header {text-align:center; margin-bottom:4px}
        .sub-header  {text-align:center; color:#555; margin-top:0}
        .result-table {height: 120px; overflow-y: auto}
    """,
) as demo:
    gr.Markdown(
        "<h1 class='main-header'>Chemical & Precursor Declaration Checker</h1>"
        "<p class='sub-header'>"
        "Tra cứu nghĩa vụ khai báo hóa chất & tiền chất theo quy định Việt Nam"
        "</p>"
    )

    with gr.Tabs():
        # ---- Tab 1: Nhập text ----
        with gr.TabItem("Nhập mã CAS"):
            with gr.Row():
                with gr.Column(scale=4):
                    cas_input = gr.Textbox(
                        label="Nhập mã CAS",
                        placeholder="VD: 7664-39-3, 128-37-0  (phân cách bằng dấu phẩy)",
                        lines=2,
                    )
                with gr.Column(scale=1, min_width=140):
                    text_btn = gr.Button("Tra cứu", variant="primary", size="lg")

            gr.Examples(examples=EXAMPLE_INPUTS, inputs=cas_input, label="Ví dụ")

            gr.Markdown("### Kết quả")
            text_table = gr.Dataframe(
                headers=["CAS", "Tên hóa chất", "Yêu cầu pháp lý", "Cơ sở pháp lý"],
                datatype=["str", "str", "str", "str"],
                interactive=False, wrap=True,
                elem_classes=["result-table"],
            )
            text_detail = gr.Markdown()

            text_btn.click(fn=analyse_text, inputs=cas_input, outputs=[text_table, text_detail])
            cas_input.submit(fn=analyse_text, inputs=cas_input, outputs=[text_table, text_detail])

        # ---- Tab 2: Upload ảnh ----
        with gr.TabItem("Upload ảnh"):
            with gr.Row():
                with gr.Column(scale=4):
                    img_input = gr.Image(
                        label="Upload ảnh chứa danh sách mã CAS",
                        type="filepath",
                    )
                with gr.Column(scale=1, min_width=140):
                    img_btn = gr.Button("Trích xuất & Tra cứu", variant="primary", size="lg")

            extracted_cas = gr.Textbox(label="Mã CAS trích xuất được", interactive=False)

            gr.Markdown("### Kết quả")
            img_table = gr.Dataframe(
                headers=["CAS", "Tên hóa chất", "Yêu cầu pháp lý", "Cơ sở pháp lý"],
                datatype=["str", "str", "str", "str"],
                interactive=False, wrap=True,
                elem_classes=["result-table"],
            )
            img_detail = gr.Markdown()

            img_btn.click(
                fn=analyse_image, inputs=img_input,
                outputs=[extracted_cas, img_table, img_detail],
            )

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