import argparse
from src.provider import AppProvider
from src.config import OPENAI_API_KEY
from openai import OpenAI

import asyncio
import torch
import subprocess


print("CUDA available:", torch.cuda.is_available())
print("CUDA device count:", torch.cuda.device_count())

if torch.cuda.is_available():
    print("Current device:", torch.cuda.current_device())
    print("Device name:", torch.cuda.get_device_name(torch.cuda.current_device()))

openai_client = OpenAI(api_key = OPENAI_API_KEY)

app = AppProvider(openai_client)
agents = app.provide_agents()

routing_expert = {
    "customer_service" : {
        "instance": agents.provide_cs_agent_rag(),
        "chat_template" : app.provide_chat_template("customer_service")
    },
    "bank_crud": {
        "instance": agents.provide_banking_crud_agent(),
        "chat_template" : {
            "api_banking_template" : app.provide_chat_template("api_banking"),
            "responder_template" : app.provide_chat_template("responder_banking")
        }
    }
}


agents.set_routing_expert(routing_expert)



chatbot_ui = app.provide_chatbot().provide_chatbot_ui()

rtc = app.provide_rtc()
rtc_handler_cs = rtc.provide_rtc_handler_cs()
rtc_handler_banking = rtc.provide_rtc_handler_banking()
# rtc_gpt_handler = rtc.provide_rtc_gpt_handler()

parser = argparse.ArgumentParser()
parser.add_argument("--mode", choices=[
    "extract-documents",
    "rtc-cs-server", 
    "rtc-cs-ui", 
    "rtc-banking-server",
    "rtc-banking-ui",
    "chatbot",
    ], required=True)

parser.add_argument("--port", default=7861, required=True)
args = parser.parse_args()
match args.mode:
    case "extract-documents":
        asyncio.run(agents.provide_rag_executor_agent().load_documents())
    case "chatbot":
        print("Launching Chabot UI Regardly :)")
        chatbot_ui.launch(port = int(args.port))
    case "rtc-cs-server":
        print("launching RTC Server Mode ... ")
        rtc_handler_cs.start_server(port = int(args.port))
    case "rtc-cs-ui":
        print("launching RTC UI Mode ... ")
        rtc_handler_cs.launch_ui(port = int(args.port))
    case "rtc-banking-server":
        print("RTC GPT Server mode ...")
        rtc_handler_banking.start_server(port = int(args.port))
    case "rtc-banking-ui":
        print("RTC GPT UI mode ...")
        routing_expert["customer_service"]["instance"] = agents.provide_cs_agent_gpt()
        agents.set_routing_expert(routing_expert)
        rtc_handler_banking.launch_ui(port = int(args.port))
    case _:
        print("ERROR : INVALID ARGUMENT | PLEASE CHOOSE ONE BETWEEN chatbot / rtc-server/ rtc-ui / rtc-gpt-server / rtc-gpt-ui mode ")