import gradio as gr
import requests
import json

def get_answers(question, top_k_reader, top_k_retriever):
    headers = {
        'accept': 'application/json',
        'Content-Type': 'application/json',
    }
    data = {
        'question': question, 
        'num_answers': top_k_reader, 
        'num_docs': top_k_retriever
    }
    
    try:
        response = requests.post(
            'https://8080-01jr2stv02evq6hmj7k295vb70.cloudspaces.litng.ai/query',
            headers=headers,
            data=json.dumps(data))
        response.raise_for_status()
        result = response.json()
        
        output_html = ""
        
        if 'answer' in result and 'answers' in result['answer']:
            for each in result['answer']['answers']:
                title = each['meta'].get('title', 'No title')
                url = each['meta'].get('url', '').split(';')[0]
                context = each.get('context', '')
                answer_text = each.get('answer', '')
                
                # Ambil offset
                offsets = each.get('offsets_in_context', [{}])[0]
                offset_start = offsets.get('start', 0)
                offset_end = offsets.get('end', 0)
                
                # Validasi offset
                valid_offset = (
                    0 <= offset_start < offset_end <= len(context)
                )
                
                # Format highlight
                if valid_offset:
                    highlighted = (
                        f"{context[:offset_start]}"
                        f"<span style='background-color: #ffeb3b; padding: 2px; border-radius: 3px;'>"
                        f"{context[offset_start:offset_end]}"
                        f"</span>"
                        f"{context[offset_end:]}"
                    )
                else:
                    highlighted = context
                
                # Bangun HTML output
                output_html += f"""
                <div style="margin: 20px 0; padding: 15px; border: 1px solid #ddd; border-radius: 5px;">
                    <div style="display: flex; justify-content: space-between; margin-bottom: 10px;">
                        <h3 style="margin: 0;">
                            {title} 
                            <a href="{url}" target="_blank" style="text-decoration: none;">🔗</a>
                        </h3>
                        <div style="background: #e0e0e0; padding: 5px 10px; border-radius: 5px;">
                            Confidence: {int(each.get('score', 0)*100)}%
                        </div>
                    </div>
                    
                    <div style="color: #666; margin-bottom: 10px;">
                        📅 {each["meta"].get("publish_time", "Unknown")} | 
                        👤 {each["meta"].get("authors", "Unknown")}
                    </div>
                    
                    <div style="padding: 10px; background: #f8f9fa; border-radius: 3px;">
                        {highlighted}
                    </div>
                </div>
                """
        
        return output_html if output_html else "No answers found."
    
    except Exception as e:
        return f"Error: {str(e)}"

# UI Components
with gr.Blocks(title="COVID QA", theme=gr.themes.Default()) as demo:
    gr.Markdown("# COVID QA")
    
    with gr.Row():
        with gr.Column(scale=3):
            question = gr.Textbox(label="Question", value="What is COVID?")
            top_k_reader = gr.Slider(label="Max. number of answers", minimum=1, maximum=10, value=3)
            top_k_retriever = gr.Slider(label="Max. number of documents", minimum=1, maximum=10, value=3)
            submit_btn = gr.Button("Get Answers")
        
        with gr.Column(scale=7):
            output = gr.HTML(label="Answers")

    submit_btn.click(
        fn=get_answers,
        inputs=[question, top_k_reader, top_k_retriever],
        outputs=output
    )

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