import gradio as gr
import asyncio
import functools
from concurrent.futures import ThreadPoolExecutor
import contextvars
import logging
from model.contextual_response.contextual_response import RAGManager, get_contextual_response
from model.multi_layer_operation_predictor.operation_predictor import get_operation_definition

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

rag_manager = RAGManager()

_executor = ThreadPoolExecutor(max_workers=3)

def to_thread(func):
    @functools.wraps(func)
    async def wrapper(*args, **kwargs):
        loop = asyncio.get_running_loop()
        ctx = contextvars.copy_context()
        func_call = functools.partial(ctx.run, func, *args, **kwargs)
        return await loop.run_in_executor(_executor, func_call)
    return wrapper

async def run_contextual_response(currentLine: str) -> str:
    """Run contextual response in a safe manner."""
    try:
        return await get_contextual_response(currentLine)
    except Exception as e:
        logger.error(f"Error in contextual response: {str(e)}")
        return ""

def predict_operation(currentLine, language):
    """
    Predicts the operation based on user input and selected language
    """
    result = get_operation_definition(currentLine, language.lower())
    return result if result else "No matching operation found"

iface_operation = gr.Interface(
    fn=predict_operation,
    inputs=[
        gr.Textbox(label="Enter Operation Description"),
        gr.Dropdown(choices=["python", "javascript", "typescript", "php", "java"], label="Select Programming Language")
    ],
    outputs=gr.Textbox(label="Predicted Operation"),
    title="Operation Prediction",
    description="Enter an operation description and select a programming language to get the predicted operation.",
)

iface_contextual = gr.Interface(
    fn=lambda question: asyncio.run(run_contextual_response(question)),
    inputs=gr.Textbox(label="Enter your question"),
    outputs=gr.Textbox(label="Contextual Response"),
    title="Contextual Response",
    description="Enter a question to get a contextual response.",
)

app = gr.TabbedInterface([iface_operation, iface_contextual], ["Operation Prediction", "Contextual Response"])

if __name__ == "__main__":
    app.launch()