import os
import sys

import argparse
import stripe
from collections import defaultdict, deque
import re
import logging
from dataclasses import dataclass, field, asdict

import os
from dotenv import load_dotenv
import json
import asyncio

from tqdm import tqdm
from datetime import datetime, timedelta
import time

import json
from urllib.parse import urlparse
from urllib.parse import quote_plus


load_dotenv(dotenv_path=".env")  # idk if required

working_dir = os.getcwd()

project_root = os.path.abspath(
    os.path.join(os.path.dirname(__file__), "../../studio_api")
)
os.chdir(project_root)
sys.path.append(project_root)

os.environ.setdefault("DJANGO_SETTINGS_MODULE", "studio_api.settings")
os.environ[
    "DJANGO_ALLOW_ASYNC_UNSAFE"
] = "true"  # https://docs.djangoproject.com/en/4.1/topics/async/#async-safety

database_url = f"postgresql://{os.getenv('DATABASE_USER')}:{quote_plus(os.getenv('DATABASE_PASSWORD'))}@{os.getenv('DATABASE_HOST')}/{os.getenv('DATABASE_NAME')}"
replica_database_url = database_url

print("Using database:")
print(f"- host: {os.getenv('DATABASE_HOST')}")
print(f"- username: {os.getenv('DATABASE_USER')}")
print(
    f"- password: {'*' * len(os.getenv('DATABASE_PASSWORD')) if os.getenv('DATABASE_PASSWORD') else 'None'}"
)
print(f"- database: {os.getenv('DATABASE_NAME')}")

if not database_url or not replica_database_url:
    raise ValueError(
        "DATABASE_URL and REPLICA_DATABASE_URL must be set in environment variables"
    )

os.environ["DATABASE_URL"] = database_url
os.environ["REPLICA_DATABASE_URL"] = replica_database_url

import django

django.setup()

from django.contrib.auth.models import Group, User
from studio_api.bots import models
from django.utils import timezone
from django.db.models import F


def load_file(file_path):
    with open(file_path, "r") as file:
        return [s.strip() for s in file.readlines()]


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Add users to botjail")
    parser.add_argument(
        "--user-id-file",
        type=str,
        required=True,
        help="Path to file containing user IDs (one per line)",
    )
    parser.add_argument(
        "--copycat-name",
        type=str,
        required=True,
        help="Name of the copycat app/bot to store in the DB for reference.",
    )
    parser.add_argument(
        "--execute",
        action="store_true",
        help="Actually execute changes. Without this flag, runs in dry-run mode",
    )

    args = parser.parse_args()

    should_execute = args.execute

    # combine working directory with user id file path using os.path.join
    user_id_file_path = os.path.join(working_dir, args.user_id_file)

    # load user ids
    user_ids = load_file(user_id_file_path)
    user_ids = [int(user_id) for user_id in user_ids]

    print(f"Loaded {len(user_ids)} user ids from {user_id_file_path}")

    BATCH_SIZE = 50

    total_batches = len(user_ids) // BATCH_SIZE + 1

    total_updated = 0
    total_new = 0

    # prompt for user confirmation if we are actually executing
    if should_execute:
        user_input = input(
            f"Are you sure you want to add {len(user_ids)} users to botjail? (y/n): "
        )
        if user_input.lower() != "y":
            print("Aborting.")
            exit(0)

    for i in range(0, len(user_ids), BATCH_SIZE):
        batch_user_ids = user_ids[i : i + BATCH_SIZE]
        batch_number = i // BATCH_SIZE + 1
        print(f"Processing batch {batch_number} of {total_batches}")

        existing_botjailed_user_ids = models.BotJail.objects.filter(  # .exclude(status=models.BotJail.Status.RELEASED)
            user_id__in=batch_user_ids
        ).values_list("user_id", flat=True)

        if len(existing_botjailed_user_ids) > BATCH_SIZE:
            raise Exception(
                f"Found {len(existing_botjailed_user_ids)} botjailed users, more than expected. This should not happen."
            )

        # find existing users and update their entries, excluding released users
        existing_botjailed_users = models.BotJail.objects.exclude(
            status=models.BotJail.Status.RELEASED
        ).filter(user_id__in=existing_botjailed_user_ids)

        total_updated += len(existing_botjailed_users)

        if should_execute:
            now = timezone.now()

            # update existing botjailed users to copycat so they get watermarked
            existing_botjailed_users.update(
                status=models.BotJail.Status.JAILED,
                total_jailed_count=F("total_jailed_count") + 1,
                latest_jailed_reason=models.BotJail.Reason.COPYCAT,
                latest_jailed_at=now,
                total_copycat_violation_count=F("total_copycat_violation_count") + 1,
                copycat_name=args.copycat_name.strip(),
                latest_copycat_violation_at=now,
            )
            for botjail_user in existing_botjailed_users:
                print(
                    f"Updated {botjail_user.user_id} to be botjailed for copycat {args.copycat_name}."
                )
        else:
            print(
                f"Would have updated {len(existing_botjailed_users)} existing botjailed users"
            )
            for botjail_user in existing_botjailed_users:
                print(
                    f"Would have updated {botjail_user.user_id} ({botjail_user.user.username} - {botjail_user.user.email}) to be botjailed for copycat {args.copycat_name}. Current botjail status: {botjail_user.latest_jailed_reason}"
                )

        # add new users to botjail
        new_user_ids_to_jail = set(batch_user_ids) - set(existing_botjailed_user_ids)

        botjail_entries_to_create = []

        users_to_jail = User.objects.filter(id__in=new_user_ids_to_jail)

        now = timezone.now()

        for user in users_to_jail:
            botjail_entries_to_create.append(
                models.BotJail(
                    username=user.username,
                    user=user,
                    status=models.BotJail.Status.JAILED,
                    total_jailed_count=1,
                    latest_jailed_reason=models.BotJail.Reason.COPYCAT,
                    latest_jailed_at=now,
                    total_copycat_violation_count=1,
                    copycat_name=args.copycat_name.strip(),
                    latest_copycat_violation_at=now,
                )
            )

        total_new += len(botjail_entries_to_create)

        if should_execute:
            models.BotJail.objects.bulk_create(botjail_entries_to_create)
            for user in users_to_jail:
                print(f"Added {user.id} to botjail for copycat {args.copycat_name}")
        else:
            print(
                f"Would have added {len(botjail_entries_to_create)} new botjailed users"
            )
            for user in users_to_jail:
                print(
                    f"Would have added {user.id} ({user.username} - {user.email}) to botjail for copycat {args.copycat_name}"
                )

        # time.sleep(1)

    print(
        f"🎶🎶🎶🎶🎶📣🤖🤖🤖 Updated {total_updated} existing botjailed users and added {total_new} new botjailed users. 🎶🎶🎶🎶🎶📣🤖🤖🤖"
    )