import gradio as gr
import catboost as cb
import torch
from transformers import AutoTokenizer, AutoModel
import pandas as pd
import numpy as np
import os
import datetime
import uuid
import threading
import re
import time
from huggingface_hub import HfApi, login

# --- CONFIGURATION ---
EMBEDDING_MODEL_NAME = "mixedbread-ai/deepset-mxbai-embed-de-large-v1"
BINARY_MODEL_PATH = "catboost_tweet_multiclass_model_tuned.cbm"
NARRATIVE_MODEL_PATH = "catboost_narrative_model.cbm"

# Data Collection
DATASET_REPO_ID = "RuDisinfo-EN/RUDisinfo_Feedback"
HF_TOKEN = os.environ.get("HF_TOKEN")

# Thresholds
BINARY_LOWER = 0.40
BINARY_UPPER = 0.60
NARRATIVE_MARGIN = 0.15

MAX_CHAR_LIMIT = 2000

# --- SETUP: HF Login ---
api = None

def get_api():
    global api
    if api is not None:
        return api
    token = os.environ.get("HF_TOKEN")
    if not token:
        return None
    try:
        api = HfApi(token=token)
        user = api.whoami()
        print(f"Logged in as: {user['name']}")
        return api
    except Exception as e:
        print(f"HF API init failed: {type(e).__name__}: {e}")
        api = None
        return None

# --- Thread-safe Model Loading ---
tokenizer = None
embedding_model = None
model_binary = None
model_narrative = None
_models_loaded = False
_model_lock = threading.Lock()
_inference_lock = threading.Lock()

def load_models():
    global tokenizer, embedding_model, model_binary, model_narrative, _models_loaded
    if _models_loaded:
        return
    with _model_lock:
        if _models_loaded:  # Double-checked locking
            return
        print("Loading models on first request...")
        tokenizer = AutoTokenizer.from_pretrained(EMBEDDING_MODEL_NAME)
        embedding_model = AutoModel.from_pretrained(EMBEDDING_MODEL_NAME)
        embedding_model.eval()
        model_binary = cb.CatBoostClassifier()
        model_binary.load_model(BINARY_MODEL_PATH)
        model_narrative = cb.CatBoostClassifier()
        model_narrative.load_model(NARRATIVE_MODEL_PATH)
        _models_loaded = True
        print("Models loaded.")

NARRATIVE_MAP = {
    0: 'Elites vs. People', 1: 'Critique of the West', 2: 'Victimhood',
    3: 'Lost Sovereignty', 4: 'Imminent Collapse of West', 5: 'Hahaganda (Mockery)',
    6: 'Delegitimize Ukraine', 7: 'Biolabs Conspiracy', 8: 'Just Asking Questions',
    9: 'Aggressor Inversion', 10: 'Historical Revisionism', 12: 'Nuclear Coercion'
}

# --- HELPER FUNCTIONS ---
def get_embedding(text):
    with _inference_lock:
        encoded_input = tokenizer([text], padding=True, truncation=True, return_tensors='pt', max_length=512)
        with torch.no_grad():
            model_output = embedding_model(**encoded_input)
            token_embeddings = model_output.last_hidden_state
            attention_mask = encoded_input['attention_mask']
            input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
            sum_embeddings = torch.sum(token_embeddings * input_mask_expanded, 1)
            sum_mask = torch.clamp(input_mask_expanded.sum(1), min=1e-9)
            return (sum_embeddings / sum_mask)[0].numpy()

def sanitize_for_csv(text):
    if not text:
        return text
    text = re.sub(r'[\r\n\t]', ' ', text)
    if text[0] in ('=', '+', '-', '@'):
        return "'" + text
    return text

