from langchain_openai import OpenAIEmbeddings, AzureOpenAIEmbeddings
from akasha.utils.prompts.format import language_dict
from langchain_core.embeddings import Embeddings
from typing import Union, Callable, Tuple, List, Dict, Any
from langchain.schema import Document
import jieba
import json
import re
import opencc

jieba.setLogLevel(jieba.logging.INFO)  ## ignore logging jieba model information


cc = opencc.OpenCC("s2twp")


def separate_name(name: str):
    """separate type:name by ':'

    Args:
        **name (str)**: string with format "type:name" \n

    Returns:
        (str, str): res_type , res_name
    """
    sep = name.split(":")

    if len(sep) > 2:
        res_type = sep[0].lower()
        res_name = ":".join(sep[1:])
    elif len(sep) < 2:
        ### if the format type not equal to type:name ###
        res_type = sep[0].lower()
        res_name = ""
    else:
        res_type = sep[0].lower()
        res_name = sep[1]

    return res_type, res_name


def decide_embedding_type(embeddings: Embeddings) -> str:
    """check the embedding type and return the type:name

    Args:
        embeddings (Embeddings): embedding class

    Raises:
        Exception: _description_

    Returns:
        str: type:name
    """
    if isinstance(embeddings, OpenAIEmbeddings) or isinstance(
        embeddings, AzureOpenAIEmbeddings
    ):
        return "openai:" + embeddings.model
    else:
        from akasha.utils.models.gemi import gemini_embed

        if isinstance(embeddings, gemini_embed):
            return "gemini:" + embeddings.model_name

        from akasha.utils.models.custom import custom_embed

        if isinstance(embeddings, custom_embed):
            return embeddings.model_name

        else:
            from langchain_huggingface import HuggingFaceEmbeddings

            if isinstance(embeddings, HuggingFaceEmbeddings):
                return "hf:" + embeddings.model_name

            from langchain_community.embeddings import TensorflowHubEmbeddings

            if isinstance(embeddings, TensorflowHubEmbeddings):
                return "tf:" + embeddings.model_url

            else:
                raise Exception("can not find the embeddings type.")


def get_embedding_type_and_name(
    embeddings: Union[Embeddings, str, Callable],
) -> Tuple[str, str]:
    """get the type and name of the embeddings"""

    if callable(embeddings):
        embeddings_name = embeddings.__name__
    elif isinstance(embeddings, str):
        embeddings_name = embeddings
    else:
        embeddings_name = decide_embedding_type(embeddings)

    embed_type, embed_name = separate_name(embeddings_name)

    return embed_type, embed_name


def sim_to_trad(text: str) -> str:
    """convert simplified chinese to traditional chinese

    Args:
        **text (str)**: simplified chinese\n

    Returns:
        str: traditional chinese
    """
    global cc
    return cc.convert(text)


def get_doc_length(language: str, text: str) -> int:
    """calculate the length of terms in a giving Document

    Args:
        **language (str)**: 'ch' for chinese and 'en' for others, default 'ch'\n
        **doc (Document)**: Document object\n

    Returns:
        doc_length: int Docuemtn length
    """

    if "chinese" in language_dict[language]:
        doc_length = len(list(jieba.cut(text)))
    else:
        doc_length = len(text.split())
    return doc_length


def get_docs_length(language: str, docs: List[Document]) -> int:
    """calculate the total length of terms in giving documents

    Args:
        language (str): 'ch' for chinese and 'en' for others, default 'ch'\n
        docs (list): list of Documents\n

    Returns:
        docs_length: int total Document length
    """
    docs_length = 0
    for doc in docs:
        docs_length += get_doc_length(language, doc.page_content)
    return docs_length


