# grad_bot2.py
import sys
import logging

# 1) Configure logging to go to stdout
logging.basicConfig(
    level=logging.INFO,
    stream=sys.stdout,
    format="%(asctime)s %(levelname)s %(name)s | %(message)s",
)

# 2) (optional) Redirect built-in print → logging.info
import builtins
_original_print = builtins.print
def _patched_print(*args, **kwargs):
    msg = " ".join(str(a) for a in args)
    logging.info(msg)
builtins.print = _patched_print


import gradio as gr
import aiohttp
import builtins

# Configure logging to capture all stdout and stderr
debug_format = '%(asctime)s - %(levelname)s - %(message)s'
logging.basicConfig(stream=sys.stdout, level=logging.INFO, format=debug_format)

# Redirect other loggers (uvicorn, gradio internal) to our root logger
for logger_name in ['uvicorn', 'uvicorn.error', 'uvicorn.access', 'fastapi', 'gradio']:
    logger = logging.getLogger(logger_name)
    logger.handlers = logging.getLogger().handlers
    logger.propagate = True

# Redirect stdout and stderr to logging
class LoggerWriter:
    def __init__(self, level_func):
        self.level_func = level_func
    def write(self, message):
        message = message.strip()
        if message:
            self.level_func(message)
    def flush(self):
        pass

sys.stdout = LoggerWriter(logging.info)
sys.stderr = LoggerWriter(logging.error)

# Also capture any print calls into logging
_orig_print = builtins.print
builtins.print = lambda *args, **kwargs: logging.info(' '.join(str(a) for a in args))

# Configuration for the licensing service
LICENSING_SERVICE_BASE = 'https://gbdlicensing01containerapp.delightfulocean-00db0bf9.westeurope.azurecontainerapps.io'

async def call_licensing_api(
    endpoint: str,
    index_name: str,
    api_key: str,
    input_text: str = None,
    question: str = None,
) -> dict:
    """
    Sends a POST to either /create_index or /answer_question on your FastAPI service.
    """
    if endpoint not in ("create_index", "answer_question"):
        raise ValueError("`endpoint` must be either 'create_index' or 'answer_question'")
    if endpoint == "create_index" and input_text is None:
        raise ValueError("When calling 'create_index', you must provide `input_text`")
    if endpoint == "answer_question" and question is None:
        raise ValueError("When calling 'answer_question', you must provide `question`")

    headers = {"X-API-Key": api_key, "Content-Type": "application/json"}
    payload = {"index_name": index_name}
    if endpoint == "create_index":
        payload["input_text"] = input_text
    else:
        payload["question"] = question

    timeout = aiohttp.ClientTimeout(total=60)
    async with aiohttp.ClientSession(timeout=timeout) as session:
        url = f"{LICENSING_SERVICE_BASE}/{endpoint}"
        logging.info(f"Calling licensing API at {url} with payload: {payload}")
        async with session.post(url, json=payload, headers=headers) as resp:
            resp.raise_for_status()
            result = await resp.json()
            logging.info(f"Received response: {result}")
            return result

async def get_answer(api_key: str, final_selection: str, question: str) -> str:
    """Fetches an answer for the question based on the final topic selection."""
    index_map = {
        'Office 365': 'licensing-fundamentals'
    }
    index_name = index_map.get(final_selection)
    logging.info(f"Final selection: {final_selection} -> index: {index_name}")
    response = await call_licensing_api(
        endpoint="answer_question",
        index_name=index_name,
        api_key=api_key,
        question=question
    )
    answer = response.get('answer', '')
    logging.info(f"Extracted answer: {answer}")
    return answer


def main():
    with gr.Blocks() as demo:
        # API key input
        api_key_input = gr.Textbox(
            label='Licensing Service API Key',
            placeholder='Paste your API key here',
            type='password'
        )

        # Step 1: Primary topic selection
        primary_topic = gr.Radio(
            choices=['License Program', 'Products', 'Customer info'],
            label='What would you like to learn about?'
        )

        # Step 2: Product type, shown only if 'Products' selected
        product_container = gr.Column(visible=False)
        with product_container:
            product_type = gr.Radio(
                choices=['Software', 'Online services'],
                label='Which products?'
            )

        # Step 3: Service type, shown only if 'Online services' selected
        service_container = gr.Column(visible=False)
        with service_container:
            service_type = gr.Radio(
                choices=['Office 365'],
                label='Select an online service'
            )

        # Step 4: QA, shown only after service_type selection
        qa_container = gr.Column(visible=False)
        with qa_container:
            question = gr.Textbox(
                lines=3,
                placeholder='Type your question here...',
                label='Your Question'
            )
            answer = gr.Textbox(
                lines=5,
                interactive=False,
                label='Answer'
            )
            submit = gr.Button('Submit')
            submit.click(
                fn=get_answer,
                inputs=[api_key_input, service_type, question],
                outputs=[answer]
            )

        # Callbacks to reveal steps
        primary_topic.change(
            fn=lambda sel: gr.update(visible=(sel=='Products')),
            inputs=[primary_topic],
            outputs=[product_container]
        )
        primary_topic.change(
            fn=lambda sel: gr.update(visible=False),
            inputs=[primary_topic],
            outputs=[service_container, qa_container]
        )

        product_type.change(
            fn=lambda sel: gr.update(visible=(sel=='Online services')),
            inputs=[product_type],
            outputs=[service_container]
        )
        product_type.change(
            fn=lambda sel: gr.update(visible=False),
            inputs=[product_type],
            outputs=[qa_container]
        )

        service_type.change(
            fn=lambda sel: gr.update(visible=(sel=='Office 365')),
            inputs=[service_type],
            outputs=[qa_container]
        )

        # Launch with public sharing enabled
        demo.launch(share=True)

if __name__ == '__main__':
    main()