#### To launch the script
# gradio wcag_validator_ui.py
# python wcag_validator_ui.py

import gradio as gr
import requests
from pathlib import Path
import sys
import pandas as pd

parent_dir = Path(__file__).parent.parent
sys.path.insert(0, str(parent_dir))
from dotenv import load_dotenv, find_dotenv
from dependences.utils import (
    call_API_urlibrequest,
    create_folder,
    db_persistence_startup,
    db_persistence_insert,
    return_from_env_valid,
)
from dependences_ui.utils import *
import logging
import time
import json
import urllib.request
import urllib.parse
import os
import sqlite3


WCAG_VALIDATOR_RESTSERVER_HEADERS = [("Content-Type", "application/json")]


def process_dataframe(db_path, url, updated_df, user_state={}):

    print("Processing dataframe to adjust columns...")
    column_rating_name = "User Assessment for LLM Proposal"

    # Get the assessment column
    try:
        updated_df[column_rating_name] = updated_df[column_rating_name].astype(int)
    except ValueError:
        return "Error: User Assessment for LLM Proposal must be an integer"

    if (updated_df[column_rating_name] < 1).any() or (
        updated_df[column_rating_name] > 5
    ).any():
        return "Error: User Assessment for LLM Proposal must be between 1 and 5"

    dataframe_json = updated_df.to_json(orient="records")
    connection_db = sqlite3.connect(db_path)
    json_user_str = json.dumps({"username": user_state["username"]}, ensure_ascii=False)
    try:
        # insert after everything to keep datetime aligned
        db_persistence_insert(
            connection_db=connection_db,
            insert_type="wcag_user_llm_alttext_assessments",
            page_url=url,
            user=json_user_str,
            llm_model="",
            json_in_str=dataframe_json,  # to improve
            json_out_str="done via UI",
            table="wcag_user_assessments",
        )
    except Exception as e:
        print("Error inserting user assessment into database:", str(e))
    finally:
        if connection_db:
            connection_db.close()
        return "User assessment saved successfully!"


