import os
import gradio as gr
import json
import pandas as pd
import requests as req

# Retrieve HF space secrets
auth_key = os.getenv('AUTH_KEY')
api_url = os.getenv('API_URL')
api_port = os.getenv('API_PORT')
FEEDBACK_IP = os.getenv('FEEDBACK_IP')
FEEDBACK_PORT = os.getenv('FEEDBACK_PORT')
FEEDBACK_PATH = os.getenv('FEEDBACK_PATH')
API_KEY = os.getenv('API_KEY')

HEADERS = {
    'Content-Type': 'application/json'
}

# Define feedback function to send like/dislike feedback
def send_feedback(request_data, response_data, like_reaction, dislike_reaction):
    print("Sending feedback...", request_data, response_data, like_reaction, dislike_reaction)
    # Construct the feedback payload
    feedback_payload = {
        "tool_id": 3,
        "request": json.dumps(request_data),
        "result": json.dumps(response_data),
        "like": like_reaction,
        "dislike": dislike_reaction
    }
    headers = {
        'Content-Type': 'application/json',
        'x-api-key': API_KEY
    }
    try:
        # Construct feedback URL and send the POST request
        feedback_url = f"http://{FEEDBACK_IP}:{FEEDBACK_PORT}{FEEDBACK_PATH}"
        response = req.post(feedback_url, json=feedback_payload, headers=headers)
        response.raise_for_status()  # Raise an error for bad responses
        print("Feedback sent successfully.")
        return {"message": "Feedback sent successfully"}
    except req.RequestException as e:
        print("Error sending feedback:", e)
        return {"error": str(e)}

# Define feedback toggle functionality
def toggle_feedback(request_data, response_data, like_clicked, dislike_clicked):
    print("Toggling feedback...", like_clicked, dislike_clicked)

    # Determine feedback type
    like_reaction = True if like_clicked else False
    dislike_reaction = True if dislike_clicked else False

    # Send feedback to the backend
    feedback_response = send_feedback(request_data, response_data, like_reaction, dislike_reaction)

    # Return appropriate message based on the feedback response
    if 'error' in feedback_response:
        return f"Failed to send feedback: {feedback_response['error']}"
    else:
        return "Feedback sent successfully!"

def preprocess_and_flatten(json_results, mode, meta_fields=None):
    # Ensure 'meta_fields' is a list or set default fields
    if meta_fields is None:
        meta_fields = ['doc_id', 'details', 'domain']

    # Check if json_results is a valid dictionary
    if not isinstance(json_results, dict):
        print(f"Invalid JSON results: Expected a dictionary but got {type(json_results)}")
        return pd.DataFrame()  # Return an empty DataFrame if json_results is not a dictionary

    # Collect flattened data
    flattened_data = []

    # Mode-based logic
    if mode == 'news_analysis':
        # Handle 'claim_objects' for news analysis mode
        claim_objects = json_results.get('claim_objects', [])
        if isinstance(claim_objects, list):
            for item in claim_objects:
                flattened_data.append({
                    'doc_id': json_results.get('doc_id'),
                    'details': json_results.get('details'),
                    'domain': json_results.get('domain'),
                    'topic': item.get('topic', ''),
                    'claim': item.get('claim', ''),
                    'claimer': item.get('claimer', '')
                })

    elif mode == 'claim_verification':
        # Handle 'support', 'refute', 'no_info' for claim verification mode
        nested_fields = ['support', 'refute', 'no_info']
        for field in nested_fields:
            nested_items = json_results.get(field, [])
            if not isinstance(nested_items, list):
                continue

            # Loop over each item in the nested field and flatten
            for item in nested_items:
                flattened_data.append({
                    'doc_id': json_results.get('doc_id'),
                    'details': json_results.get('details'),
                    'category': field,  # Mark which category the item belongs to (support/refute/no_info)
                    'sentence': item.get('sentence', ''),
                    'doi': item.get('doi', '')
                })

    # Convert to DataFrame
    dataframe_results = pd.DataFrame(flattened_data)

    # Capitalize column names
    dataframe_results.columns = [col.capitalize() for col in dataframe_results.columns]

    # Rename columns at the end of the function, conditionally if they exist
    rename_columns = {}
    
    # Conditionally add renaming based on the mode and column existence
    if 'doc_id' in dataframe_results.columns:
        rename_columns['doc_id'] = 'DOC ID'
    
    if mode == 'claim_verification':
        if 'doi' in dataframe_results.columns:
            rename_columns['doi'] = 'DOI'
        if 'sentence' in dataframe_results.columns:
            rename_columns['sentence'] = 'Sentence'
    
    # Apply the renaming if there are any columns to rename
    if rename_columns:
        dataframe_results.rename(columns=rename_columns, inplace=True)

    return dataframe_results
    
