"""
Functions to read in the data for bilbystats
"""

from importlib import resources
from pathlib import Path
import pandas as pd
from dotenv import load_dotenv
import os
import socket
from datetime import datetime, timedelta, date
from zoneinfo import ZoneInfo

import json
import pandas as pd
from typing import List, Dict, Any
import pickle
import textwrap
import re

import warnings


def read_data_path(data_name):
    """
    Retrieve the file path for a given data file name.

    Depending on the value of `data_name`, this function constructs the 
    appropriate file path from the `bilbystats.data` or 
    `bilbystats.data.prompts` module.

    ---------------------------------------------------------------------------
    ARGUMENTS:
    data_name : str
        The name of the data file to locate. If it is 
        'simple_sentiment_prompt.txt', the function looks inside 
        'bilbystats.data.prompts'; otherwise, it defaults to 'bilbystats.data'.

    ---------------------------------------------------------------------------
    OUTPUT:
    path : str
        The resolved file path as a string.
    ---------------------------------------------------------------------------
    AUTHORS: Samuel Davenport
    ---------------------------------------------------------------------------
    """
    if data_name == 'simple_sentiment_prompt.txt':
        path = str(resources.files(
            "bilbystats.data.prompts").joinpath("simple_sentiment_prompt.txt"))
    elif data_name == 'sentiment_classification.txt':
        path = str(resources.files(
            "bilbystats.data.prompts").joinpath("sentiment_classification.txt"))
    else:
        path = str(resources.files("bilbystats.data").joinpath(data_name))

    return path

def bs_path():
    path = str(resources.files("bilbystats"))
    return path


def read_data(data_name):
    """
    Load data from a specified file based on its extension.

    This function reads a data file given its name by first resolving its path
    using `read_data_path`. It supports reading `.txt`, `.csv`, and `.parquet` 
    files, and returns the contents accordingly.

    ---------------------------------------------------------------------------
    ARGUMENTS:
    data_name : str
        The name of the data file to read. The file extension must be one of 
        '.txt', '.csv', '.parquet' or '.json'.

    ---------------------------------------------------------------------------
    OUTPUT:
    data : str or pandas.DataFrame
        The contents of the file. Returns a string if the file is a `.txt`,
        and a pandas DataFrame if it is a `.csv` or `.parquet`.

    ---------------------------------------------------------------------------
    AUTHORS: Samuel Davenport
    ---------------------------------------------------------------------------
    """
    try:
        path = read_data_path(data_name)

        file_ending = Path(path).suffix
        if file_ending == '.txt':
            with open(path, 'r') as file:
                data = file.read()
        elif file_ending == '.csv':
            data = pd.read_csv(path)
        elif file_ending == '.parquet':
            data = pd.read_parquet(path)
        elif file_ending == '.json':
            with open(path, 'r') as f:
                data = json.load(f)
        elif file_ending == '.html':
            with open(path, 'r', encoding='utf-8') as f:
                data = f.read()
                
    except:
        path = data_name
        file_ending = Path(path).suffix
        if file_ending == '.txt':
            with open(path, 'r') as file:
                data = file.read()
        elif file_ending == '.csv':
            data = pd.read_csv(path)
        elif file_ending == '.parquet':
            data = pd.read_parquet(path)
        elif file_ending == '.json':
            with open(path, 'r') as f:
                data = json.load(f)
        elif file_ending == '.html':
            with open(path, 'r', encoding='utf-8') as f:
                data = f.read()
            
    return data

def save2json(data, saveloc, encoding="utf-8"):
    with open(saveloc, "w", encoding=encoding) as f:
        json.dump(data, f, ensure_ascii=False, indent=4)

def read_txt_file(filepath, as_lines=False, strip_lines=True, encoding="utf-8"):
    """
    Reads a .txt file and returns its contents.

    Parameters:
        filepath (str): Path to the text file.
        as_lines (bool): If True, returns a list of lines. If False, returns a full string.
        strip_lines (bool): If True and as_lines=True, strips whitespace from each line.
        encoding (str): Encoding used to open the file (default: 'utf-8').

    Returns:
        str or list[str]: File contents as a string or list of lines.
    """
    with open(filepath, "r", encoding=encoding) as file:
        if as_lines:
            lines = file.readlines()
            if strip_lines:
                return [line.strip() for line in lines]
            return lines
        else:
            return file.read()


