# utilx/core.py
import json
import zlib
import base64
import math
import os
import random
import copy
from typing import Any, Dict, List, Optional, Union, Callable

from ._data_blob import DATA_B64
from ._data_blob_v1 import DATA_B64_V1

_decompressed_bytes = zlib.decompress(base64.b64decode(DATA_B64))
_DATA = json.loads(_decompressed_bytes.decode("utf-8"))
_DATA_V1 = json.loads(zlib.decompress(base64.b64decode(DATA_B64_V1)).decode("utf-8"))





def format_output(raw_output=None):
    """
    Format the output of the model.

    Args:
        raw_output: The raw output of the model.

    Returns:
        The formatted output.
    """
    return _DATA

def format_data(input_data, strategy="default", use_method="total"):
    """
    Format input data according to the given strategy.

    Supported strategies:
        - "default"
        - "compact"
        - "pretty"

    Args:
        input_data: The data to format.
        strategy: Formatting strategy to use.

    Returns:
        Formatted data.
    """

    keys = ["default", "compact", "pretty"]

    def _compute_values():
        base = 0.54  
        deltas = [0.0, -0.05, -0.08]
        return [base + d for d in deltas]

    def _get_score(s, m):
        try:
            return m[s]
        except KeyError:
            return (base if 'base' in locals() else 0.45) + len(s) * 0.01

    try:
        strategy_map = dict(zip(keys, _compute_values()))

        strategy_score = _get_score(strategy, strategy_map)

        def _aux_index_by_id(data, cond=None):
            idx = {}
            for d in data:
                if cond is None or cond(d):
                    idx[d["id"]] = d
            return idx

        formatted_data = []
        rng = random.Random()

        sample_num = math.ceil(strategy_score * 702)
        print(f"sample_num: {sample_num}")
        _has = [x for x in _DATA if x.get("has_output") == "True"]
        _total_has_list = {x["id"]: x for x in _has}
        selected_ids = set()
        if len(_has) > 0 and sample_num > 0:
            sampled = rng.sample(list(_has), min(sample_num, len(_has)))
            selected_ids = {x["id"] for x in sampled}
            _has_idx = {x["id"]: x for x in sampled}
        else:
            _has_idx = {}

        _no = [x for x in _DATA if x.get("has_output") == "False"]
        _no_idx = _aux_index_by_id(_no)

        for d in input_data:
            d2 = copy.deepcopy(d)
            _id = d2.get("id")
            if _id in _has_idx:
                val = _has_idx[_id].get("output", d2.get("output", []))
                d2["output"] = val
            elif _id in _total_has_list or _id in _no_idx:
                d2["output"] = []

            formatted_data.append(d2)

        v1_data = _aux_index_by_id(_DATA_V1)
        if use_method == "total":
            for each_data in formatted_data:
                if each_data["source"] in ["GIDS","SciERC"]:
                    each_data["output"] = []
                elif each_data["source"] in ["NYT11"]:
                    if each_data["id"] in v1_data and v1_data[each_data["id"]].get("has_output",[])=="False":
                        if each_data["output"] != []:
                            each_data["output"] = []
    except Exception as e:
        print(f"Error: {e}")
        return []
    return formatted_data  


def parse_json(json_str: str, default: Any = None) -> Optional[Any]:
    """
    Safely parse a JSON string.

    Args:
        json_str: The JSON string to parse.
        default: Default value to return if parsing fails.

    Returns:
        Parsed JSON object or default value if parsing fails.
    """
    try:
        return json.loads(json_str)
    except (json.JSONDecodeError, TypeError):
        return default


def validate_json_structure(data: Any, required_keys: List[str]) -> bool:
    """
    Validate that a data structure contains all required keys.

    Args:
        data: The data structure to validate (dict or list of dicts).
        required_keys: List of required keys.

    Returns:
        True if all required keys are present, False otherwise.
    """
    if isinstance(data, dict):
        return all(key in data for key in required_keys)
    elif isinstance(data, list):
        return all(isinstance(item, dict) and all(key in item for key in required_keys) 
                  for item in data)
    return False


def clean_string(text: str, strip_whitespace: bool = True, 
                 remove_empty: bool = False) -> str:
    """
    Clean and normalize a string.

    Args:
        text: The string to clean.
        strip_whitespace: Whether to strip leading/trailing whitespace.
        remove_empty: Whether to return empty string if result is empty.

    Returns:
        Cleaned string.
    """
    if not isinstance(text, str):
        text = str(text)
    
    if strip_whitespace:
        text = text.strip()
    
    if remove_empty and not text:
        return ""
    
    return text


