import time
import rsazure_openai_toolkit as rschat
from rsazure_openai_toolkit.core import integration as rschat_core


def main():
    user_input = "Say hello!"

    # Carrega variáveis do .env com override explícito
    rschat.load_env(override=True)

    context_data = rschat.get_context_messages(user_input=user_input)
    messages = context_data["messages"]

    model_config = rschat.get_model_config()
    config = rschat.get_cli_config()

    start = time.time()
    response = rschat_core.main(
        api_key=config["api_key"],
        azure_endpoint=config["endpoint"],
        api_version=config["version"],
        deployment_name=config["deployment_name"],
        messages=messages,
        **model_config
    )
    elapsed = round(time.time() - start, 2)

    if not response.choices:
        print("❌ No response received from the model.")
        return

    content = response.choices[0].message.content
    usage = response.usage.model_dump() if response.usage else {}

    input_tokens = usage.get("prompt_tokens", 0)
    output_tokens = usage.get("completion_tokens") or rschat.estimate_input_tokens(
        messages=[{"role": "assistant", "content": content}],
        deployment_name=config["deployment_name"]
    )
    total_tokens = usage.get("total_tokens", input_tokens + output_tokens)

    result = rschat.ChatResult(
        question=user_input,
        response_text=content,
        system_prompt=config["system_prompt"],
        model=response.model,
        seed=model_config.get("seed"),
        input_tokens=input_tokens,
        output_tokens=output_tokens,
        total_tokens=total_tokens,
        elapsed_time=elapsed,
        model_config=model_config,
        raw_response=response.model_dump()
    )

    result.print()

    logger = rschat.get_logger()
    if logger.enabled:
        logger.log(result.to_log_dict())
    else:
        print(f"📭 Logging is disabled ({logger})\n")


if __name__ == "__main__":
    main()