def save_to_dataset(text, prob_misinfo, top_narrative, prob_narrative_1, prob_narrative_2, reason="Hard Sample"):
    hf_api = get_api()
    if hf_api is None:
        return "Setup error: HF API not initialized."

    unique_id = str(uuid.uuid4())
    timestamp = datetime.datetime.now().date().isoformat()
    filename = f"feedback_{unique_id}.csv"

    new_row = pd.DataFrame([{
        "timestamp": timestamp, "text": sanitize_for_csv(text), "prob_misinfo": prob_misinfo,
        "top_narrative": top_narrative, "conf_1": prob_narrative_1, "conf_2": prob_narrative_2,
        "reason": reason
    }])

    temp_path = f"/tmp/{filename}"
    new_row.to_csv(temp_path, index=False)

    try:
        hf_api.upload_file(
            path_or_fileobj=temp_path,
            path_in_repo=filename,
            repo_id=DATASET_REPO_ID,
            repo_type="dataset"
        )
        return "Saved."
    except Exception as e:
        print(f"Upload failed: {type(e).__name__}")
        return "Save failed (server error)."
    finally:
        if os.path.exists(temp_path):
            os.remove(temp_path)

# --- Session-based Feedback Rate Limiting ---
# State: {"feedback_given": bool, "last_feedback_ts": float}

def _can_give_feedback(session_state):
    """Returns (allowed: bool, reason: str)"""
    if not session_state:
        return True, ""
    if session_state.get("feedback_given"):
        return False, "Feedback already submitted for this analysis."
    last_ts = session_state.get("last_feedback_ts", 0)
    if time.time() - last_ts < 5:
        return False, "Please wait a moment before submitting feedback."
    return True, ""

# --- BUTTON CALLBACKS ---

def on_like(state_data, session_state):
    allowed, msg = _can_give_feedback(session_state)
    if not allowed:
        return msg, session_state
    if not state_data:
        return "No analysis data found. Please run an analysis first.", session_state
    result_msg = save_to_dataset(
        state_data["text"], state_data["prob_misinfo"], state_data["top_narrative"],
        state_data["conf_1"], state_data["conf_2"],
        reason="User Feedback: Confirmed correct"
    )
    new_session = {"feedback_given": True, "last_feedback_ts": time.time()}
    return f"Feedback submitted: {result_msg}", new_session

def show_feedback_options():
    return gr.Group(visible=True)

def on_detailed_dislike(state_data, specific_reason, session_state):
    allowed, msg = _can_give_feedback(session_state)
    if not allowed:
        return msg, gr.Group(visible=False), session_state
    if not state_data:
        return "No analysis data found. Please run an analysis first.", gr.Group(visible=False), session_state
    result_msg = save_to_dataset(
        state_data["text"],
        state_data["prob_misinfo"],
        state_data["top_narrative"],
        state_data["conf_1"],
        state_data["conf_2"],
        reason=f"User Feedback: {specific_reason}"
    )
    new_session = {"feedback_given": True, "last_feedback_ts": time.time()}
    return f"Feedback submitted ({specific_reason}): {result_msg}", gr.Group(visible=False), new_session

