import json
import os
import pickle as pkl
import re
import shutil
import string
from collections import Counter
from pathlib import Path

import numpy as np
import torch
from transformers import AutoModel, AutoTokenizer


from pathlib import Path

# Core Application URL
SERVER_URL = "http://localhost:8000/"

# Maximum length for user queries
MAX_USER_QUERY_LEN = 128

# Base Directories
CURRENT_DIR = Path(__file__).parent
DEPLOYMENT_DIR = CURRENT_DIR / "deployment"
DATA_PATH = CURRENT_DIR / "files"

# Deployment Directories
CLIENT_DIR = DEPLOYMENT_DIR / "client_dir"
SERVER_DIR = DEPLOYMENT_DIR / "server_dir"
KEYS_DIR = DEPLOYMENT_DIR / ".fhe_keys"

# All Directories
ALL_DIRS = [KEYS_DIR, CLIENT_DIR, SERVER_DIR]

# Model and Data Files
LOGREG_MODEL_PATH = CURRENT_DIR / "models" / "cml_logreg.model"
ORIGINAL_FILE_PATH = DATA_PATH / "original_document.txt"
ANONYMIZED_FILE_PATH = DATA_PATH / "anonymized_document.txt"
MAPPING_UUID_PATH = DATA_PATH / "original_document_uuid_mapping.json"
MAPPING_ANONYMIZED_SENTENCES_PATH = DATA_PATH / "mapping_clear_to_anonymized.pkl"
MAPPING_ENCRYPTED_SENTENCES_PATH = DATA_PATH / "mapping_clear_to_encrypted.pkl"
MAPPING_DOC_EMBEDDING_PATH = DATA_PATH / "mapping_doc_embedding_path.pkl"

PROMPT_PATH = DATA_PATH / "chatgpt_prompt.txt"


# List of example queries for easy access
DEFAULT_QUERIES = {
    "Example Query 1": "What is the amount of the contract between David and Kate?",
    "Example Query 2": "How many people are engaged in the contract?",
    "Example Query 3": "Does Kate have an international bank account?",
}

# Load tokenizer and model
TOKENIZER = AutoTokenizer.from_pretrained("obi/deid_roberta_i2b2")
EMBEDDINGS_MODEL = AutoModel.from_pretrained("obi/deid_roberta_i2b2")

PUNCTUATION_LIST = list(string.punctuation)
PUNCTUATION_LIST.remove("%")
PUNCTUATION_LIST.remove("$")
PUNCTUATION_LIST = "".join(PUNCTUATION_LIST) + '°'


def clean_directory() -> None:
    """Clear direcgtories"""

    print("Cleaning...\n")
    for target_dir in ALL_DIRS:
        if os.path.exists(target_dir) and os.path.isdir(target_dir):
            shutil.rmtree(target_dir)
        target_dir.mkdir(exist_ok=True, parents=True)


def get_batch_text_representation(texts, model, tokenizer, batch_size=1):
    """Get mean-pooled representations of given texts in batches."""
    mean_pooled_batch = []
    for i in range(0, len(texts), batch_size):
        batch_texts = texts[i : i + batch_size]
        inputs = tokenizer(batch_texts, return_tensors="pt", padding=True, truncation=True)
        with torch.no_grad():
            outputs = model(**inputs, output_hidden_states=False)
        last_hidden_states = outputs.last_hidden_state
        input_mask_expanded = (
            inputs["attention_mask"].unsqueeze(-1).expand(last_hidden_states.size()).float()
        )
        sum_embeddings = torch.sum(last_hidden_states * input_mask_expanded, 1)
        sum_mask = input_mask_expanded.sum(1)
        mean_pooled = sum_embeddings / sum_mask
        mean_pooled_batch.extend(mean_pooled.cpu().detach().numpy())
    return np.array(mean_pooled_batch)


def is_user_query_valid(user_query: str) -> bool:
    """
    Check if the `user_query` is None and not empty.
    Args:
        user_query (str): The input text to be checked.
    Returns:
        bool: True if the `user_query` is None or empty, False otherwise.
    """
    # If the query is not part of the default queries
    is_default_query = user_query in DEFAULT_QUERIES.values()

    # Check if the query exceeds the length limit
    is_exceeded_max_length = user_query is not None and len(user_query) <= MAX_USER_QUERY_LEN

    return not is_default_query and not is_exceeded_max_length


def compare_texts_ignoring_extra_spaces(original_text, modified_text):
    """Check if the modified_text is identical to the original_text except for additional spaces.

    Args:
        original_text (str): The original text for comparison.
        modified_text (str): The modified text to compare against the original.

    Returns:
        (bool): True if the modified_text is the same as the original_text except for
            additional spaces; False otherwise.
    """
    normalized_original = " ".join(original_text.split())
    normalized_modified = " ".join(modified_text.split())

    return normalized_original == normalized_modified


def is_strict_deletion_only(original_text, modified_text):

    # Define a regex pattern that matches a word character next to a punctuation
    # or a punctuation next to a word character, without a space between them.
    pattern = r"(?<=[\w])(?=[^\w\s])|(?<=[^\w\s])(?=[\w])"

    # Replace instances found by the pattern with a space
    original_text = re.sub(pattern, " ", original_text)
    modified_text = re.sub(pattern, " ", modified_text)

    # Tokenize the texts into words, considering also punctuation
    original_words = Counter(original_text.lower().split())
    modified_words = Counter(modified_text.lower().split())

    base_words = all(item in original_words.keys() for item in modified_words.keys())
    base_count = all(original_words[k] >= v for k, v in modified_words.items())

    return base_words and base_count


def read_txt(file_path):
    """Read text from a file."""
    with open(file_path, "r", encoding="utf-8") as file:
        return file.read()


def write_txt(file_path, data):
    """Write text to a file."""
    with open(file_path, "w", encoding="utf-8") as file:
        file.write(data)


def write_pickle(file_path, data):
    """Save data to a pickle file."""
    with open(file_path, "wb") as f:
        pkl.dump(data, f)


def read_pickle(file_name):
    """Load data from a pickle file."""
    with open(file_name, "rb") as file:
        return pkl.load(file)


def read_json(file_name):
    """Load data from a json file."""
    with open(file_name, "r") as file:
        return json.load(file)


def write_json(file_name, data):
    """Save data to a json file."""
    with open(file_name, "w", encoding="utf-8") as file:
        json.dump(data, file, indent=4, sort_keys=True)


def write_bytes(path, data):
    """Save binary data."""
    with path.open("wb") as f:
        f.write(data)


def read_bytes(path):
    """Load data from a binary file."""
    with path.open("rb") as f:
        return f.read()