from __future__ import annotations

from huggingface_hub import HfFileSystem
import gradio as gr
import random
import os
import bcrypt
from typing import Optional, Any
from loguru import logger
from IPython.display import display, Markdown
from datetime import datetime, timezone
import yaml
import uuid
from glob import glob

from evidence_seeker import EvidenceSeeker
from evidence_seeker import RetrievalConfig
from evidence_seeker import ConfirmationAnalyzerConfig
from evidence_seeker import ClaimPreprocessingConfig

def get_sources(documents, confirmation_by_document) -> str | None:
    grouped = {}
    for doc in documents:
        if doc["metadata"]["file_name"] not in grouped.keys():
            grouped[doc["metadata"]["file_name"]] = {"author": doc["metadata"]["author"],
                                                    "url": doc["metadata"]["url"],
                                                    "title": doc["metadata"]["title"].replace("{","").replace("}",""),
                                                    "texts": [{"original_text": doc["metadata"]["original_text"], 
                                                               "conf": confirmation_by_document[doc["uid"]],
                                                               "full_text": doc["text"]}]}
        else:
            grouped[doc["metadata"]["file_name"]]["texts"].append({"original_text": doc["metadata"]["original_text"], "conf": confirmation_by_document[doc["uid"]], "full_text": doc["text"]})

    t = []
    for doc in grouped.keys():
        grouped[doc]["texts"] = sorted(grouped[doc]["texts"], key=lambda item: item["conf"], reverse=True)
        t.append(f"- {grouped[doc]['author']}: *{grouped[doc]['title']}* ([Link]({grouped[doc]['url']})):")
        for text in grouped[doc]["texts"]:
            short = f'"{text["original_text"].strip().replace("\n"," ").replace('"',"'")}" **[{round(text["conf"],5)}]**'
            detailed = f'"  {text["full_text"].strip().replace("\n","  ").replace('"',"'")}"'
            part = f"  - {short}\n    <details>\n    <summary>Mehr Details</summary>\n\n    - {detailed}\n\n    </details>"
            t.append(part)
    if len(t) == 0:
        return None
    else:
        return "\n\n".join(t) + "\n\n"

def describe_result(claim: str, results: list) -> str:
    preamble_template = (
        '## EvidenceSeeker Results\n\n'
        '### Input\n\n'
        '**Submitted claim:** {claim}\n\n'
        '### Results\n\n'
    )
    result_template = (
        '**Clarified claim:** <font color="orange">{text}</font> [_{statement_type}_]\n\n'
        '**Status**: {verbalized_confirmation}\n\n'
        '|Metric|Value|\n'
        '|:---|---:|\n'
        '|Average confirmation|{average_confirmation:.3f}|\n'
        '|Evidential divergence|{evidential_uncertainty:.3f}|\n'
        '|Width of evidential base|{n_evidence}|\n\n'
    )
    markdown = []
    markdown.append(preamble_template.format(claim=claim))
    for claim_dict in results:
        rdict = claim_dict.copy()
        rdict["statement_type"] = rdict["statement_type"].value
        markdown.append(result_template.format(**claim_dict))
        s = get_sources(rdict["documents"],rdict["confirmation_by_document"])
        if s: 
            s = "\n\n<details>\n\n<summary>Used documents</summary>\n\n" + s + "</details>\n\n"
            markdown.append(s)
    return "\n".join(markdown)