def load_images_from_json(json_input):
    """Extract URLs and alt text from JSON and create HTML gallery"""
    try:
        data = json_input

        if "images" not in data or not data["images"]:
            return "No images found in JSON", ""

        images = data["images"]
        info_text = f"Found {len(images)} image(s)"
        print(f"Found {len(data['images'])} image(s)")

        # Create HTML gallery with checkboxes and assessment forms
        html = """
        <style>
            .image-gallery {
                display: grid;
                grid-template-columns: repeat(auto-fill, minmax(250px, 1fr));
                gap: 20px;
                padding: 20px;
            }
            .image-card {
                border: 2px solid #e0e0e0;
                border-radius: 8px;
                padding: 10px;
                background: white;
            }
            .image-card:has(input[type="checkbox"]:checked) {
                border-color: #2196F3;
                background: #a7c1c1;
            }
            .image-card img {
                width: 100%;
                height: 200px;
                object-fit: scale-down;
                border-radius: 4px;
            }
            .image-info {
                margin-top: 10px;
            }
            .checkbox-label {
                display: flex;
                align-items: center;
                gap: 8px;
                cursor: pointer;
                font-weight: 500;
            }
            .image-checkbox {
                width: 18px;
                height: 18px;
                cursor: pointer;
                accent-color: #2196F3;
            }
            .alt-text {
                font-size: 14px;
                color: #666;
                margin-top: 5px;
            }
            .assessment-panel {
                display: none;
                margin-top: 15px;
                padding: 10px;
                background: #7896b9;
                border-radius: 4px;
                border: 1px solid #2196F3;
            }
            .assessment-panel.visible {
                display: block;
            }
            .form-group {
                margin: 10px 0;
            }
            .form-group label {
                display: block;
                font-weight: 500;
                margin-bottom: 5px;
                font-size: 13px;
            }
            
            .radio-container {
                display: flex;
                gap: 15px;
                align-items: center;
            }

            .radio-option {
                display: flex;
                align-items: center;
                gap: 5px;
                cursor: pointer;
            }

            .radio-label {
                font-weight: 500;
            }
            textarea {
                width: 100%;
                padding: 8px;
                border: 1px solid #ccc;
                border-radius: 4px;
                font-size: 13px;
                font-family: inherit;
                resize: vertical;
            }
        </style>
        <div class="image-gallery">
        """

        for idx, img_data in enumerate(images):
            url = img_data.get("url", "")
            alt_text = img_data.get("alt_text", "No description")

            html += f"""
            <div class="image-card">
                <img src="{url}" alt="{alt_text}" loading="lazy" onerror="this.src='data:image/svg+xml,%3Csvg xmlns=%22http://www.w3.org/2000/svg%22 width=%22200%22 height=%22200%22%3E%3Crect fill=%22%23ddd%22 width=%22200%22 height=%22200%22/%3E%3Ctext x=%2250%25%22 y=%2250%25%22 text-anchor=%22middle%22 dy=%22.3em%22 fill=%22%23999%22%3EImage not found%3C/text%3E%3C/svg%3E'">
                <div class="image-info">
                    <label class="checkbox-label">
                        <input type="checkbox" class="image-checkbox" data-imgurl="{url}" data-index="{idx}" 
                               onchange="
                                   const panel = document.getElementById('panel-{idx}');
                                   const checkedCount = document.querySelectorAll('.image-checkbox:checked').length;
                                   if (this.checked) {{
                                       if (checkedCount > 3) {{
                                           this.checked = false;
                                           alert('Maximum 3 images can be selected!');
                                           return;
                                       }}
                                       panel.classList.add('visible');
                                   }} else {{
                                       panel.classList.remove('visible');
                                   }}
                               ">
                        Select #{idx + 1}
                    </label>
                    <div class="alt-text">Current alt_text: {alt_text}</div>
                    
                    <div id="panel-{idx}" class="assessment-panel">
                        <div class="form-group">
                            <label>Rate current alt-text:</label>
                            <div class="radio-container">
                            <label class="radio-option">
                                <input type="radio" name="assessment-{idx}" value="1" data-index="{idx}">
                                <span class="radio-label">1</span>
                            </label>
                            <label class="radio-option">
                                <input type="radio" name="assessment-{idx}" value="2" data-index="{idx}">
                                <span class="radio-label">2</span>
                            </label>
                            <label class="radio-option">
                                <input type="radio" name="assessment-{idx}" value="3" data-index="{idx}" checked>
                                <span class="radio-label">3</span>
                            </label>
                            <label class="radio-option">
                                <input type="radio" name="assessment-{idx}" value="4" data-index="{idx}">
                                <span class="radio-label">4</span>
                            </label>
                            <label class="radio-option">
                                <input type="radio" name="assessment-{idx}" value="5" data-index="{idx}">
                                <span class="radio-label">5</span>
                            </label>
                        </div>
                        </div>
                        <div class="form-group">
                            <label>New alt-text:</label>
                            <textarea class="new-alt-text" data-index="{idx}" rows="3" placeholder="Enter improved alt-text...">{alt_text}</textarea>
                        </div>
                    </div>
                    
                    <input type="hidden" class="original-alt" data-index="{idx}" value="{alt_text}" />
                </div>
            </div>
            """
            # info_text += f"✓ Image {idx+1} alt_text: {alt_text}\n"
        html += "</div>"

        return info_text, html

    except json.JSONDecodeError as e:
        return f"Error: Invalid JSON format - {str(e)}", ""
    except Exception as e:
        return f"Error: {str(e)}", ""


