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


def chat_loop():
    print("🔁 Chat loop started")
    print("💡 Type 'exit' to quit\n")

    # Carrega variáveis do .env (força override)
    rschat.load_env(override=True)
    config = rschat.get_cli_config()
    logger = rschat.get_logger()

    while True:
        user_input = input("You: ").strip()
        if user_input.lower() == "exit":
            print("👋 Goodbye!")
            break
        if not user_input:
            continue

        # Prepara mensagens com ou sem contexto persistente
        context_data = rschat.get_context_messages(user_input=user_input)
        messages = context_data["messages"]

        model_config = rschat.get_model_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.")
            continue

        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)

        # Constrói resultado e imprime no formato padrão
        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()

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


def main():
    chat_loop()


if __name__ == "__main__":
    main()