def analyze_and_process(text):
    # Returns 6 values: bin, narr_top, narr_full, status, analysis_state, feedback_session_reset
    if not text or len(text.strip()) < 5:
        return None, None, None, "Please enter valid text.", None, {}

    text = " ".join(text.split())

    if len(text) > MAX_CHAR_LIMIT:
        return None, None, None, f"Input too long. Maximum {MAX_CHAR_LIMIT} characters allowed.", None, {}

    try:
        load_models()

        vector = get_embedding(text)
        binary_probs = model_binary.predict_proba([vector])[0]
        prob_misinfo = round(float(binary_probs[1]), 3)
        prob_safe = round(float(binary_probs[0]), 3)

        narr_probs = model_narrative.predict_proba([vector])[0]
        narrative_dict_full = {NARRATIVE_MAP[i]: round(float(p), 3) for i, p in enumerate(narr_probs) if i in NARRATIVE_MAP}

        sorted_indices = np.argsort(narr_probs)[::-1]
        sorted_indices_filtered = [i for i in sorted_indices if i in NARRATIVE_MAP]
        top1_idx = sorted_indices_filtered[0]
        top2_idx = sorted_indices_filtered[1]
        margin = narr_probs[top1_idx] - narr_probs[top2_idx]

        is_binary_ambiguous = BINARY_LOWER <= prob_misinfo <= BINARY_UPPER
        is_narrative_confused = margin < NARRATIVE_MARGIN

        save_status = "Submit feedback below to help improve the model."
        if is_binary_ambiguous and is_narrative_confused:
            auto_result = save_to_dataset(
                text, prob_misinfo, NARRATIVE_MAP[top1_idx],
                narr_probs[top1_idx], narr_probs[top2_idx],
                reason="Auto: Hard Sample"
            )
            save_status = f"Ambiguous result logged for research ({auto_result})"

        if prob_misinfo > BINARY_UPPER:
            main_status = f"High disinformation risk ({prob_misinfo*100:.1f}%)"
        elif prob_misinfo > BINARY_LOWER:
            main_status = f"Moderate disinformation risk ({prob_misinfo*100:.1f}%)"
        else:
            main_status = f"Low disinformation risk ({prob_misinfo*100:.1f}%)"

        final_status = f"{main_status}\n{save_status}"

        state_data = {
            "text": text,
            "prob_misinfo": prob_misinfo,
            "top_narrative": NARRATIVE_MAP[top1_idx],
            "conf_1": float(narr_probs[top1_idx]),
            "conf_2": float(narr_probs[top2_idx])
        }

        # Reset feedback session atomically with new analysis result
        return ({"No Disinformation Detected": prob_safe, "Disinformation Detected": prob_misinfo},
                narrative_dict_full, narrative_dict_full, final_status, state_data, {})

    except Exception as e:
        print(f"Internal error: {type(e).__name__}: {str(e)}")
        return None, None, None, "An internal error occurred. Please try again later.", None, {}

# --- GUI ---
theme = gr.themes.Ocean(
    primary_hue="cyan",
    secondary_hue="blue",
    neutral_hue="slate",
    font=[gr.themes.GoogleFont("Roboto"), "ui-sans-serif", "system-ui", "sans-serif"],
).set(
    button_primary_background_fill="*primary_600",
    button_primary_background_fill_hover="*primary_500",
    block_title_text_weight="600",
    container_radius="radius_lg",
    button_border_width="1px",
)

custom_css = """
.gradio-container {font-family: 'Roboto', sans-serif;}
#status_box {
    background-color: #f0f9ff;
    border: 1px solid #bae6fd;
    padding: 15px;
    border-radius: 12px;
    box-shadow: 0 4px 6px -1px rgb(0 0 0 / 0.1);
}
.footer {text-align: center; margin-top: 30px; font-size: 0.9em; color: #4b5563; line-height: 1.6;}
"""

