# MIT License
#
# Copyright (c) 2025, Loong Ma
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

import json
import os
import re
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, Any, List
import pandas as pd
import yaml
from datasets import Dataset
from deprecated import deprecated
from fastmcp.tools import Tool
from pydantic import BaseModel


class DefaultToolFnSchema(BaseModel):
    """Default tool function Schema."""

    input: str


@dataclass
class ToolMetadata:
    description: str
    name: str | None = None
    fn_schema: type[BaseModel] | None = DefaultToolFnSchema
    return_direct: bool = False

    def get_parameters_dict(self) -> dict:
        if self.fn_schema is None:
            parameters = {
                "type": "object",
                "properties": {
                    "input": {"title": "input query string", "type": "string"},
                },
                "required": ["input"],
            }
        else:
            parameters = self.fn_schema.model_json_schema()
            parameters = {
                k: v for k, v in parameters.items() if k in ["type", "properties", "required", "definitions", "$defs"]
            }
        return parameters

    @property
    def fn_schema_str(self) -> str:
        """Get fn schema as string."""
        if self.fn_schema is None:
            raise ValueError("fn_schema is None.")
        parameters = self.get_parameters_dict()
        return json.dumps(parameters, ensure_ascii=False)

    def get_name(self) -> str:
        """Get name."""
        if self.name is None:
            raise ValueError("name is None.")
        return self.name

    @deprecated("Deprecated in favor of `to_openai_tool`, which should be used instead.")
    def to_openai_function(self) -> dict[str, Any]:
        """
        Deprecated and replaced by `to_openai_tool`.
        The name and arguments of a function that should be called, as generated by the
        model.
        """
        return {
            "name": self.name,
            "description": self.description,
            "parameters": self.get_parameters_dict(),
        }

    def to_openai_tool(self, skip_length_check: bool = False) -> dict[str, Any]:
        """To OpenAI tool."""
        if not skip_length_check and len(self.description) > 1024:
            raise ValueError(
                "Tool description exceeds maximum length of 1024 characters. "
                "Please shorten your description or move it to the prompt."
            )
        return {
            "type": "function",
            "function": {
                "name": self.name,
                "description": self.description,
                "parameters": self.get_parameters_dict(),
            },
        }

    def to_gemini_tool(self) -> dict[str, Any]:
        """To Google Gemini tool."""
        return {
            "function_declarations": [
                {
                    "name": self.name,
                    "description": self.description,
                    "parameters": self.get_parameters_dict(),
                }
            ]
        }


def convert_to_mcp_tools(function_docs: list[dict[str, Any]]) -> dict[str, list[dict[str, Any]]]:
    tools = []
    for f in function_docs:
        func = f["function"]
        tool = {
            "name": func["name"],
            "description": func.get("description", ""),
            "input_schema": func.get("parameters", {}),
        }
        tools.append(tool)
    return {"tools": tools}


def convert_to_openai_tools(tools: list[Tool]) -> dict[str, list[dict[str, Any]]]:
    functions = []
    for tool in tools:
        function = {
            "type": "function",
            "function": {
                "name": tool.name,
                "description": tool.description or tool.name,
                "parameters": tool.inputSchema or {},
            },
        }
        functions.append(function)
    return {"tools": functions}


def convert_to_gemini_tools(tools: list[Tool]) -> dict[str, list[dict[str, Any]]]:
    functions = []
    for tool in tools:
        function = {
            "function_declarations": [
                {
                    "name": tool.name,
                    "description": tool.description or tool.name,
                    "parameters": tool.inputSchema or {},
                }
            ]
        }
        functions.append(function)
    return {"tools": functions}


def parse_hermes_tool_calls(message: Any) -> List[Dict[str, Any]]:
    """
    Parse one or multiple Hermes-style tool calls.

    Returns:
        [
            {
                "name": str,
                "arguments": dict
            },
            ...
        ]
    """

    tool_calls = []

    # ---------- Case 1: string with <tool_call> blocks ----------
    print("message: ", message)
    if isinstance(message, str):
        print("message: ", message)
        matches = re.findall(
            r"<tool_call>\s*(\{.*?\})\s*</tool_call>",
            message,
            re.S
        )
        print("matches: ", matches)

        for block in matches:
            try:
                payload = json.loads(block)
                args = payload.get("arguments", {})
                tool_calls.append({
                    "name": payload["name"],
                    "arguments": json.loads(args) if isinstance(args, str) else args
                })
            except json.JSONDecodeError as e:
                raise ValueError(f"Invalid JSON in tool_call: {block}") from e

        if tool_calls:
            return tool_calls

        # ---------- Fallback: {"tool_call": {...}} or {"tool_calls": [...]} ----------
        try:
            data = json.loads(message)
            if "tool_calls" in data:
                return [
                    {
                        "name": tc["name"],
                        "arguments": tc.get("arguments", {})
                    }
                    for tc in data["tool_calls"]
                ]
            if "tool_call" in data:
                return [{
                    "name": data["tool_call"]["name"],
                    "arguments": data["tool_call"].get("arguments", {})
                }]
        except json.JSONDecodeError:
            pass

    # ---------- Case 2: dict input ----------
    if isinstance(message, dict):
        # OpenAI-style
        if "tool_calls" in message and message["tool_calls"]:
            print("tool_calls: ", message["tool_calls"])
            for tool in message["tool_calls"]:
                function = tool.get("function", {})
                args = function.get("arguments", {})
                if isinstance(args, str):
                    args = json.loads(args)
                tool_calls.append({
                    "name": function["name"],
                    "arguments": args
                })
            return tool_calls

        # OpenAI legacy
        if "function_call" in message and message["function_call"]:
            fc = message["function_call"]
            args = fc.get("arguments", {})
            if isinstance(args, str):
                args = json.loads(args)
            return [{
                "name": fc["name"],
                "arguments": args
            }]

        # Pure Hermes JSON
        if "name" in message and "arguments" in message:
            return [{
                "name": message["name"],
                "arguments": message["arguments"]
            }]

        # hermes raw content
        if "content" in message:
            return parse_hermes_tool_calls(message["content"])

    return []