def extract_json(text: str) -> Union[Dict[str, Any], List[Dict[str, Any]], None]:
    """
    Extract JSON data from text generated by LLMs.
    Handles both individual dictionaries and lists of dictionaries.

    Args:
        text (str): The text containing JSON-like data

    Returns:
        Union[Dict[str, Any], List[Dict[str, Any]], None]: Parsed JSON data or None if parsing fails
    """
    # Clean up the text by removing markdown formatting if present
    cleaned_text = re.sub(r"```(?:json)?\n", "", text)
    cleaned_text = re.sub(r"```", "", cleaned_text)

    # Try to extract content within square brackets first (for list of dictionaries)
    list_match = re.search(r"\[(.*?)\]", cleaned_text, re.DOTALL)

    if list_match:
        try:
            # Try to parse as a list of JSON objects
            json_str = list_match.group(0)
            # Fix common issues: replace single quotes with double quotes
            json_str = re.sub(r"'([^']*)':", r'"\1":', json_str)  # Fix keys
            json_str = re.sub(
                r":\s*\'([^\']*)\'", r': "\1"', json_str
            )  # Fix string values
            return json.loads(json_str)
        except json.JSONDecodeError:
            # If list parsing fails, we'll fall through to try single object parsing
            pass

    # Try to extract the most promising JSON object
    # This looks for balanced braces { } with content inside
    matches = []
    stack = []
    start = -1

    for i, char in enumerate(cleaned_text):
        if char == "{":
            if not stack:  # Start of a potential JSON object
                start = i
            stack.append("{")
        elif char == "}":
            if stack and stack[-1] == "{":
                stack.pop()
                if not stack:  # We've found a complete balanced {} section
                    matches.append(cleaned_text[start : i + 1])

    # Process matches from longest to shortest (assuming more complete is better)
    matches.sort(key=len, reverse=True)

    for match in matches:
        try:
            # Fix common issues that LLMs make in JSON generation
            json_str = match
            # Replace single quotes with double quotes for keys and string values
            json_str = re.sub(r"'([^']*)':", r'"\1":', json_str)  # Fix keys
            json_str = re.sub(
                r":\s*\'([^\']*)\'", r': "\1"', json_str
            )  # Fix string values
            # Remove trailing commas before closing brackets
            json_str = re.sub(r",\s*}", "}", json_str)
            json_str = re.sub(r",\s*]", "]", json_str)
            # Fix boolean values
            json_str = re.sub(r":\s*True", r": true", json_str)
            json_str = re.sub(r":\s*False", r": false", json_str)
            # Fix None values
            json_str = re.sub(r":\s*None", r": null", json_str)

            return json.loads(json_str)
        except json.JSONDecodeError:
            continue

    # If all attempts failed, return None
    return None


def extract_multiple_json(text: str) -> List[Dict[str, Any]]:
    """
    Extract multiple JSON objects from text generated by LLMs.

    Args:
        text (str): The text containing multiple JSON objects

    Returns:
        List[Dict[str, Any]]: List of parsed JSON objects
    """
    # Clean up the text
    cleaned_text = re.sub(r"```(?:json)?\n", "", text)
    cleaned_text = re.sub(r"```", "", cleaned_text)

    results = []
    matches = []
    stack = []
    start = -1

    # Find all potential JSON objects with balanced braces
    for i, char in enumerate(cleaned_text):
        if char == "{":
            if not stack:
                start = i
            stack.append("{")
        elif char == "}":
            if stack and stack[-1] == "{":
                stack.pop()
                if not stack:  # Complete balanced object found
                    matches.append(cleaned_text[start : i + 1])

    # Try to parse each potential JSON object
    for match in matches:
        try:
            # Apply the same fixes as in the single object function
            json_str = match
            json_str = re.sub(r"'([^']*)':", r'"\1":', json_str)
            json_str = re.sub(r":\s*\'([^\']*)\'", r': "\1"', json_str)
            json_str = re.sub(r",\s*}", "}", json_str)
            json_str = re.sub(r",\s*]", "]", json_str)
            json_str = re.sub(r":\s*True", r": true", json_str)
            json_str = re.sub(r":\s*False", r": false", json_str)
            json_str = re.sub(r":\s*None", r": null", json_str)

            parsed = json.loads(json_str)
            results.append(parsed)
        except json.JSONDecodeError:
            continue

    return results