def load_llm_assessment_from_json(json_input):

    try:
        # Parse JSON input
        data = json_input

        if "mllm_validations" not in data or not data["mllm_validations"]:
            print("no mllm_validations found")
            return pd.DataFrame()

        info_text = f"Assessment done on {len(data['mllm_validations']['mllm_alttext_assessments'])} image(s)\n\n"
        print(
            f"Assessment done on {len(data['mllm_validations']['mllm_alttext_assessments'])} image(s)"
        )

        data_frame = []
        for idx, img_data in enumerate(
            data["mllm_validations"]["mllm_alttext_assessments"], 1
        ):

            original_alt_text_assessment = img_data["mllm_response"].get(
                "original_alt_text_assessment", "No description"
            )
            new_alt_text = img_data["mllm_response"].get(
                "new_alt_text", "No description"
            )
            alt_text_original = img_data.get("alt_text", "No alt_text provided")

            data_frame.append(
                {
                    "Original Alt Text": alt_text_original,
                    "LLM Assessment": original_alt_text_assessment,
                    "LLM Proposed Alt Text": new_alt_text,
                }
            )

        df = pd.DataFrame(data_frame)
        return df

    except json.JSONDecodeError as e:
        return f"Error: Invalid JSON format - {str(e)}", []
    except Exception as e:
        return f"Error: {str(e)}", []


def make_alttext_llm_assessment_api_call(
    url,
    selected_images_json=[],
    db_path=None,
    wcag_rest_server_url="http://localhost:8000",
    user_state={},
    number_of_images=30,
):

    print(
        f"Making API call for llm assessment for {url} to {wcag_rest_server_url}/wcag_alttext_validation"
    )
    selected_images = json.loads(selected_images_json) if selected_images_json else []
    # print("selected_images:", selected_images)

    if not selected_images or len(selected_images) == 0:
        info_text = "No images selected"
        print(info_text)
        return "LLM assessment not started", pd.DataFrame()

    # prepare data for insertion
    json_in_str = {}
    json_out_str = {}
    selected_urls = []
    selected_alt_text_original = []
    user_assessments = []
    user_new_alt_texts = []
    selected_image_id = []
    user_assessments_llm_proposal = []
    for img in selected_images:
        selected_urls.append(img["image_url"])
        selected_alt_text_original.append(img["original_alt_text"])
        user_assessments.append(img["assessment"])
        user_new_alt_texts.append(img["new_alt_text"])
        selected_image_id.append(
            int(img["image_index"]) + 1
        )  # add the id selected (+1 for index alignment)
        user_assessments_llm_proposal.append(3)  # default value for now
    json_in_str["images_urls"] = selected_urls
    json_in_str["images_alt_text_original"] = selected_alt_text_original
    json_out_str["user_assessments"] = user_assessments
    json_out_str["user_new_alt_texts"] = user_new_alt_texts
    json_in_str = json.dumps(json_in_str, ensure_ascii=False)
    json_out_str = json.dumps(json_out_str, ensure_ascii=False)
    json_user_str = json.dumps({"username": user_state["username"]}, ensure_ascii=False)
    connection_db = sqlite3.connect(db_path)
    # ---------

    try:

        response = call_API_urlibrequest(
            data={
                "page_url": url,
                "number_of_images": number_of_images,
                "context_levels": 5,
                "pixel_distance_threshold": 200,
                "save_images": "True",
                "save_elaboration": "True",
                "specific_images_urls": selected_urls,
            },
            url=wcag_rest_server_url + "/wcag_alttext_validation",
            headers=WCAG_VALIDATOR_RESTSERVER_HEADERS,
        )
        # return response
        info_dataframe = load_llm_assessment_from_json(response)

        # add the UI ids and other fields to to api response
        info_dataframe.insert(
            0, "Image #", selected_image_id
        )  # add the UI ids from to api response
        info_dataframe.insert(2, "User Assessment", user_assessments)

        info_dataframe.insert(3, "User Proposed Alt Text", user_new_alt_texts)
        info_dataframe["User Assessment for LLM Proposal"] = (
            user_assessments_llm_proposal
        )

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

    try:
        # insert after everything to keep datetime aligned
        db_persistence_insert(
            connection_db=connection_db,
            insert_type="wcag_user_alttext_assessments",
            page_url=url,
            user=json_user_str,
            llm_model="",
            json_in_str=json_in_str,
            json_out_str=json_out_str,
            table="wcag_user_assessments",
        )
    except Exception as e:
        print("Error inserting user assessment into database:", str(e))
    finally:
        if connection_db:
            connection_db.close()
        return "LLM assessment completed", info_dataframe