def extract_format(format: str = "json", content: str = "") -> Any:
    import json
    import re

    if format == "json":
        # Strategy 1: Look for ```json ... ``` code blocks
        pattern = re.compile(r"```json\s*([\s\S]*?)\s*```")
        match = pattern.search(content)
        if match:
            try:
                data = json.loads(match.group(1))
                return data
            except json.JSONDecodeError:
                pass
        
        # Strategy 2: Look for ``` ... ``` code blocks (without language specifier)
        pattern_plain = re.compile(r"```\s*([\s\S]*?)\s*```")
        match_plain = pattern_plain.search(content)
        if match_plain:
             try:
                data = json.loads(match_plain.group(1))
                return data
             except json.JSONDecodeError:
                pass

        # Strategy 3: Try to parse the entire content
        try:
            return json.loads(content)
        except json.JSONDecodeError:
            pass

        # Strategy 4: Try to find the first JSON object using regex
        # This regex looks for { ... } ensuring balanced braces is hard with regex, 
        # but we can try to find the first { and the last }
        try:
            start_idx = content.find('{')
            end_idx = content.rfind('}')
            if start_idx != -1 and end_idx != -1 and end_idx > start_idx:
                json_candidate = content[start_idx : end_idx + 1]
                return json.loads(json_candidate)
        except json.JSONDecodeError:
            pass
            
    elif format == "":
        pass
    return None


def tool_format_convert(mcp_tools: list[dict[str, Any]], model: str) -> list[dict[str, Any]]:
    if model.startswith("gpt"):
        return convert_to_openai_tools(mcp_tools)
    if model.startswith("gemini"):
        return convert_to_gemini_tools(mcp_tools)
    if model.startswith("claude"):
        return mcp_tools
    return convert_to_openai_tools(mcp_tools)


def read_yaml(path: Path) -> dict[str, Any]:
    if not path.exists():
        return {}
    if yaml is None:
        raise RuntimeError("PyYAML is required to read YAML config files. Install with `pip install pyyaml`.")
    with open(path, encoding="utf-8") as f:
        return yaml.safe_load(f) or {}


def read_json_file(file_path):
    """
    读取 JSON 文件
    :param file_path: JSON 文件路径
    :return: 解析后的字典/列表数据
    """
    try:
        if not os.path.exists(file_path):
            raise FileNotFoundError(f"文件 {file_path} 不存在")
        with open(file_path, encoding="utf-8") as f:
            data = json.load(f)
        print("JSON 文件读取成功")
        return data
    except json.JSONDecodeError as e:
        print(f"JSON 文件格式错误: {e}")
    except Exception as e:
        print(f"读取 JSON 文件时出错: {e}")


def write_json_file(data, file_path):
    """
    写入 JSON 文件
    :param data: 要写入的数据（字典或列表）
    :param file_path: JSON 文件路径
    """
    try:
        with open(file_path, "w", encoding="utf-8") as f:
            json.dump(data, f, ensure_ascii=False, indent=4)
        print("JSON 文件写入成功")
    except Exception as e:
        print(f"写入 JSON 文件时出错: {e}")


def read_jsonl(file_path):
    """通用的读取jsonl文件为列表字典格式的数据"""
    data = []
    with open(file_path, encoding="utf-8") as fin:
        for line in fin:
            if line.strip():  # 跳过空行
                data.append(json.loads(line))
    if len(data) > 1:
        if isinstance(data[0], dict):
            print("data[0].keys()", data[0].keys())
    return data


def read_jsonl_to_df(file_path):
    data = []
    with open(file_path, encoding="utf-8") as fin:
        for line in fin:
            if line.strip():  # 跳过空行
                data.append(json.loads(line))
    if len(data) > 1:
        if isinstance(data[0], dict):
            print("data[0].keys()", data[0].keys())
    df = pd.DataFrame(data)
    print(df.shape, df.columns)
    return df


