from __future__ import annotations

import asyncio
import logging
import uuid
import yaml

import gradio as gr  # type: ignore

from logikon.backends.chat_models_with_grammar import create_logits_model, LogitsModel, LLMBackends
from logikon.guides.proscons.recursive_balancing_guide import RecursiveBalancingGuide, RecursiveBalancingGuideConfig

from backend.config import process_config
from backend.messages_processing import add_details, history_to_langchain_format
from backend.svg_processing import postprocess_svg

logging.basicConfig(level=logging.INFO)


EXAMPLES = [
    ("We're a nature-loving family with three kids, have some money left, and no plans "
     "for next week-end. Should we visit Disneyland?"),
    "Should I stop eating animals?",
    "Bob needs a reliable and cheap car. Should he buy a Mercedes?",
    ('Gavin has an insurance policy that includes coverage for "General Damages," '
     'which includes losses from "missed employment due to injuries that occur '
     'under regular working conditions."\n\n'
     'Gavin works as an A/C repair technician in a small town. One day, Gavin is '
     'hired to repair an air conditioner located on the second story of a building. '
     'Because Gavin is an experienced repairman, he knows that the safest way to '
     'access the unit is with a sturdy ladder. While climbing the ladder, Gavin '
     'loses his balance and falls, causing significant injury. Because of this, he '
     'subsequently has to stop working for weeks. Gavin files a claim with his '
     'insurance company for lost income.\n\n'
     'Does Gavin\'s insurance policy cover his claim for lost income?'),
     "How many arguments did you consider in your internal reasoning? (Brief answer, please.)",
     "Did you consider any counterarguments in your internal reasoning?",
     "From all the arguments you considered and assessed, which one is the most important?",
     "Did you refute any arguments or reasons for lack of plausibility?"
]

TITLE = """<div align=left>
    <h1>🪂 Logikon <i>Guided Reasoning™️</i> Demo Chatbot</h1>
    </div>
    <p>This is a <b>TEMPLATE</b>:<br/>
    ➡️ Duplicate this space and configure your own inference endpoints
    in the <code>config.yaml</code> file to get started!</p>
    """

TERMS_OF_SERVICE ="""<h2>Terms of Service</h2>

    <p>This app is provided by Logikon AI for educational and research purposes only. 
    The app is powered by Logikon's <i>Guided Reasoning™️</i>&nbsp; technology, which is a novel approach to 
    reasoning with language models. The app is a work in progress and may not always provide accurate or reliable information. 
    By accepting these terms of service, you agree not to use the app:</p>

    <ol>
    <li>In any way that violates any applicable national, federal, state, local or international law or regulation;</li>
    <li>For the purpose of exploiting, harming or attempting to exploit or harm minors in any way;</li>
    <li>To generate and/or disseminate malware (e.g. ransomware) or any other content to be used for the purpose of harming electronic systems;</li>
    <li>To generate or disseminate verifiably false information and/or content with the purpose of harming others;</li>
    <li>To generate or disseminate personal identifiable information that can be used to harm an individual;</li>
    <li>To generate or disseminate information and/or content (e.g. images, code, posts, articles), and place the information and/or content in any public context (e.g. bot generating tweets) without expressly and intelligibly disclaiming that the information and/or content is machine generated;</li>
    <li>To defame, disparage or otherwise harass others;</li>
    <li>To impersonate or attempt to impersonate (e.g. deepfakes) others without their consent;</li>
    <li>For fully automated decision making that adversely impacts an individual’s legal rights or otherwise creates or modifies a binding, enforceable obligation;</li>
    <li>For any use intended to or which has the effect of discriminating against or harming individuals or groups based on online or offline social behavior or known or predicted personal or personality characteristics;</li>
    <li>To exploit any of the vulnerabilities of a specific group of persons based on their age, social, physical or mental characteristics, in order to materially distort the behavior of a person pertaining to that group in a manner that causes or is likely to cause that person or another person physical or psychological harm;</li>
    <li>For any use intended to or which has the effect of discriminating against individuals or groups based on legally protected characteristics or categories;</li>
    <li>To provide medical advice and medical results interpretation;</li>
    <li>To generate or disseminate information for the purpose to be used for administration of justice, law enforcement, immigration or asylum processes, such as predicting an individual will commit fraud/crime commitment (e.g. by text profiling, drawing causal relationships between assertions made in documents, indiscriminate and arbitrarily-targeted use).            </li>
    </ol>

    """