# Define the functions to handle the inputs and outputs
def news_analysis(text):
    try:
        response = req.post(
            f"{api_url}:{api_port}/news_analysis",
            json={
                'doc_id': '1',
                'text': text,
                'auth_key': auth_key
            },
            headers=HEADERS
        )
        response.raise_for_status()
        # Prepare results for JSON output
        json_results = response.json()
        # Flatten 'claim_objects' field
        dataframe_results = preprocess_and_flatten(json_results, mode='news_analysis')
        return json_results, dataframe_results
    except Exception as e:
        results = {'error': str(e)}
        return results, pd.DataFrame()

def claim_verification(text):
    try:
        response = req.post(
            f"{api_url}:{api_port}/claim_verification",
            json={
                'doc_id': '1',
                'text': text,
                'auth_key': auth_key
            },
            headers=HEADERS
        )
        response.raise_for_status()
        # Prepare results for JSON output
        json_results = response.json()
        # Flatten 'support', 'refute', and 'no_info' fields
        dataframe_results = preprocess_and_flatten(json_results, mode='claim_verification')
        return json_results, dataframe_results
    except Exception as e:
        results = {'error': str(e)}
        return results, pd.DataFrame()

# Define reusable feedback and export binding function
def bind_feedback_buttons(like_button, dislike_button, json_output, feedback_message):
    like_button.click(
        toggle_feedback,
        inputs=[json_output, json_output, gr.Textbox(visible=False, value='True'), gr.Textbox(visible=False, value='False')],
        outputs=[feedback_message]
    )

    dislike_button.click(
        toggle_feedback,
        inputs=[json_output, json_output, gr.Textbox(visible=False, value='False'), gr.Textbox(visible=False, value='True')],
        outputs=[feedback_message]
    )

def bind_export_buttons(export_csv_button, export_json_button, table_output, json_output):
    export_csv_button.click(
        export_results,
        inputs=[table_output, gr.Textbox(visible=False, value='csv'), json_output],
        outputs=[gr.File()]
    )

    export_json_button.click(
        export_results,
        inputs=[table_output, gr.Textbox(visible=False, value='json'), json_output],
        outputs=[gr.File()]
    )

# export function for results
def export_results(results, export_type, original_json):
    print("Exporting results...", export_type)
    try:
        if export_type == 'csv':
            # Ensure results is a DataFrame before exporting
            try:
                if not isinstance(results, pd.DataFrame):
                    results = pd.DataFrame(results)
            except ValueError as e:
                print("Error converting results to DataFrame:", e)
                return gr.File(None), f"Error: Unable to convert results to DataFrame - {str(e)}"
            
            csv_file_path = "exported_results.csv"
            results.to_csv(csv_file_path, index=False)
            print("CSV export successful:", csv_file_path)
            return gr.File(csv_file_path)

        elif export_type == 'json':
            # Ensure original_json is serializable
            if not isinstance(original_json, (dict, list)):
                raise ValueError("Invalid data for JSON export")

            json_file_path = "exported_results.json"
            with open(json_file_path, "w") as f:
                json.dump(original_json, f, indent=4)
            print("JSON export successful:", json_file_path)
            return gr.File(json_file_path)

        else:
            print("Error: Unsupported export type or no data available.")
            return gr.File(None), "Error: Unsupported export type or no data available."
    except (IOError, ValueError) as e:
        print("Error during export:", e)
        return gr.File(None), f"Error: {str(e)}"

# CSS for styling the interface
common_css = """
.unpadded_box {
  display: none !important;
}

#like-dislike-container, #claim-like-dislike-container {
    display: flex;
    justify-content: flex-start;
    margin-top: 20px; /* Increased margin to add more space between rows */
    gap: 15px; /* Add gap between like and dislike buttons */
}

#like-btn, #dislike-btn, #like-claim-btn, #dislike-claim-btn, #export-csv-btn, #export-json-btn, 
#export-claim-csv-btn, #export-claim-json-btn, #submit-btn, #submit-claim-btn {
    background-color: #e0e0e0;
    font-size: 18px;
    border-radius: 8px;
    padding: 12px; /* Increased padding for better look and feel */
    margin: 10px; /* Added margin for spacing between buttons */
    max-width: 250px;
    cursor: pointer;
    border: 1px solid transparent;
    transition: background-color 0.3s, box-shadow 0.3s; /* Smooth hover transition */
}

#like-btn:hover, #dislike-btn:hover, #like-claim-btn:hover, #dislike-claim-btn:hover, 
#submit-btn:hover, #submit-claim-btn:hover, #export-csv-btn:hover, #export-json-btn:hover, 
#export-claim-csv-btn:hover, #export-claim-json-btn:hover {
    background-color: #d0d0d0;
    box-shadow: 0px 4px 8px rgba(0, 0, 0, 0.1); /* Add shadow on hover for depth effect */
}

.active {
    background-color: #c0c0c0;
    font-weight: bold;
    border-color: #000;
}

.feedback-message {
    font-size: 16px; /* Slightly larger for readability */
    color: #4CAF50;
    margin-top: 10px; /* Space between feedback message and buttons */
}

.gr-textbox, .gr-markdown {
    margin-top: 15px; /* Space between input elements and titles */
}

#export-container {
    margin-top: 20px; /* Add space above the export container */
    gap: 15px; /* Add gap between export buttons */
}

.output-container {
    margin-top: 30px; /* Add space above the output container */
}

.gr-row {
    margin-top: 20px; /* Spacing for each row */
}
"""