def format_value_counts(df, key_column):
    # compute counts and proportions
    counts = df[key_column].value_counts()
    proportions = df[key_column].value_counts(normalize=True)

    # merge counts and proportions into a DataFrame
    result = pd.DataFrame(
        {
            "count": counts,
            "proportion": proportions,
        }
    ).reset_index()

    # rename columns
    result.columns = [key_column, "count", "proportion"]

    # optional: keep two decimal places
    result["proportion"] = result["proportion"].apply(lambda x: f"{x * 100:.2f}%")

    total_count = result["count"].sum()

    # add total row
    total_row = pd.DataFrame({key_column: ["Total"], "count": [total_count], "proportion": ["100%"]})

    # merge results
    result = pd.concat([result, total_row], ignore_index=True)

    return result


def pick_unique(
    dataset: list[dict] | pd.DataFrame | Dataset, field: str, k: int = 4
) -> list[dict] | pd.DataFrame | Dataset:
    """
    select K samples from dataset where the specific field values are all different.

    parameters：
        dataset: list[dict] / pandas.DataFrame / HuggingFace Dataset
        field: the unique field name
        k:  select how many elements

    """
    # list of dict
    if isinstance(dataset, list):
        seen = set()
        result = []
        for item in dataset:
            value = item[field]
            if value not in seen:
                seen.add(value)
                result.append(item)
            if len(result) == k:
                break
        return result

    # pandas.DataFrame
    elif isinstance(dataset, pd.DataFrame):
        return dataset.drop_duplicates(subset=[field]).head(k)

    # HuggingFace Dataset
    elif Dataset is not None and isinstance(dataset, Dataset):
        df = dataset.to_pandas().drop_duplicates(subset=[field]).head(k)
        return Dataset.from_pandas(df)

    else:
        raise TypeError("Unsupported dataset type: must be list[dict], pandas.DataFrame, or HuggingFace Dataset.")




def detect_language(input_text):
    if isinstance(input_text, str):
        s = input_text.strip()
        # German check: e.g., contains ü, ö, ä, ß or typical German words
        if re.search(r"[üöäß]|[ÄÖÜ]", s) or re.search(
            r"\b(der|die|das|und|ein|nicht|ist|zu|in|den|mit|auf|für|sie|es|dem|von)\b", s, re.I
        ):
            return "ger"
        # Japanese check: hiragana, katakana, or frequent Japanese-specific patterns
        # Only classify as Japanese if contains **hiragana or katakana** (NOT CJK, which includes Chinese)
        if re.search(r"[\u3040-\u309F\u30A0-\u30FF]", s):
            return "jap"
        # English check: contains [a-zA-Z] and high ASCII ratio, avoid misclassifying numbers/symbols
        if re.search(r"[a-zA-Z]", s):
            ascii_ratio = sum(1 for c in s if ord(c) < 128) / max(1, len(s))
            if ascii_ratio > 0.7:
                return "en"
    # Default: assume Chinese if not above
    return "zh"


# 填补language列为空的行
def complete_language_column(row):
    if isinstance(row["language"], str) and row["language"].strip():
        return row["language"]
    else:
        return detect_language(row["input"])


async def get_mcp_tools(mcp_cfg: dict) -> list[dict]:
    """Get tools from MCP server."""
    from fastmcp import Client

    mcp_cfg = mcp_cfg
    client = Client(**mcp_cfg)
    async with client:
        tools = await client.list_tools()
    return tools


if __name__ == "__main__":
    data = extract_format(content='```json\n{"a": 1}\n```')
    print(f"extract format: {data}")

    # example 
    function_docs = [
        {
            "type": "function",
            "function": {
                "name": "get_weather",
                "description": "Retrieves current weather for the given location.",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "location": {"type": "string", "description": "City and country e.g. Bogotá, Colombia"},
                        "units": {
                            "type": "string",
                            "enum": ["celsius", "fahrenheit"],
                            "description": "Units the temperature will be returned in.",
                        },
                    },
                    "required": ["location", "units"],
                    "additionalProperties": False,
                },
                "strict": True,
            },
        },
        {
            "type": "function",
            "function": {
                "name": "get_local_time",
                "description": "Get the local time of a given location",
                "strict": True,
                "parameters": {
                    "type": "object",
                    "required": ["location", "timezone"],
                    "properties": {
                        "location": {
                            "type": "string",
                            "description": "The name or coordinates of the location for which to get the local time",
                        },
                        "timezone": {
                            "type": "string",
                            "description": "The timezone of the location, \
                                defaults to the location's timezone if not provided",
                        },
                    },
                    "additionalProperties": False,
                },
            },
        },
    ]

    mcp_tools = convert_to_mcp_tools(function_docs)

    import json

    print(json.dumps(mcp_tools, indent=2, ensure_ascii=False))
    with open("function_docs.json", "w") as f:
        json.dump(mcp_tools, f, indent=2, ensure_ascii=False)