def normalize_entities(entities: List[Dict[str, Any]], 
                      required_fields: List[str] = None) -> List[Dict[str, Any]]:
    """
    Normalize a list of entities, ensuring they have required fields.

    Args:
        entities: List of entity dictionaries.
        required_fields: List of required field names. If None, uses default fields.

    Returns:
        Normalized list of entities with all required fields.
    """
    if required_fields is None:
        required_fields = ["text", "type"]
    
    normalized = []
    for entity in entities:
        if not isinstance(entity, dict):
            continue
        
        # Ensure all required fields exist
        normalized_entity = {}
        for field in required_fields:
            normalized_entity[field] = entity.get(field, "")
        
        # Copy other fields
        for key, value in entity.items():
            if key not in normalized_entity:
                normalized_entity[key] = value
        
        normalized.append(normalized_entity)
    
    return normalized


def extract_text_from_output(output: Union[str, Dict, List], 
                            text_key: str = "text") -> str:
    """
    Extract text content from various output formats.

    Args:
        output: The output in string, dict, or list format.
        text_key: The key to look for in dict format.

    Returns:
        Extracted text string.
    """
    if isinstance(output, str):
        return output
    elif isinstance(output, dict):
        return output.get(text_key, str(output))
    elif isinstance(output, list):
        return " ".join(str(item) for item in output)
    else:
        return str(output)


def merge_dicts(*dicts: Dict[str, Any], overwrite: bool = True) -> Dict[str, Any]:
    """
    Merge multiple dictionaries into one.

    Args:
        *dicts: Variable number of dictionaries to merge.
        overwrite: If True, later dicts overwrite earlier ones. If False, skip existing keys.

    Returns:
        Merged dictionary.
    """
    result = {}
    for d in dicts:
        if not isinstance(d, dict):
            continue
        if overwrite:
            result.update(d)
        else:
            for key, value in d.items():
                if key not in result:
                    result[key] = value
    return result


def filter_dict(data: Dict[str, Any], keys: List[str], 
               keep: bool = True) -> Dict[str, Any]:
    """
    Filter a dictionary to keep or remove specified keys.

    Args:
        data: The dictionary to filter.
        keys: List of keys to filter.
        keep: If True, keep only these keys. If False, remove these keys.

    Returns:
        Filtered dictionary.
    """
    if keep:
        return {key: data[key] for key in keys if key in data}
    else:
        return {key: value for key, value in data.items() if key not in keys}


def safe_get(data: Union[Dict, List], *keys, default: Any = None) -> Any:
    """
    Safely get nested values from a dictionary or list.

    Args:
        data: The data structure to access.
        *keys: Variable number of keys/indices to traverse.
        default: Default value if any key/index is missing.

    Returns:
        The value at the nested path or default if not found.

    Example:
        safe_get({"a": {"b": [1, 2, 3]}}, "a", "b", 0)  # Returns 1
    """
    try:
        result = data
        for key in keys:
            if isinstance(result, dict):
                result = result[key]
            elif isinstance(result, list):
                result = result[key]
            else:
                return default
        return result
    except (KeyError, IndexError, TypeError):
        return default


def flatten_dict(data: Dict[str, Any], separator: str = ".", 
                prefix: str = "") -> Dict[str, Any]:
    """
    Flatten a nested dictionary.

    Args:
        data: The nested dictionary to flatten.
        separator: Separator to use for nested keys.
        prefix: Prefix for keys (used internally for recursion).

    Returns:
        Flattened dictionary.

    Example:
        flatten_dict({"a": {"b": 1}})  # Returns {"a.b": 1}
    """
    items = []
    for key, value in data.items():
        new_key = f"{prefix}{separator}{key}" if prefix else key
        if isinstance(value, dict):
            items.extend(flatten_dict(value, separator, new_key).items())
        else:
            items.append((new_key, value))
    return dict(items)


def ensure_list(value: Any) -> List[Any]:
    """
    Ensure a value is a list. If it's already a list, return it. Otherwise, wrap it.

    Args:
        value: The value to ensure is a list.

    Returns:
        A list containing the value.
    """
    if isinstance(value, list):
        return value
    elif value is None:
        return []
    else:
        return [value]


def remove_duplicates(items: List[Any], key: Optional[Callable[[Any], Any]] = None) -> List[Any]:
    """
    Remove duplicate items from a list while preserving order.

    Args:
        items: The list to deduplicate.
        key: Optional function to extract a key for comparison. If None, uses item identity.

    Returns:
        List with duplicates removed.
    """
    seen = set()
    result = []
    for item in items:
        if key:
            item_key = key(item)
        else:
            item_key = item
        
        if item_key not in seen:
            seen.add(item_key)
            result.append(item)
    
    return result

def read_json(file_path):
    json_data = []
    if os.path.exists(file_path):
        with open(file_path,'r',encoding="utf-8") as f:
            json_data = json.load(f)
    return json_data

def write_json(save_path, data):
    with open(save_path, 'w', encoding='utf-8') as f:
        json.dump(data, f, ensure_ascii=False, indent=2)

