import html  # Added for escaping HTML
import json
import logging  # Added for status check logging
import os  # Added for environment variables

import gradio as gr
import numpy as np
from sentence_transformers import SentenceTransformer

# Added HfApi for endpoint check
from huggingface_hub import InferenceClient
from dotenv import load_dotenv  # Added for .env loading

import utils.interface_utils as interface_utils
import utils.llm_utils as llm_utils

# Load environment variables from .env file
load_dotenv()

# REMOVED Endpoint name constant (will be in llm_utils)
# LLM_ENDPOINT_NAME = "phi-4-max"

# REMOVED Endpoint Status Check Function

# --- Load Data and Models ---
# These should be loaded once when the app starts.
print("Loading data and models...")

# Load data from files
processed_docs = json.load(
    open("docs_passages_storage/processed_docs.json", encoding="utf-8")
)
passages = json.load(open("docs_passages_storage/passages.json", encoding="utf-8"))
doc_embeds = np.load("docs_passages_storage/passage_embeddings.npy")

# Load the embedding model - Force CPU
print("Loading embedding model on CPU...")
model_name = "Snowflake/snowflake-arctic-embed-l-v2.0"
embed_model = SentenceTransformer(model_name, device="cpu")
print("Embedding model loaded.")

# Initialize HF Client globally using env vars
hf_api_token = os.getenv("HF_TOKEN")
if not hf_api_token:
    print(
        "Warning: HF_TOKEN environment variable not set. Inference client might fail."
    )
hf_client = InferenceClient(token=hf_api_token)

print("Data and models loaded.")