CHATBOT_INSTRUCTIONS = (
    "1️⃣ In the first turn, ask a question or present a decision problem.\n"
    "2️⃣ In the following turns, ask the chatbot to explain its reasoning.\n\n"
    "💡 Note that this demo bot is hard-wired to deliberate with Guided Reasoning™️ "
    "in the first turn only.\n\n"
    "🔐 Chat conversations and feedback are logged (anonymously).\n"
    "Please don't share sensitive or identity revealing information.\n\n"
    "🙏 Benjamin is powered by the free API inference services of 🤗.\n"
    "In case you encounter issues due to rate limits... simply try again later.\n"
    "[We're searching sponsors to run Benjamin on 🚀 dedicated infrastructure.]\n\n"
    "💬 We'd love to hear your feedback!\n"
    "Please use the 👋 Community tab above to reach out.\n"
)

# config
with open("config.yaml") as stream:
    try:
        demo_config = yaml.safe_load(stream)
        logging.info(f"Config: {demo_config}")
    except yaml.YAMLError as exc:
        logging.error(f"Error loading config: {exc}")
        gr.Error("Error loading config: {exc}")

try:
    client_kwargs, guide_kwargs = process_config(demo_config)
except Exception as exc:
    logging.error(f"Error processing config: {exc}")
    gr.Error(f"Error processing config: {exc}")

logging.info(f"Reasoning guide expert model is {guide_kwargs['expert_model']}.")



def new_conversation_id():
    conversation_id = str(uuid.uuid4())
    print(f"New conversation with conversation ID: {conversation_id}")
    return conversation_id


def setup_client_llm(**client_kwargs) -> LogitsModel | None:
    try:
        llm = create_logits_model(**client_kwargs)
    except Exception as e:
        logging.error(f"When setting up client llm: Error: {e}")
        return False
    return llm


async def log_like_dislike(conversation_id: gr.State, x: gr.LikeData):
    print(conversation_id, x.index, x.liked)
    # add your own feedback logging here


def add_message(history, message, conversation_id):
    if len(history) == 0:
        # reset conversation id
        conversation_id = new_conversation_id()

    print(f"add_message: {history} \n {message}")
    if message["text"] is not None:
        history.append((message["text"], None))
    return history, gr.MultimodalTextbox(value=None, interactive=False), conversation_id