def setup(index_path: Optional[str] = None) -> EvidenceSeeker:
    # TODO: Load configs from YAML files. (should, perhaps, be saved under ./config/)
    retrieval_config = RetrievalConfig(
        embed_backend_type="huggingface_inference_api",
        embed_model_name="sentence-transformers/paraphrase-multilingual-mpnet-base-v2",
        embed_base_url="https://router.huggingface.co/hf-inference/models/sentence-transformers/paraphrase-multilingual-mpnet-base-v2",
        api_key_name="hf_debatelab_inference_provider",
        hub_key_name="hf_evse_data",
        bill_to="DebateLabKIT",
        #index_hub_path="DebateLabKIT/apuz-index-es",
        # Can be used if indexed is stored locally.
        index_persist_path="../evidence-seeker/TMP/APUZ/storage",
    )
    confirmation_analysis_config = ConfirmationAnalyzerConfig(
        timeout=1200,
        used_model_key="together.ai"
    )
    
    preprocessing_config = ClaimPreprocessingConfig(
        timeout=1200,
        used_model_key="together.ai"
    )

    evidence_seeker = EvidenceSeeker(retrieval_config=retrieval_config,confirmation_analysis_config=confirmation_analysis_config,preprocessing_config=preprocessing_config)
    return evidence_seeker

def draw_example(examples : list[str]) -> str:
    random.shuffle(examples)
    return examples[0]
    
def check_password(input_password: str, hash : str) -> bool:
    return bcrypt.checkpw(input_password.encode(), hash.encode())

def initialize_result_state(pipeline):
    if pipeline:
        last_result = {"pipeline": {}}
        preprocessing_config = pipeline.preprocessor.workflow.config.model_dump()
        preprocessor = preprocessing_config["models"][preprocessing_config["used_model_key"]]
        confirmation_analysis_config = pipeline.analyzer.workflow.config.model_dump()
        analyzer = confirmation_analysis_config["models"][confirmation_analysis_config["used_model_key"]]
        last_result["pipeline"]["models"] = {
            "preprocessing": {"model": preprocessor["model"], "backend_type": preprocessor["backend_type"], "max_tokens": preprocessor["max_tokens"], "temperature": preprocessor["temperature"]},
            "confirmation_analysis": {"model": analyzer["model"], "backend_type": analyzer["backend_type"], "max_tokens": analyzer["max_tokens"], "temperature": analyzer["temperature"]},
            "retrieval": {"model": pipeline.retriever.embed_model_name, "backend_type": pipeline.retriever.embed_backend_type, "embed_batch_size": pipeline.retriever.embed_batch_size, "top_k": pipeline.retriever.similarity_top_k}
        }
        index_path = "https://huggingface.co/datasets/" + pipeline.retriever.index_hub_path if pipeline.retriever.index_hub_path else "local"
        last_result["pipeline"]["index"] = index_path
        last_result["pipeline"]["ignored_statement_types"] = [t for t in pipeline.retriever.ignore_statement_types]
    else:
        last_result = {'pipeline': None}
    return gr.State(last_result)

def reactivate(check_btn, statement):
    if statement.strip() != "":
        check_btn = gr.Button(visible=True,interactive=True) 
    pos_btn = gr.Button(visible=True,interactive=True)
    neg_btn = gr.Button(visible=True,interactive=True)
    feedback = gr.Markdown(visible=True)
    return feedback, check_btn, pos_btn, neg_btn

def deactivate():
    check_btn = gr.Button(interactive=False)
    pos_btn = gr.Button(interactive=False, variant="secondary", visible=False)
    neg_btn = gr.Button(interactive=False, variant="secondary", visible=False)
    feedback = gr.Markdown(visible=False)
    return check_btn, pos_btn, neg_btn, feedback

def log_feedback(feedback_btn, last_result):
    choice = "positive" if feedback_btn == "👍" else "negative"
    last_result["feedback"] = None if last_result["feedback"] == choice else choice
    logger.log('INFO',f"{last_result["feedback"]} feedback on results")
    return gr.Button(variant="primary" if last_result["feedback"] else "secondary"), gr.Button(variant="secondary")