# --- Main Gradio Function ---
def get_results(query: str, progress=gr.Progress(track_tqdm=True)) -> list:
    """Processes query, retrieves passages, processes each, formats output, tracks progress."""
    # Define placeholders for 10 outputs (9 results + 1 summary)
    if not query:
        # Return default values for all 10 output slots
        return ["-"] * 9 + ["### Summary\n-"]

    # --- Check Endpoint Status FIRST ---
    # Use token loaded previously for the client
    endpoint_status = llm_utils.check_endpoint_status(
        token=hf_api_token
    )  # Call function from llm_utils
    if endpoint_status["status"] == "error":
        error_message = endpoint_status["ui_message"]
        logging.error(f"Endpoint status error: {error_message}")
        progress(1, desc="Endpoint Error")
        # Display error in first result slot and summary slot
        return [
            "### Endpoint Error",
            f"<p style='color: red;'>{html.escape(error_message)}</p>",
            "_",
            "-",
            "-",
            "-",
            "-",
            "-",
            "-",
            "### Summary\n_Endpoint not ready._",
        ]

    print(f"Processing query: {query}")
    try:
        progress(0, desc="Retrieving relevant documents...")
        # Step 1: Retrieve top excerpts
        retrieved_excerpts = llm_utils.retrieve_passages(
            query=query,
            doc_embeds=doc_embeds,
            passages=passages,
            processed_docs=processed_docs,
            embed_model=embed_model,
            max_docs=3,
        )
        progress(0.1, desc="Retrieved documents.")  # Update progress after retrieval

        if not retrieved_excerpts:
            print("No passages retrieved.")
            progress(1, desc="No relevant passages found.")  # Update progress
            return [
                "### Document:\n-",
                "<p><i>No relevant passages found.</i></p>",
                "_",
            ] * 3 + [
                "### Summary\n_No passages found._"
            ]  # Add summary placeholder

        # Step 2: Process each excerpt individually
        processed_data = []
        num_excerpts = len(retrieved_excerpts)
        for i, excerpt in enumerate(retrieved_excerpts):
            # Update progress description before processing each excerpt
            progress(
                0.1 + (i * 0.8 / num_excerpts),
                desc=f"Processing excerpt {i + 1}/{num_excerpts}...",
            )
            # Pass the globally initialized hf_client
            processed_result = llm_utils.process_single_excerpt(
                i, excerpt, query, hf_client
            )
            processed_data.append({"excerpt": excerpt, **processed_result})

        progress(
            0.9, desc="Formatting results..."
        )  # Mark processing complete, start formatting

        # Step 3: Format results for Gradio output components
        final_outputs = [
            "### Document:\n-",
            "<p><i>No result</i></p>",
            "_No passage text_",
        ] * 3 + [
            "### Summary\n_Generating..._"
        ]  # Reset outputs + summary placeholder

        global_quote_counter = 0  # Initialize global counter
        snippets_for_llm_summary = []  # List to store formatted snippets for the LLM
        all_spans_with_hover_info = (
            {}
        )  # Dict to store spans per passage index: {0: [(s, e, info), ...], 1: ...}

        for i in range(min(len(processed_data), 3)):
            result = processed_data[i]
            excerpt = result["excerpt"]
            citations = result["citations"]
            parse_successful = result["parse_successful"]
            raw_error_response = result["raw_error_response"]
            passage_text = excerpt.get("passage_text", "")
            doc_url = excerpt.get("document_url", "#")

            # 1. Format Document URL Markdown
            doc_url_md = (
                f"### Document {i+1}:\n[{html.escape(doc_url)}]({html.escape(doc_url)})"
            )

            # 2. Format Quotes HTML
            quotes_html_parts = []
            if parse_successful and citations:
                quotes_html_parts.append(
                    "<p><strong>Relevant Quotes:</strong> (hover for details)</p>"
                )
                quotes_html_parts.append(
                    "<ul style='list-style-type: none; margin-left: 10px; padding-left: 0;'>"  # Use none for list type
                )
                for cit in citations:
                    global_quote_counter += 1  # Increment counter
                    quote_id = global_quote_counter
                    # Store id in citation dict (optional, but might be useful later)
                    cit["global_id"] = quote_id

                    quote = cit.get("quote", "N/A")
                    context = cit.get("context", "N/A")
                    relevance = cit.get("relevance", "N/A")
                    hover_text = f"Context: {html.escape(context, quote=True)}\nRelevance: {html.escape(relevance, quote=True)}"
                    # Update HTML to include the ID
                    quotes_html_parts.append(
                        f"<li style='margin-bottom: 5px;' title='{hover_text}'>[{quote_id}]: <i>{html.escape(quote)}</i></li>"
                    )

                    # Prepare hover text for the highlighted span in the passage
                    span_hover_text = f"Quote ID: {quote_id}\nContext: {context}\nRelevance: {relevance}"
                    # Get spans for this specific citation
                    citation_spans = cit.get("char_spans", [])
                    # Associate hover text with these spans for the current passage (index i)
                    if i not in all_spans_with_hover_info:
                        all_spans_with_hover_info[i] = []
                    for start, end in citation_spans:
                        all_spans_with_hover_info[i].append(
                            (start, end, span_hover_text)
                        )

                    # Add formatted snippet to list if parsing was successful
                    if (
                        parse_successful
                    ):  # Ensure we only add successfully parsed citations
                        snippets_for_llm_summary.append(
                            {
                                "id": quote_id,
                                "context": context,
                                "relevance": relevance,
                                "quote": quote,
                                "document_url": doc_url,  # Added document URL here
                            }
                        )
                quotes_html_parts.append("</ul>")
            elif not parse_successful and raw_error_response:
                quotes_html_parts.append(
                    "<p style='color: red;'><strong>Error parsing citations:</strong></p>"
                )
                # Limit error display length
                error_display = html.escape(raw_error_response[:1000]) + (
                    "..." if len(raw_error_response) > 1000 else ""
                )
                quotes_html_parts.append(
                    f"<pre style='background-color: #f8d7da; color: #721c24; padding: 5px; border-radius: 4px; white-space: pre-wrap; word-wrap: break-word;'><code>{error_display}</code></pre>"
                )
            else:
                quotes_html_parts.append("<p><i>No specific quotes identified.</i></p>")
            quotes_html = "".join(quotes_html_parts)
            if not quotes_html:
                quotes_html = "<p><i>No quotes processed.</i></p>"

            # 3. Format Passage Markdown using the collected spans with hover info
            spans_for_this_passage = all_spans_with_hover_info.get(
                i, []
            )  # Get spans for index i
            passage_md = interface_utils.generate_highlighted_markdown(
                passage_text, spans_for_this_passage
            )
            if not passage_md:
                passage_md = "_Passage text unavailable._"

            # Update the final_outputs list
            final_outputs[i * 3 + 0] = doc_url_md
            final_outputs[i * 3 + 1] = quotes_html
            final_outputs[i * 3 + 2] = passage_md

        # Step 4: Generate LLM summary
        progress(0.95, desc="Generating summary...")  # New progress step
        summary_text = "### Summary\n_Error generating summary._"  # Default error text
        if snippets_for_llm_summary:
            # Create a lookup for quote details by ID
            snippet_lookup = {s["id"]: s for s in snippets_for_llm_summary}

            # Pass the globally initialized hf_client
            summary_result = llm_utils.generate_summary_answer(
                snippets=snippets_for_llm_summary, query=query, hf_client=hf_client
            )
            if summary_result["parse_successful"]:
                summary_items = []
                for sentence_data in summary_result["answer_sentences"]:
                    sentence = sentence_data.get("sentence", "")
                    citation_ids = sentence_data.get("citations", [])

                    # Generate HTML links for citations
                    citation_links = []
                    for c_id in citation_ids:
                        # Look up snippet details
                        snippet_info = snippet_lookup.get(c_id)
                        if snippet_info:
                            url = snippet_info.get("document_url", "#")
                            quote_text = snippet_info.get("quote", "")
                            escaped_quote = html.escape(quote_text, quote=True)
                            link = f"<a href='{url}' title='{escaped_quote}' target='_blank'>[{c_id}]</a>"
                            citation_links.append(link)
                        else:
                            citation_links.append(
                                f"[{c_id}]"
                            )  # Fallback if ID not found

                    # Format citation string like [link1, link2]
                    citation_str = (
                        f" [{', '.join(citation_links)}]" if citation_links else ""
                    )
                    if sentence:
                        # Append sentence and linked citations
                        summary_items.append(f"{html.escape(sentence)}{citation_str}")
                if summary_items:
                    summary_text = "### Generated Answer\n" + " ".join(summary_items)
                else:
                    summary_text = (
                        "### Generated Answer\n_LLM did not generate any sentences._"
                    )
            else:
                # Display parsing error from summary LLM
                summary_text = (
                    "### Summary Generation Error\n"
                    + f"<p style='color: red;'>Could not generate summary:</p><pre style='background-color: #f8d7da; color: #721c24; padding: 5px; border-radius: 4px; white-space: pre-wrap; word-wrap: break-word;'><code>{html.escape(summary_result['raw_error_response'])}</code></pre>"
                )
        else:
            summary_text = "### Summary\n_No valid quotes found to generate summary._"

        final_outputs[9] = summary_text  # Add summary to the 10th slot

        progress(1, desc="Done!")  # Final progress update
        return final_outputs

    except Exception as e:
        print(f"Error in get_results: {e}")
        # Display error in the first result slot, update summary slot
        error_outputs = (
            [
                "### Error Processing Query",
                f"<p style='color: red;'>An unexpected error occurred: {html.escape(str(e))}</p>",
                "_Error details above._",
            ]
            + ["-"] * 6
            + ["### Summary\n_Error occurred._"]
        )  # Add summary placeholder for error
        progress(1, desc="Error!")
        return error_outputs