def check_dir(addon: str = ''):
    """
    Determine the model checkpoint directory path with optional suffix.

    This function attempts to load a local `.env` file to retrieve the 
    `MODEL_CHECKPOINTS_DIR` environment variable. If the variable is found, 
    it appends the optional `addon` string to the path. If not, it defaults 
    to the current directory with the `addon` appended.

    ---------------------------------------------------------------------------
    ARGUMENTS:
    addon : str, optional (default='')
        A string to append to the resolved directory path.

    ---------------------------------------------------------------------------
    OUTPUT:
    checkpoint_path : str
        The resolved checkpoint directory path, either from the environment 
        variable or defaulting to the current directory.
    ---------------------------------------------------------------------------
    AUTHORS: Samuel Davenport
    ---------------------------------------------------------------------------
    """
    try:
        env_path = resources.files(
            "bilbystats.defaults").joinpath("local_defaults.env")

        # Load environment variables from the .env file
        load_dotenv(dotenv_path=env_path)

        checkpoint_path = os.getenv("MODEL_CHECKPOINTS_DIR") + addon
    except:
        # If a local checkpoint path is not set defaults to the current directory.
        warnings.warn("MODEL_CHECKPOINTS_DIR is not set in bilbystats/defaults/local_defaults.env. So reverting to using " + addon + " as the directory." + '''
                      To configure MODEL_CHECKPOINTS_DIR you need to run
                      cp bilbystats/defaults/local_defaults_example.env bilbystats/defaults/local_defaults.env
                      and set MODEL_CHECKPOINTS_DIR within local_defaults.env''')
        checkpoint_path = addon

    return checkpoint_path

def inprod(addon):
    try:
        env_path = resources.files(
            "bilbystats.defaults").joinpath("local_defaults.env")

        # Load environment variables from the .env file
        load_dotenv(dotenv_path=env_path)

        checkpoint_path = os.getenv("PRODUCTION_STAGE")
    except:
        # If a local checkpoint path is not set defaults to the current directory.
        warnings.warn("MODEL_CHECKPOINTS_DIR is not set in bilbystats/defaults/local_defaults.env. So reverting to using " + addon + " as the directory." + '''
                      To configure MODEL_CHECKPOINTS_DIR you need to run
                      cp bilbystats/defaults/local_defaults_example.env bilbystats/defaults/local_defaults.env
                      and set MODEL_CHECKPOINTS_DIR within local_defaults.env''')
        checkpoint_path = addon

    return checkpoint_path    

def json2df(json_strings: List[str], columns = None, orig_df = None, rw = False, direct_eval = True) -> pd.DataFrame:
    """
    Convert a list of JSON strings to a pandas DataFrame with arbitrary columns.

    This function parses JSON strings and creates a structured DataFrame with columns
    based on the keys found in the JSON objects. All JSON strings should have the same
    set of keys. It handles parsing errors gracefully by inserting None values for 
    malformed JSON entries.

    ---------------------------------------------------------------------------
    ARGUMENTS:
    json_strings : List[str]
        List of JSON strings, each containing the same set of key-value pairs.
        Expected format: '{"key1": "value1", "key2": "value2", ...}'

    ---------------------------------------------------------------------------
    OUTPUT:
    df : pd.DataFrame
        DataFrame with columns corresponding to the keys in the JSON objects.

    ---------------------------------------------------------------------------
    AUTHORS: Samuel Davenport
    ---------------------------------------------------------------------------
    """
    if not json_strings:
        return pd.DataFrame()

    # Parse first JSON to get column names
    if not columns:
        try:
            first_string = json_strings[0]
            if direct_eval:
                first_data = eval(first_string)
            else:
                first_data = extract_json(first_string)
            columns = list(first_data.keys())
        except json.JSONDecodeError as e:
            a = 4

    # Initialize dictionary to store column data
    data_dict = {col: [] for col in columns}

    for i, json_str in enumerate(json_strings):
        try:
            # Parse the JSON string
            if direct_eval:
                parsed_data = eval(json_str)
            else:
                parsed_data = extract_json(json_str)
            #parsed_data = extract_json(json_str)

            # Extract values for each column
            for col in columns:
                value = parsed_data.get(col, None)
                data_dict[col].append(value)

        except json.JSONDecodeError as e:
            # Add None values for all columns for failed parsing
            for col in columns:
                data_dict[col].append(None)
        except Exception as e:
            print(f"Warning: Unexpected error at index {i}: {e}")
            #print(f"Warning: Unexpected error at index {i}: {e}")
            for col in columns:
                data_dict[col].append(None)

    # Create DataFrame
    df = pd.DataFrame(data_dict)
    
    if orig_df != None:
        out = pd.concat([df, orig_df], axis=1)
    else:
        out = df

    return out