def log_result(last_result : dict[str, Any]):
    global RUNS_ON_SPACES, HF_BASE, LOCAL_BASE, RESULT_DIR
    fn = f"evseek_request_{last_result['request_id']}.yaml"
    subdir = _current_subdir("daily")
    if RUNS_ON_SPACES:
        if "EVIDENCE_SEEKER_APPDATA" in os.environ.keys():
            fs = HfFileSystem(token=os.environ["EVIDENCE_SEEKER_APPDATA"])
        else:
            fs = HfFileSystem()
        try:
            files = fs.glob("/".join([HF_BASE, RESULT_DIR, "**", fn]))
            assert len(files) < 2
            if len(files) == 0:
                filepath = "/".join([HF_BASE, RESULT_DIR, subdir, fn])
                fs.makedirs("/".join([HF_BASE, RESULT_DIR, subdir]), exist_ok=True)
            else:
                filepath = files[0]
        except FileNotFoundError:
            filepath = "/".join([HF_BASE, RESULT_DIR, subdir, fn])
            fs.makedirs("/".join([HF_BASE, RESULT_DIR, subdir]), exist_ok=True)
        with fs.open(filepath, encoding="utf-8", mode="w") as f:
            yaml.dump(last_result, f, allow_unicode=True, default_flow_style=False, encoding='utf-8')
    else:
        files = glob("/".join([LOCAL_BASE, RESULT_DIR, "**", fn]), recursive=True)
        assert len(files) < 2
        if len(files) == 0:
            filepath = "/".join([LOCAL_BASE, RESULT_DIR, subdir, fn])
            os.makedirs("/".join([LOCAL_BASE, RESULT_DIR, subdir]), exist_ok=True)
        else:
            filepath = files[0]
        with open(filepath, encoding="utf-8", mode="w") as f:
            yaml.dump(last_result, f, allow_unicode=True, default_flow_style=False, encoding='utf-8')

def _current_subdir(subdirectory_construction : str) -> str:
    now = datetime.now()
    if subdirectory_construction == "monthly":
        subdirectory_path = now.strftime("y%Y_m%m")
    elif subdirectory_construction == "weekly":
        year, week, _ = now.isocalendar()
        subdirectory_path = f"y{year}_w{week}"
    elif subdirectory_construction == "yearly":
        subdirectory_path = now.strftime("y%Y")
    elif subdirectory_construction == "daily":
        subdirectory_path = now.strftime("%Y_%m_%d")
    else:
        subdirectory_path = ""
    return subdirectory_path

EXAMPLE_INPUTS = [
    "Die Osterweiterung hat die EU-Institutionen nachhaltig geschwächt.",
    "In den knapp 70 Jahren seit ihrer Gründung hat es in der Bundeswehr immer wieder rechtsextremistische Vorfälle gegeben.",
    "In der Bundeswehr gibt es keinen politischen Extremismus.",
    "Die BRICS-Staaten sorgen für eine Veränderung der westlich geprägten Weltordnung.",
    "Die Genfer Konventionen sind oft hinter ihrem Anspruch, die Zivilbevölkerung zu schützen, zurückgeblieben.",
    "Die Anzahl hybrider Kriege hat zugenommen.",
    "Es ist für Deutschland wirtschaftlich ein Nachteil, dass viele Frauen in Teilzeit arbeiten.",
    "Premier Modi hat Putin als seinen Freund bezeichnet.",
    "Eine Minderheit der deutschen Bevölkerung befürwortet einen autoritären deutschen Staat.",
]

HF_BASE = 'datasets/DebateLabKIT/evidence-seeker-appdata'
RESULT_DIR = 'result_data'
LOCAL_BASE = '../TMP'

if "RUNS_ON_SPACES" not in os.environ.keys():
    RUNS_ON_SPACES = False
    # TODO: Remove - path to .env file should be set in the configt 
    # (if it is set, it will be loaded automatically)
    from dotenv import load_dotenv
    load_dotenv()
    logger.info("Gradioapp runs locally.")
else:
    RUNS_ON_SPACES = os.environ["RUNS_ON_SPACES"] == "True"
    logger.info("Gradioapp runs on 🤗.")

EVIDENCE_SEEKER = setup()

