import gradio as gr
import json
import os
import random
from datetime import datetime

# --- Configuration ---
DATA_PATH = '/home/mshahidul/readctrl/data/extracting_subclaim/extracted_subclaims_syn_data_with_gs_summary_en.json'
SAVE_PATH = '/home/mshahidul/readctrl/data/thresold_finding/annotated_subclaims_triplet.json'

with open(DATA_PATH, 'r') as f:
    data = json.load(f)

# --- Logic Functions ---
def load_example(index):
    if index >= len(data):
        return [
            gr.update(value="### 🎉 All Done!"), gr.update(value="You have completed all records."), 
            [], "0%", "0%", "0%", gr.update(choices=[], value=[]), 
            gr.update(choices=[], value=[]), gr.update(choices=[], value=[]), ""
        ]
    
    record = data[index]
    source_type = random.choice(["Full Original Text", "Gold Summary"])
    
    if source_type == "Full Original Text":
        text_content, subclaims = record['fulltext'], record['fulltext_subclaims']
    else:
        text_content, subclaims = record['summary'], record['summary_subclaims']
        
    source_info = f"### Instance: {index + 1}/{len(data)} | Source: **{source_type}**"
    
    return [
        source_info, text_content, subclaims, "0%", "0%", "0%", 
        gr.update(choices=subclaims, value=[]), 
        gr.update(choices=subclaims, value=[]), 
        gr.update(choices=subclaims, value=[]),
        "" # Clear warning box
    ]

def calc_pct_and_validate(low, inter, prof, total_list):
    if not total_list: return "0%", "0%", "0%", ""
    
    l_pct = (len(low)/len(total_list)) * 100
    i_pct = (len(inter)/len(total_list)) * 100
    p_pct = (len(prof)/len(total_list)) * 100
    
    warning = ""
    if not (l_pct <= i_pct <= p_pct):
        warning = "⚠️ **Hierarchy Warning:** Information density should be: Low ≤ Intermediate ≤ Proficient."
        
    return f"{l_pct:.1f}%", f"{i_pct:.1f}%", f"{p_pct:.1f}%", warning

def save_and_next(username, index, source_info, low_sel, int_sel, prof_sel, subclaims):
    if not username or username.strip() == "":
        gr.Warning("Please enter your name before submitting!")
        return [index] + load_example(index)

    now = datetime.now()
    timestamp = now.strftime("%Y-%m-%d %H:%M:%S")
    stype = "Full Original Text" if "Full Original Text" in source_info else "Gold Summary"
    
    result = {
        "annotator": username,
        "timestamp": timestamp,
        "index": index,
        "source_type": stype,
        "annotations": {
            "low": {"subclaims": low_sel, "pct": len(low_sel)/len(subclaims)},
            "intermediate": {"subclaims": int_sel, "pct": len(int_sel)/len(subclaims)},
            "proficient": {"subclaims": prof_sel, "pct": len(prof_sel)/len(subclaims)}
        }
    }
    
    existing = []
    if os.path.exists(SAVE_PATH):
        try:
            with open(SAVE_PATH, 'r') as f: existing = json.load(f)
        except: pass
    
    existing.append(result)
    with open(SAVE_PATH, 'w') as f: json.dump(existing, f, indent=4)
    
    gr.Info(f"Saved successfully at {timestamp}!")
    return [index + 1] + load_example(index + 1)

# --- UI Definition ---
with gr.Blocks(theme=gr.themes.Soft(), title="Medical Literacy Annotation Tool") as demo:
    index_state = gr.State(0)
    subclaim_list_state = gr.State([])
    with open("/home/mshahidul/readctrl/code/interface/instructions", "r") as f:
        instructions_text = f.read()
    gr.Markdown(instructions_text)
   
    with gr.Row():
        # LEFT SIDEBAR: Instructions
        with gr.Column(scale=1, variant="panel"):
            user_input = gr.Textbox(label="Annotator Name", placeholder="e.g., Shama", interactive=True)
            
            
            
            gr.HTML("<hr>")
            source_display = gr.Markdown("### Initializing...")
            text_viewer = gr.Textbox(label="Reference Text", interactive=False, lines=15)

        # RIGHT MAIN: Annotation Area
        with gr.Column(scale=2):
            hierarchy_warning = gr.Markdown(value="", visible=True)
            
            with gr.Row():
                with gr.Column():
                    gr.Markdown("### 🟢 Low")
                    low_pct = gr.Label(label="Coverage", value="0%")
                    low_check = gr.CheckboxGroup(label="Subclaims", choices=[])
                
                with gr.Column():
                    gr.Markdown("### 🟡 Intermediate")
                    int_pct = gr.Label(label="Coverage", value="0%")
                    int_check = gr.CheckboxGroup(label="Subclaims", choices=[])
                
                with gr.Column():
                    gr.Markdown("### 🔴 Proficient")
                    prof_pct = gr.Label(label="Coverage", value="0%")
                    prof_check = gr.CheckboxGroup(label="Subclaims", choices=[])

            submit_btn = gr.Button("Submit & Next Record", variant="primary", size="lg")

    # --- Event Handlers ---
    demo.load(load_example, [index_state], [source_display, text_viewer, subclaim_list_state, low_pct, int_pct, prof_pct, low_check, int_check, prof_check, hierarchy_warning])
    
    # Real-time update for percentages and hierarchy warning
    for check_sys in [low_check, int_check, prof_check]:
        check_sys.change(
            calc_pct_and_validate, 
            [low_check, int_check, prof_check, subclaim_list_state], 
            [low_pct, int_pct, prof_pct, hierarchy_warning]
        )

    submit_btn.click(
        save_and_next,
        [user_input, index_state, source_display, low_check, int_check, prof_check, subclaim_list_state],
        [index_state, source_display, text_viewer, subclaim_list_state, low_pct, int_pct, prof_pct, low_check, int_check, prof_check, hierarchy_warning]
    )

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