with gr.Blocks(theme=theme, css=custom_css, title="Misinfo Research Tool") as demo:
    current_analysis_state = gr.State()
    feedback_session_state = gr.State({})

    gr.Markdown("# AI Disinformation Analysis Tool")
    gr.Markdown("A collaborative research prototype for detecting disinformation narratives. Optimized for German/English text.")

    with gr.Accordion("Technical Details & Methodology", open=False):
        gr.Markdown("Architecture: Distilled, Decoupled Multi-Stage Inference Pipeline using two CatBoost models (Qwen3-Next-80B as teacher) with Active Learning Loop. Optimized for German/English text analysis.")
        gr.Markdown("""
        The tool analyzes text using two independent models running in parallel:

        **Step 1: Risk Assessment**
        The first model estimates the probability that a text contains disinformation:
        - **High Risk:** above 60%
        - **Moderate Risk:** between 40% and 60%
        - **Low Risk:** below 40%

        **Step 2: Narrative Fingerprinting**
        The second model maps the text against 13 specific propaganda themes (e.g. "Biolabs Conspiracy", "Victimhood"). This runs regardless of risk level, allowing detection of subtle persuasive patterns even in low-risk text.

        > **Data Notice:** When the model is uncertain about a result, the input text may be automatically logged to our research dataset to support model improvement. By using this tool, you consent to this use. No personal data is collected.
        """)

    with gr.Row():
        with gr.Column(scale=1):
            input_text = gr.Textbox(label="Input Text", placeholder="Paste tweet or text here...", lines=8)
            btn = gr.Button("Analyze", variant="primary")
            status_output = gr.Textbox(label="Status", elem_id="status_box", lines=3)
            with gr.Row():
                btn_like = gr.Button("Correct prediction", variant="secondary")
                with gr.Column(scale=1):
                    btn_dislike_trigger = gr.Button("Incorrect prediction", variant="stop")
                    with gr.Group(visible=False) as feedback_group:
                        feedback_reason = gr.Radio(
                            choices=[
                                "Wrong Narrative Detected",
                                "Completely Safe Text (False Positive)",
                                "Other Error"
                            ],
                            label="What is incorrect?",
                            value="Wrong Narrative Detected"
                        )
                        btn_submit_dislike = gr.Button("Submit Feedback", size="sm")

        with gr.Column(scale=1):
            lbl_bin = gr.Label(label="Risk Assessment", num_top_classes=2)
            lbl_narr_top = gr.Label(label="Dominant Narratives", num_top_classes=3)

    with gr.Accordion("Detailed Analysis", open=False):
        json_output = gr.Label(label="Full Narrative Spectrum", num_top_classes=13)

    btn.click(
        fn=analyze_and_process,
        inputs=input_text,
        outputs=[lbl_bin, lbl_narr_top, json_output, status_output, current_analysis_state, feedback_session_state]
    )
    btn_like.click(
        fn=on_like,
        inputs=[current_analysis_state, feedback_session_state],
        outputs=[status_output, feedback_session_state],
        concurrency_limit=1
    )
    btn_dislike_trigger.click(fn=show_feedback_options, inputs=None, outputs=feedback_group)
    btn_submit_dislike.click(
        fn=on_detailed_dislike,
        inputs=[current_analysis_state, feedback_reason, feedback_session_state],
        outputs=[status_output, feedback_group, feedback_session_state]
    )

    gr.HTML(
        """
        <div class="footer">
            <hr style="margin-top: 30px; margin-bottom: 20px; border: 0; border-top: 1px solid #e5e7eb;">
            <h3>Project Consortium &amp; Development Team</h3>
            <p style="margin-bottom: 15px;">
                <b>Project Contributors:</b> Olga Slivko (RSM), Raphaela Andres (ZEW Mannheim), Jacob Schildknecht (ZEW Mannheim), Heiko Paulheim (Uni Mannheim)
            </p>
            <p style="margin-bottom: 15px;">
                <b>Expert Validators:</b>
                <a href="https://www.phil.uni-mannheim.de/zeitgeschichte/team/assoc-prof-dr-ivan-balykin/" target="_blank">Prof. Dr. Ivan Balykin</a>,
                <a href="https://www.researchgate.net/profile/Viktoriia-Makovska" target="_blank">Viktoriia Makovska</a>,
                <a href="https://www.tu.berlin/en/qu/ueber-uns/team-personen/senior-researchers/dr-veronika-solopova" target="_blank">Dr. Veronika Solopova</a>
            </p>
            <p style="font-style: italic; color: #6b7280;">
                A joint research initiative by:<br>
                <b>Rotterdam School of Management,</b>
                <b>ZEW Mannheim,</b>
                <b>University of Mannheim</b>
            </p>
            <p style="font-size: 0.8em; margin-top: 15px;">
                &copy; 2026 Research Group "Datengestutzte Analyse von Hate Speech und die Auswirkungen der Regulierung in Deutschland". Powered by Hugging Face &amp; CatBoost.
                This project was partially funded by the German Foundation for Peace Research (DSF) in conjunction with funding provided by the German Federal Ministry of Education and Research (BMBF).
            </p>
        </div>
        """
    )

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