import os
import logging
import requests
import csv
import datetime
from dotenv import load_dotenv
from telegram import Update
from telegram.ext import Application, CommandHandler, MessageHandler, filters, ContextTypes, CallbackContext

load_dotenv()

logging.basicConfig(
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    level=logging.INFO
)

TELEGRAM_TOKEN = os.getenv("TELEGRAM_TOKEN")
FOLDER_ID = os.getenv("FOLDER_ID")
OAUTH_TOKEN = os.getenv("OAUTH_TOKEN")


STATS_FILE = "bot_statistics.csv"

def get_iam_token():
    response = requests.post(
        'https://iam.api.cloud.yandex.net/iam/v1/tokens',
        json={'yandexPassportOauthToken': OAUTH_TOKEN}
    )

    response.raise_for_status()
    return response.json()['iamToken']

def init_stats_file():
    if not os.path.exists(STATS_FILE):
        with open(STATS_FILE, 'w', newline='') as file:
            writer = csv.writer(file)
            writer.writerow(['timestamp', 'user_id', 'username', 'status'])
        logging.info(f"Statistics file created at {STATS_FILE}")

def save_statistics(user_id, username, type="success"):
    with open(STATS_FILE, 'a', newline='') as file:
        writer = csv.writer(file)
        timestamp = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        writer.writerow([timestamp, user_id, username, type])

async def start(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
    """Send a welcome message when the command /start is issued."""
    user = update.effective_user
    await update.message.reply_html(
        f"Hi {user.mention_html()}! I'm your AI assistant powered by YandexGPT. Ask me anything!"
    )
    
    save_statistics(
        user_id=user.id,
        username=user.username or "unknown",
        type="command_start"
    )

async def process_message(update: Update, context: CallbackContext) -> None:
    user_text = update.message.text

    iam_token = get_iam_token()

    data = {}
    data["modelUri"] = f"gpt://{FOLDER_ID}/yandexgpt"
    data["completionOptions"] = {"temperature": 0.3, "maxTokens": 1000}
    data["messages"] = [
        {"role": "user", "text": f"{user_text}"},
    ]

    URL = "https://llm.api.cloud.yandex.net/foundationModels/v1/completion"
    response = requests.post(
        URL,
        headers={
            "Accept": "application/json",
            "Authorization": f"Bearer {iam_token}"
        },
        json=data,
    ).json()

    print(response)

    save_statistics(
        user_id=update.message.from_user.id,
        username=update.effective_user.username,
        type="Message"
    )

    answer = response.get('result', {})\
                     .get('alternatives', [{}])[0]\
                     .get('message', {})\
                     .get('text', {})

    await update.message.reply_text(answer)

def main() -> None:
    """Start the bot."""
    init_stats_file()

    if not TELEGRAM_TOKEN:
        print("No telegram token.")
        return
    
    print(TELEGRAM_TOKEN)
    application = Application.builder().token(TELEGRAM_TOKEN).build()

    application.add_handler(CommandHandler("start", start))
    
    application.add_handler(MessageHandler(filters.TEXT & ~filters.COMMAND, process_message))

    application.run_polling()
    logging.info("Bot started")

if __name__ == "__main__":
    main()