def make_image_extraction_api_call(
    url,
    number_of_images=30,
    wcag_rest_server_url="http://localhost:8000",
):
    print(
        f"Making API call for image_extraction for {url} to {wcag_rest_server_url}/extract_images"
    )
    try:

        response = call_API_urlibrequest(
            data={
                "page_url": url,
                "number_of_images": number_of_images,
            },
            url=wcag_rest_server_url + "/extract_images",
            headers=WCAG_VALIDATOR_RESTSERVER_HEADERS,
        )
        # return response
        info_text, gallery_images = load_images_from_json(response)

        return info_text, gallery_images
    except Exception as e:
        return {"error": str(e)}


# ------- Gradio Interface -------#


# Create Gradio interface
with gr.Blocks(theme=gr.themes.Glass(), title="WCAG AI Validator") as demo:

    env_path = find_dotenv(filename=".env")
    if env_path == "":
        print("env path not found: service starting with the default params values")
    _ = load_dotenv(env_path)  # read .env file
    db_path = return_from_env_valid("DB_PATH", "persistence/wcag_validator_ui.db")
    print("db_path:", db_path)
    wcag_rest_server_url = return_from_env_valid(
        "WCAG_REST_SERVER_URL", "http://localhost:8000"
    )

    default_urls = [
        "https://amazon.com",
        "https://ebay.com",
    ]
    url_list_str = return_from_env_valid("URL_LIST", json.dumps(default_urls))
    url_list = json.loads(url_list_str)

    print("wcag_rest_server_url:", wcag_rest_server_url)

    connection_db = db_persistence_startup(
        db_name_and_path=db_path, table="wcag_user_assessments"
    )
    print("Database connection reference available:", connection_db)
    connection_db.close()

    gr.Markdown("# WCAG AI Validator UI")

    # login section
    user_state = gr.State({"logged_in": False, "username": None})
    with gr.Accordion(label="Register & Login", open=True) as register_and_login:
        with gr.Column(visible=True) as login_section:
            gr.Markdown("## Login / Register")

            with gr.Tab("Login"):
                login_username = gr.Textbox(
                    label="Username", placeholder="Enter your username"
                )
                login_password = gr.Textbox(
                    label="Password", type="password", placeholder="Enter your password"
                )
                login_btn = gr.Button("Login", variant="primary")
                login_msg = gr.Textbox(label="Login Status", interactive=False)

            with gr.Tab("Register"):
                reg_username = gr.Textbox(
                    label="Username", placeholder="Choose a username"
                )
                reg_password = gr.Textbox(
                    label="Password",
                    type="password",
                    placeholder="Choose a password (min 6 characters)",
                )
                reg_confirm = gr.Textbox(
                    label="Confirm Password",
                    type="password",
                    placeholder="Confirm your password",
                )
                reg_btn = gr.Button("Register", variant="primary")
                reg_msg = gr.Textbox(label="Registration Status", interactive=False)

        with gr.Column(visible=False) as protected_section:

            content_display = gr.Textbox(
                label="Your account", lines=5, interactive=False
            )
            logout_btn = gr.Button("Logout", variant="stop")

    # end login section

    with gr.Tab("Alt Text Assessment", visible=False) as alttext_assessment:

        db_path_state = gr.State(value=db_path)  # Store path in State\
        wcag_rest_server_url_state = gr.State(value=wcag_rest_server_url)
        with gr.Row():
            with gr.Column():

                with gr.Row():
                    with gr.Column():
                        url_input = gr.Dropdown(
                            url_list,
                            value=url_list[0],
                            multiselect=False,
                            label="Select an URL",
                            info="Select an URL to load in iframe",
                        )
                        images_number = gr.Slider(
                            5,
                            100,
                            value=50,
                            step=5,
                            label="Max number of images to retrieve",
                            visible=False,
                        )
                    with gr.Column():

                        image_extraction_api_call_btn = gr.Button(
                            "Extract Images & Alt Texts", variant="primary"
                        )
                        alttext_api_call_btn = gr.Button(
                            "Start LLM Assessment",
                            variant="secondary",
                            interactive=False,
                        )
                        image_info_output = gr.Textbox(
                            label="Activity tracking", lines=1
                        )

        with gr.Row(visible=False) as alttext_results_row:

            # Use DataFrame for tabular output
            alttext_info_output = gr.DataFrame(
                headers=[
                    "Image #",
                    "Original Alt Text",
                    "User Assessment",
                    "User Proposed Alt Text",
                    "LLM Assessment",
                    "LLM Proposed Alt Text",
                    "User Assessment for LLM Proposal",
                ],
                label="LLM Assessment Results",
                wrap=True,  # Wrap text in cells
                interactive=True,
                scale=7,
            )
            with gr.Column():
                save_user_assessment_btn = gr.Button(
                    "Save Your Assessment",
                    variant="secondary",
                    interactive=True,
                    scale=1,
                )
                gr.Markdown(
                    "ℹ Info: to assess the LLM output, only the values ​​for the 'User Assessment for LLM Proposal' column need to be changed."
                )

        with gr.Row():

            gallery_html = gr.HTML(label="Image Gallery")

        image_extraction_api_call_btn.click(
            fn=lambda: ("", "", gr.update(visible=False), gr.Button(interactive=False)),
            inputs=[],
            outputs=[
                image_info_output,
                gallery_html,
                alttext_results_row,
                alttext_api_call_btn,
            ],
        ).then(
            make_image_extraction_api_call,
            inputs=[url_input, images_number, wcag_rest_server_url_state],
            outputs=[image_info_output, gallery_html],
        ).then(
            fn=lambda: gr.Button(interactive=True),
            inputs=[],
            outputs=[alttext_api_call_btn],
        )

        # Process selected images with JavaScript

        alttext_api_call_btn.click(
            fn=make_alttext_llm_assessment_api_call,
            inputs=[
                url_input,
                gallery_html,
                db_path_state,
                wcag_rest_server_url_state,
                user_state,
            ],
            outputs=[image_info_output, alttext_info_output],
            js="""
         (url_input,gallery_html) => {
            const checkboxes = document.querySelectorAll('.image-checkbox:checked');
            if (checkboxes.length === 0) {
                alert('Please select at least one image!');
                return [url_input,JSON.stringify([])];
            }
            if (checkboxes.length > 3) {
                alert('Please select maximum 3 images!');
                return [url_input,JSON.stringify([])];
            }
            const selectedData = [];
            
            checkboxes.forEach(checkbox => {
                const index = checkbox.dataset.index;
                const imageUrl = checkbox.dataset.imgurl;
                const originalAlt = document.querySelector('.original-alt[data-index="' + index + '"]').value;
                const assessment = document.querySelector('input[name="assessment-' + index + '"]:checked').value;
                console.log("assessment:",assessment)
                const newAltText = document.querySelector('.new-alt-text[data-index="' + index + '"]').value;
                
                selectedData.push({
                    image_index: index,
                    image_url: imageUrl,
                    original_alt_text: originalAlt,
                    assessment: parseInt(assessment),
                    new_alt_text: newAltText
                });
            });
            
            return [url_input,JSON.stringify(selectedData)];
        }
        """,
        ).then(
            fn=lambda: gr.update(visible=True),
            inputs=[],
            outputs=[alttext_results_row],
        )

        save_user_assessment_btn.click(
            fn=process_dataframe,
            inputs=[db_path_state, url_input, alttext_info_output, user_state],
            outputs=[image_info_output],
        )

    # placed here at the end to give full contents visibility to events
    # Event handlers
    login_btn.click(
        fn=login_user,
        inputs=[login_username, login_password, user_state],
        outputs=[
            login_msg,
            reg_msg,
            user_state,
            login_section,
            protected_section,
            alttext_assessment,
            register_and_login,
        ],
    ).then(fn=protected_content, inputs=[user_state], outputs=[content_display])

    reg_btn.click(
        fn=register_user,
        inputs=[reg_username, reg_password, reg_confirm],
        outputs=[login_msg, reg_msg, user_state],
    )

    logout_btn.click(
        fn=logout_user,
        inputs=[user_state],
        outputs=[
            login_msg,
            user_state,
            login_section,
            protected_section,
            alttext_assessment,
        ],
    )


if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)