def extract_json(json_str: str):
    """
    Parse a JSON string after removing code block markers.

    This function removes markdown-style code block markers from a JSON 
    string and parses it into a Python dictionary or list using the `json` module.

    ---------------------------------------------------------------------------
    ARGUMENTS:
    json_str : str
        A string containing JSON data, potentially wrapped in markdown-style 
        code block markers (e.g., ```json ... ```).

    ---------------------------------------------------------------------------
    OUTPUT:
    parsed_data : Any
        The Python object (usually a dict or list) resulting from parsing the JSON string.
    ---------------------------------------------------------------------------
    AUTHORS: Samuel Davenport
    ---------------------------------------------------------------------------
    """
    if type(json_str) == list:
        out = []
        for i in range(len(json_str)):
            try:
                out.append(extract_json(json_str[i]))
            except:
                out.append(None)
        return out
    else:
        json_str = json_str.replace('```json', '')
        json_str = json_str.replace('```', '')
        parsed_data = json.loads(json_str)
        return parsed_data

def extract_json_dep(json_str: str):
    """
    Parse a JSON string after removing code block markers.

    This function removes markdown-style code block markers from a JSON 
    string and parses it into a Python dictionary or list using the `json` module.

    ---------------------------------------------------------------------------
    ARGUMENTS:
    json_str : str
        A string containing JSON data, potentially wrapped in markdown-style 
        code block markers (e.g., ```json ... ```).

    ---------------------------------------------------------------------------
    OUTPUT:
    parsed_data : Any
        The Python object (usually a dict or list) resulting from parsing the JSON string.
    ---------------------------------------------------------------------------
    AUTHORS: Samuel Davenport
    ---------------------------------------------------------------------------
    """
    json_str = json_str.replace('```json', '')
    json_str = json_str.replace('```', '')
    json_str = fix_json_string(json_str)
    parsed_data = json.loads(json_str)
    key_list = list(parsed_data.keys())
    for i in range(len(key_list)):
        parsed_data[key_list[i]] = undo_extra_newlines(parsed_data[key_list[i]])
    return parsed_data

def load_pickle(filename: str):
    """
    Load a Python object from a pickle file.

    This function opens a file in binary read mode and loads the 
    serialized Python object stored in it using the `pickle` module.

    ---------------------------------------------------------------------------
    ARGUMENTS:
    filename : str
        Path to the pickle file to be loaded.

    ---------------------------------------------------------------------------
    OUTPUT:
    loaded_data : Any
        The Python object deserialized from the pickle file.
    ---------------------------------------------------------------------------
    AUTHORS: Samuel Davenport
    ---------------------------------------------------------------------------
    """
    with open(filename, "rb") as f:
        loaded_data = pickle.load(f)
        
    return loaded_data


def printa(text: str, width: int = 110, use_bold = False):
    """
    Print text with line wrapping to a specified width.

    This function wraps each line of the input text to the specified width 
    and prints the formatted text to the console.

    ---------------------------------------------------------------------------
    ARGUMENTS:
    text : str
        The input text to be printed.
    width : int, optional (default=120)
        The maximum number of characters per line.

    ---------------------------------------------------------------------------
    OUTPUT:
    None
        This function prints the formatted text and does not return a value.
    ---------------------------------------------------------------------------
    AUTHORS: Samuel Davenport
    ---------------------------------------------------------------------------
    """
    if use_bold:
        BOLD = '\033[1m'
        END = '\033[0m'
        text = BOLD + text + END
    
    wrapped_text = "\n".join(textwrap.fill(line, width=width) for line in text.splitlines())
    print(wrapped_text)

def fix_json_string(json_str):
    # Find the content between quotes and escape only for JSON parsing
    def escape_for_json_only(match):
        key = match.group(1)
        value = match.group(2)
        # Escape only what's needed for valid JSON - newlines become \\n temporarily
        escaped_value = (value.replace('\\', '\\\\')
                              .replace('\n', '\\n')
                              .replace('\r', '\\r')
                              .replace('\t', '\\t')
                              .replace('"', '\\"'))
        return f'"{key}": "{escaped_value}"'
    
    # Pattern to match key-value pairs with quoted strings
    pattern = r'"([^"]+)":\s*"(.*?)"(?=\s*[,}])'
    fixed = re.sub(pattern, escape_for_json_only, json_str, flags=re.DOTALL)
    
    return fixed