with gr.Blocks(title="EvidenceSeeker") as demo:
    valid = gr.State(value=not RUNS_ON_SPACES)
    examples = gr.State(EXAMPLE_INPUTS)
    last_result = initialize_result_state(EVIDENCE_SEEKER)
    pos_btn = gr.Button(value="👍",visible=False,interactive=False,render=False)
    neg_btn = gr.Button(value="👎",visible=False,interactive=False,render=False)
    feedback = gr.Markdown(value="Wie zufrieden bist du mit der Antwort?",render=False,visible=False)

    @gr.render(inputs=[valid])
    def renderApp(v):
        if v:
            gr.Markdown("""
                    # 🕵️ EvidenceSeeker DemoApp
                    Gib eine Aussage in das Textfeld ein und lass sie durch den EvidenceSeeker prüfen:
                """, key="intro")
            with gr.Row() as row:
                statement = gr.Textbox(label="Zu prüfende Aussage:", interactive=True, value="",scale=10,lines=3,key="statement")
                with gr.Column(scale=1):
                    example_btn = gr.Button("Zufälliges Beispiel",key="example_btn")
                    check_btn = gr.Button("Prüfe Aussage", interactive=False,key="submit_btn")
            results = gr.Markdown(value="", min_height=80, key="results")
            with gr.Column():
                feedback.render()
                with gr.Row():
                    pos_btn.render()
                    neg_btn.render()

            async def check(statement : str, last_result):
                last_result["request_time"] = datetime.now(timezone.utc).strftime('%Y-%m-%d %H:%M:%S UTC')
                last_result["request"] = statement
                last_result["claims"] = []
                last_result["request_id"] = str(uuid.uuid4())
                last_result["feedback"] = None
                if valid is False:
                    md = ""
                elif EVIDENCE_SEEKER is None:
                    md = f"Kein EvidenceSeeker initialisiert um '{statement}' zu prüfen."
                else:
                    result = await EVIDENCE_SEEKER(statement)
                    last_result["claims"] = [{"text": claim["text"],"negation":claim["negation"],"statement_type": claim["statement_type"].value,"average_confirmation": claim["average_confirmation"],"verbalized_confirmation": claim["verbalized_confirmation"]} for claim in result]
                    md = describe_result(statement, result)
                logger.info(last_result)
                return md, last_result
                
            example_btn.click(fn=draw_example, inputs=examples, outputs=statement)
            check_btn.click(deactivate, None, [check_btn, pos_btn, neg_btn, feedback]).then((lambda: "### Working..."), None, results).then(check, [statement, last_result], [results, last_result]).then(log_result, last_result, None).then(reactivate, [check_btn, statement], [feedback, check_btn, pos_btn, neg_btn])
            statement.change(lambda s: (gr.Button("Prüfe Aussage", interactive=False,key="submit_btn") if s.strip() == "" else gr.Button("Prüfe Aussage", interactive=True,key="submit_btn")), statement, [check_btn])
            pos_btn.click(log_feedback, [pos_btn, last_result], [pos_btn, neg_btn]).then(log_result, last_result, None)
            neg_btn.click(log_feedback, [neg_btn, last_result], [neg_btn, pos_btn]).then(log_result, last_result, None)
        else:
            with gr.Column() as req:
                gr.Markdown("""
                    # 🕵️ EvidenceSeeker DemoApp
                """
                )
                box = gr.Textbox(label="Please enter password for access", autofocus=True, type="password", submit_btn=True)
                res = gr.Markdown(value="")  
        
                def auth(pw : str, valid : bool):
                    output = ""
                    b = gr.Textbox(value="")
                    if "APP_HASH" not in os.environ:
                        output = "Something went wrong on our end :-("
                    elif not check_password(pw, os.environ["APP_HASH"]):
                        output = "Wrong password. Try again."
                    else:
                        output = "Continuing..."
                        valid = True
                    return output, valid, b, req
                box.submit(auth,inputs=[box,valid],outputs=[res,valid, box, req])

    demo.launch()