import os
import time
from dotenv import load_dotenv
from rsazure_openai_toolkit import call_azure_openai_handler
from rsazure_openai_toolkit.utils.model_config_utils import get_model_config
from rsazure_openai_toolkit.logging.interaction_logger import InteractionLogger


load_dotenv()

messages = [
    {"role": "system", "content": "You are a helpful assistant."},
    {"role": "user", "content": "Say hello!"}
]

model_config = get_model_config()
start = time.time()

response = call_azure_openai_handler(
    api_key=os.getenv("AZURE_OPENAI_API_KEY"),
    azure_endpoint=os.getenv("AZURE_OPENAI_ENDPOINT"),
    api_version=os.getenv("AZURE_OPENAI_API_VERSION"),
    deployment_name=os.getenv("AZURE_DEPLOYMENT_NAME"),
    messages=messages,
    **model_config
)

elapsed = round(time.time() - start, 2)
content = response.choices[0].message.content
usage = response.usage.model_dump() if response.usage else {}
model_used = response.model
seed_used = model_config.get("seed")

# Token info (fallback if needed)
input_tokens = usage.get("prompt_tokens", 0)
output_tokens = usage.get("completion_tokens", 0)
total_tokens = usage.get("total_tokens", input_tokens + output_tokens)

print(f"\nAssistant:\n\t{content}")
print("\n----- REQUEST INFO -----")
print(f"📤 Input tokens: {input_tokens}")
print(f"📥 Output tokens: {output_tokens}")
print(f"🧾 Total tokens: {total_tokens}")
print(f"🧠 Model: {model_used}")
print(f"🎲 Seed: {seed_used}")
print(f"⏱️ Time: {elapsed}s\n")

# Optional logging mode: none [default], jsonl, csv
logger = InteractionLogger(
    mode=os.getenv("RSCHAT_LOG_MODE"),
    path=os.getenv("RSCHAT_LOG_PATH")
)

if logger.enabled:
    logger.log({
        "question": "Say hello!",
        "response": content,
        "model": model_used,
        "usage": usage,
        "model_config": model_config,
        "raw_response": response.model_dump()
    })
else:
    print("📭 Logging is disabled (RSCHAT_LOG_MODE is 'none' or not configured)\n")