def undo_extra_newlines(text):
    text = (text.replace('\\n', '\n')     # Convert \n to actual newline
            .replace('\\r', '\r')         # Convert \r to actual carriage return  
            .replace('\\t', '\t')         # Convert \t to actual tab
            .replace('\\\\', '\\'))       # Convert \\ to single backslash (must be last!)
    return text


def extract_json_braces(text):
    match = re.search(r'\{.*\}', text, re.DOTALL)
    return match.group(0) if match else None

def list2dict(lst, keys=None):
    if keys:
        return dict(zip(keys, lst))
    else:
        return {i: v for i, v in enumerate(lst)}

def list2df(data, columns):
    dictionary = list2dict(data, columns)
    df = pd.DataFrame(dictionary)
    return df

def dataframe_to_dataset(df):
    """Convert DataFrame to Hugging Face Dataset with tokens and NER tags"""
    sentences = []
    ner_tags = []
    
    for sent_id in df['sentence_id'].unique():
        sent_df = df[df['sentence_id'] == sent_id]
        sentences.append(sent_df['word'].tolist())
        ner_tags.append(sent_df['ner'].tolist())
    
    return sentences, ner_tags

def invert_dict(original_dict):
    inverted_dict = {value: key for key, value in original_dict.items()}
    return inverted_dict

def none_list_comb(a,b):
    out = []
    b_index = 0
    for item in a:
        if item is not None:
            out.append(b[b_index])
            b_index += 1
        else:
            out.append(None)
            
    return out

def filter_nones(a,b):
    out = []
    for i in range(len(a)):
        if a[i] is not None:
            out.append(b[i])
    return out

def count_col_none(df, covariate):
    out = df[covariate].isna().sum()
    return out

def count_none(df):
    return df.isna().sum()

def count_empty_strings(df):
    return (df == '').sum()

def filter_df(df, covariate, strings2include, strings2exclude = [], reset_index = True, case = False):
    """
    Filter a DataFrame based on string inclusion and exclusion criteria.
    
    Parameters:
    -----------
    df : pandas.DataFrame
        The input DataFrame to filter
    covariate : str
        The column name to apply filters on
    strings2include : list of str
        List of strings where at least one must be present in the covariate value (OR logic)
    strings2exclude : list of str
        List of strings that must not be present in the covariate value
    
    Returns:
    --------
    pandas.DataFrame
        Filtered DataFrame containing only rows that meet all criteria
    """
    # Start with no rows included
    inclusion_mask = pd.Series([False] * len(df), index=df.index)
    
    # Apply inclusion filters - ANY string can be present (OR logic)
    for string in strings2include:
        inclusion_mask |= df[covariate].astype(str).str.contains(string, case=case, na=False)
    
    # If no inclusion strings provided, include all rows
    if len(strings2include) == 0:
        inclusion_mask = pd.Series([True] * len(df), index=df.index)
    
    # Start with inclusion mask
    final_mask = inclusion_mask
    
    # Apply exclusion filters - NONE of these strings should be present
    for string in strings2exclude:
        final_mask &= ~df[covariate].astype(str).str.contains(string, case=case, na=False)
        
    out = df[final_mask]
    if reset_index:
        out.reset_index(drop=True, inplace=True)

    return out

def save_pickle(data, savename, saveloc):
    with open(saveloc + savename + '.pickle', 'wb') as file:
        pickle.dump(data, file)

def internet_connected(host: str = "8.8.8.8", port: int = 53, timeout: float = 3.0) -> bool:
    """
    Return True if the machine can reach *host* on *port* within *timeout* seconds.
    Default: Google public DNS (8.8.8.8) on the DNS port (53).  
    This works even if the machine has no HTTP proxy configured.
    """
    try:
        # Socket with a short timeout → returns quickly on failure
        socket.setdefaulttimeout(timeout)
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        sock.connect((host, port))
        sock.close()
        return True
    except OSError:
        return False
    
def get_date(n, date_format="%d-%m-%Y", tz=None):
    now = datetime.now(ZoneInfo(tz)) if tz else datetime.now()
    return (now - timedelta(days=-n)).strftime(date_format)

def day_of_week(text: bool = False, tz: str = None) -> int | str:
    """
    Return the current day of the week.
    
    Args:
        text: If True, return the day name. If False, return 0-6 (Monday=0, Sunday=6).
        tz: IANA timezone string (e.g., "America/New_York", "Europe/London").
            If None, uses the system's local timezone.
    
    Returns:
        The day as an integer (0-6) or string ("Monday"-"Sunday").
    """
    if tz:
        today = datetime.now(ZoneInfo(tz)).date()
    else:
        today = date.today()
    
    if text:
        return today.strftime("%A")
    
    return today.weekday()