async def bot(
    history,
    conversation_id,
    progress=gr.Progress(),
):
    
    print(f"History (conversation: {conversation_id}): {history}")
    history_langchain_format = history_to_langchain_format(history)

    # use guide always and exclusively at first turn
    if len(history_langchain_format) <= 1:

        # health check
        gr.Info("Checking availability and health of inference endpoints ...", duration=6)
        health_check = await guide.health_check()
        await asyncio.sleep(0.1)

        if health_check.get("status", None) != "ok":
            health_msg = " | ".join([f"{k}: {v}" for k, v in health_check.items()])
            logging.error(f"Guide health check failed: {health_msg}")
            gr.Error(f"LLM availability / health check failed: {health_msg}")
        logging.info(f"Health check: {health_check}")

        message = history[-1][0]

        try:
            artifacts = {}
            progress_step = 0
            async for otype, ovalue in guide.guide(message):
                logging.info(f"Guide output: {otype.value} - {ovalue}")
                if otype.value == "progress":
                    logging.info(f"Progress: {ovalue}")
                    gr.Info(ovalue, duration=12)
                    progress((progress_step,4))
                    progress_step += 1
                elif otype is not None:
                    artifacts[otype.value] = ovalue
                else:
                    break
                await asyncio.sleep(0.1)
        except asyncio.TimeoutError:
            msg = "Guided reasoning process took too long. Please try again."
            raise gr.Error(msg)
        except Exception as e:
            msg = f"Error during guided reasoning: {e}"            
            raise gr.Error(msg)
        
        svg = postprocess_svg(artifacts.get("svg_argmap"))
        protocol = artifacts.get("protocol", "I'm sorry, I failed to reason about the problem.")
        response = artifacts.pop("response", "")
        if not response:
            response = "I'm sorry, I failed to draft a response (see logs for details)."
        response = add_details(response, protocol, svg)

    # otherwise, just chat
    else:
        try:
            response = client_llm.invoke(history_langchain_format).content
        except Exception as e:
            msg = f"Error during chatbot inference: {e}"
            gr.Error(msg)
            raise ValueError(msg)
    
    print(f"Response: {response}")
    history[-1][1] = response
    
    return history



with gr.Blocks() as demo:
    
    # preamble
    gr.Markdown(TITLE)
    conversation_id = gr.State(str(uuid.uuid4()))
    tos_approved = gr.State(False)

    if not client_kwargs["inference_server_url"]:
        gr.Markdown(
            "⚠️ **Error:** Please set the client model inference endpoint in the `config.yaml` file."
        )
    if not guide_kwargs["inference_server_url"]:
        gr.Markdown(
            "⚠️ **Error:** Please set the expert model inference endpoint in the `config.yaml` file."
        )
    if not guide_kwargs["classifier_kwargs"]["inference_server_url"]:
        gr.Markdown(
            "⚠️ **Error:** Please set the classifier model inference endpoint in the `config.yaml` file."
        )
    
    # set up client and guide
    client_llm = setup_client_llm(**client_kwargs)
    guide_config = RecursiveBalancingGuideConfig(**guide_kwargs)
    guide = RecursiveBalancingGuide(tourist_llm=client_llm, config=guide_config)

    with gr.Tab(label="Chatbot", visible=False) as chatbot_tab:

        # chatbot
        chatbot = gr.Chatbot(
            [],
            elem_id="chatbot",
            bubble_full_width=False,
            placeholder=CHATBOT_INSTRUCTIONS,
        )
        chat_input = gr.MultimodalTextbox(interactive=True, file_types=["image"], placeholder="Enter message ...", show_label=False)
        clear = gr.ClearButton([chat_input, chatbot])
        gr.Examples([{"text": e, "files":[]} for e in EXAMPLES], chat_input)

        # logic
        chat_msg = chat_input.submit(add_message, [chatbot, chat_input, conversation_id], [chatbot, chat_input, conversation_id])
        bot_msg = chat_msg.then(
            bot,
            [
                chatbot,
                conversation_id
            ],
            chatbot,
            api_name="bot_response"
        )
        bot_msg.then(lambda: gr.MultimodalTextbox(interactive=True), None, [chat_input])

        chatbot.like(log_like_dislike, [conversation_id], None)

        # we're resetting conversation id when drafting first response in bot()
        # clear.click(new_conversation_id, outputs = [conversation_id])

    with gr.Tab(label="Terms of Service") as tos_tab:

        gr.HTML(TERMS_OF_SERVICE)
        tos_checkbox = gr.Checkbox(label="I agree to the terms of service")
        tos_checkbox.input(
            lambda x: (x, gr.Checkbox(label="I agree to the terms of service", interactive=False), gr.Tab("Chatbot", visible=True)),
            tos_checkbox,
            [tos_approved, tos_checkbox, chatbot_tab]
        )

if __name__ == "__main__":
    demo.queue(default_concurrency_limit=8)
    demo.launch(show_error=True)