# --- Gradio Interface ---

# Define custom CSS for scrollable accordion content
custom_css = """
.scrollable-passage-content { 
    max-height: 25vh;  /* Or a fixed height like 200px */
    overflow-y: auto !important; /* Add !important to override potential conflicts */
    display: block; /* Ensure it behaves like a block element */
    padding: 10px; /* Add some padding */
    background-color: #f9f9f9; /* Match previous background */
    border: 1px solid #ddd; /* Match previous border */
    border-radius: 5px; /* Match previous radius */
    white-space: pre-wrap; /* Preserve whitespace */
    word-wrap: break-word; /* Wrap long words */
}
"""

with gr.Blocks(css=custom_css, theme=gr.themes.Soft(), title="🤗 Policy Docs QA") as demo:
    gr.Markdown("# 🤗 Policy Docs QA")
    gr.Markdown(
        "Ask a question about the loaded policy documents to retrieve relevant passages and quotes."
    )  # Added description

    with gr.Row():
        # Column 1: Input (scale=1)
        with gr.Column(scale=1):
            question_input = gr.Textbox(
                label="Question",
                placeholder="Enter your question...",
                lines=5,
            )
            # Add example questions
            example_questions = [
                "What role does replicable evaluation play in AI regulation?",
                "How does dataset transparency help address risks of accidents and abuse of AI systems?",
            ]
            gr.Examples(
                examples=example_questions,
                inputs=question_input,
                label="Example Questions",  # Optional label
            )
            submit_button = gr.Button("Get Answer")

        # Column 2: Results (scale=3)
        with gr.Column(scale=3):
            gr.Markdown("## Retrieved Passages")  # Changed heading

            result_outputs = []  # Rename list for clarity
            for i in range(3):  # Create 3 result sections
                with gr.Group():
                    gr.Markdown(f"**Result {i+1}**")
                    doc_url_md = gr.Markdown(
                        value="### Document:\n-", label=f"Document URL {i+1}"
                    )
                    quotes_html = gr.HTML(
                        value="<p><i>Quotes will appear here...</i></p>",
                        label=f"Quotes {i+1}",
                    )
                    with gr.Accordion(f"Full Passage Context {i+1}", open=False):
                        passage_md = gr.Markdown(
                            value="_Passage text will appear here..._",
                            label=f"Passage {i+1}",
                            elem_classes="scrollable-passage-content",  # Assign the class
                        )
                    result_outputs.extend([doc_url_md, quotes_html, passage_md])

        # Column 3: Summary (scale=2)
        with gr.Column(scale=2):
            gr.Markdown("## Summary")
            gr.Markdown("***⚠️ Warning:** The text below is generated by an LLM and might not accurately reflect the policy documents. Citation links are determined by the same LLM and provided for conveninece only. For reliable information, go directly to the Retrieved Passages left.*")
            summary_output = gr.Markdown(
                value="_Summary will appear here..._", label="Summary"
            )

    # Combine all output components for the click function
    all_outputs = result_outputs + [summary_output]
    assert (
        len(all_outputs) == 10
    ), "Incorrect number of total output components created."

    submit_button.click(
        fn=get_results,
        inputs=question_input,
        outputs=all_outputs,  # Pass the list of 10 components
    )

# --- Launch App ---
if __name__ == "__main__":
    demo.launch()  # share=True for public link