# Define the interface for the first tab (News Analysis)
with gr.Blocks(css=common_css) as news_analysis_mode:
    # Input fields for news analysis
    gr.Markdown("### News Analysis")
    gr.Markdown("Classify the domain of a news article and detect major claims.")
    news_text_input = gr.Textbox(lines=10, label="News Article Text", placeholder="Enter the news article text")
    news_submit_button = gr.Button("Submit", elem_id="submit-btn")

    # Group related elements in a single container
    with gr.Group(visible=False, elem_id="output-container") as output_container:
        # Output fields for displaying results
        table_output = gr.DataFrame(label="Table View", elem_id="table_view", interactive=False)
        json_view_output = gr.JSON(label="JSON View", elem_id="json_view")

        # Feedback buttons container for user reaction
        reaction_label = gr.Markdown("**Reaction**")
        with gr.Row(elem_id="like-dislike-container"):
            like_button = gr.Button("👍 Like", elem_id="like-btn")
            dislike_button = gr.Button("👎 Dislike", elem_id="dislike-btn")
            feedback_message = gr.Markdown("")

        # Export options container
        export_label = gr.Markdown("**Export Options**")
        with gr.Row(elem_id="export-container"):
            export_csv_button = gr.Button("📄 Export as CSV", elem_id="export-csv-btn")
            export_json_button = gr.Button("📝 Export as JSON", elem_id="export-json-btn")

    # Bind export buttons to export function for News Analysis mode
    bind_export_buttons(export_csv_button, export_json_button, table_output, json_view_output)

    # Bind submit button to analyze input function
    news_submit_button.click(
        news_analysis,
        inputs=[news_text_input],
        outputs=[json_view_output, table_output]  # Ensure both outputs are specified here
    ).then(
        lambda: gr.update(visible=True),  # Show entire container after the first request
        inputs=[],
        outputs=[output_container]
    )
    
    # Bind feedback buttons for News Analysis Mode
    bind_feedback_buttons(like_button, dislike_button, json_view_output, feedback_message)

# Define the interface for the second tab (Claim Verification)
with gr.Blocks(css=common_css) as claim_verification_mode:
    gr.Markdown("### Claim Verification")
    gr.Markdown("Verify claims made in a news article.")
    claim_text_input = gr.Textbox(lines=10, label="Claim Text", placeholder="Enter the claim text")
    claim_submit_button = gr.Button("Submit", elem_id="submit-claim-btn")

    # Group related elements in a single container
    with gr.Group(visible=False) as claim_output_container:
        table_claim_output = gr.DataFrame(label="Table View", elem_id="table_view_claim", interactive=False)
        json_claim_output = gr.JSON(label="JSON View", elem_id="json_view_claim")

        claim_reaction_label = gr.Markdown("**Reaction**")
        with gr.Row(elem_id="claim-like-dislike-container"):
            like_claim_button = gr.Button("👍 Like", elem_id="like-claim-btn")
            dislike_claim_button = gr.Button("👎 Dislike", elem_id="dislike-claim-btn")
            claim_feedback_message = gr.Markdown("")

        claim_export_label = gr.Markdown("**Export Options**")
        with gr.Row(elem_id="export-claim-container"):
            export_claim_csv_button = gr.Button("📄 Export as CSV", elem_id="export-claim-csv-btn")
            export_claim_json_button = gr.Button("📝 Export as JSON", elem_id="export-claim-json-btn")

    # Bind the submit button to the function for verifying the claim text
    claim_submit_button.click(
        claim_verification,
        inputs=[claim_text_input],
        outputs=[json_claim_output, table_claim_output]
    ).then(
        lambda: gr.update(visible=True),  # Show entire container after the first request
        inputs=[],
        outputs=[claim_output_container]
    )

    # Bind feedback buttons for Claim Verification Mode
    bind_feedback_buttons(like_claim_button, dislike_claim_button, json_claim_output, claim_feedback_message)

    # Bind export buttons to export function for Claim Verification Mode
    bind_export_buttons(export_claim_csv_button, export_claim_json_button, table_claim_output, json_claim_output)

# Combine the tabs into one interface
with gr.Blocks(css=common_css) as demo:
    gr.TabbedInterface([news_analysis_mode, claim_verification_mode], ["News Analysis", "Claim Verification"])

# Launch the interface
demo.queue().launch(server_name="0.0.0.0", server_port=7860)