from json import JSONDecodeError, JSONDecoder
from json.decoder import WHITESPACE
from typing import Any, Dict, List, Literal, Optional, Tuple, Union

import orjson
import partial_json_parser
from partial_json_parser.core.options import Allow

from sglang.srt.entrypoints.openai.protocol import Tool, ToolChoice

_STANDARD_JSON_SCHEMA_TYPES = {
    "null",
    "boolean",
    "object",
    "array",
    "number",
    "string",
    "integer",
}

# Non-standard ``type`` values commonly emitted by DB/ORM-driven tool-schema
# generators. Mapped to the closest JSON Schema 2020-12 primitive so that
# ``Draft202012Validator.check_schema`` does not reject an otherwise-usable
# tool definition.
_JSON_SCHEMA_TYPE_ALIASES: Dict[str, str] = {
    "str": "string",
    "text": "string",
    "varchar": "string",
    "char": "string",
    "enum": "string",
    "uuid": "string",
    "date": "string",
    "datetime": "string",
    "time": "string",
    "timestamp": "string",
    "binary": "string",
    "blob": "string",
    "bytea": "string",
    "bytes": "string",
    "varbinary": "string",
    "bool": "boolean",
    "bigint": "integer",
    "smallint": "integer",
    "tinyint": "integer",
    "double": "number",
    "decimal": "number",
    "real": "number",
    "numeric": "number",
    "arr": "array",
    "tuple": "array",
    "set": "array",
    "map": "object",
}

# Prefix-based matching so that parameterised names like ``int32`` /
# ``float64`` / ``list[str]`` / ``dict[str, int]`` resolve. A prefix only
# matches when it spans the entire token or is followed by a non-identifier
# char, so "int" does not swallow "internal" and "list" does not swallow
# "list_price".
_PREFIX_BOUNDARY_CHARS = frozenset("0123456789[<( \t")
_PREFIX_RULES: Tuple[Tuple[Tuple[str, ...], str], ...] = (
    (("int", "uint", "long", "short", "unsigned"), "integer"),
    (("num", "float"), "number"),
    (("list",), "array"),
    (("dict",), "object"),
)


def _matches_type_prefix(base: str, prefixes: Tuple[str, ...]) -> bool:
    for p in prefixes:
        if base == p:
            return True
        if (
            len(base) > len(p)
            and base.startswith(p)
            and base[len(p)] in _PREFIX_BOUNDARY_CHARS
        ):
            return True
    return False


def _normalize_single_type(raw: Any) -> Any:
    if not isinstance(raw, str):
        return raw
    if raw in _STANDARD_JSON_SCHEMA_TYPES:
        return raw
    # ``split("(", 1)[0]`` strips parenthesized params like ``varchar(255)``
    # or ``decimal(10,2)`` without the overhead of a regex per call.
    base = raw.split("(", 1)[0].strip().lower()
    if base in _STANDARD_JSON_SCHEMA_TYPES:
        return base
    mapped = _JSON_SCHEMA_TYPE_ALIASES.get(base)
    if mapped is not None:
        return mapped
    for prefixes, target in _PREFIX_RULES:
        if _matches_type_prefix(base, prefixes):
            return target
    return raw


def _normalize_type_list(raw_items: List[Any]) -> List[Any]:
    normalized_items: List[Any] = []
    for item in raw_items:
        normalized_item = _normalize_single_type(item)
        if normalized_item not in normalized_items:
            normalized_items.append(normalized_item)
    return normalized_items


def normalize_json_schema_types(schema: Any) -> None:
    """
    Walk a JSON Schema in place and rewrite non-standard ``"type"`` values
    (e.g. ``"varchar"``, ``"enum"``, ``"int"``) to their standard JSON Schema
    equivalents.

    Acts as a compatibility layer for tool ``parameters`` schemas exported
    from database / ORM tooling, which often uses DB type names rather than
    JSON Schema types. Unknown types are left untouched so that downstream
    validation can still surface genuine errors.

    Mutates the input dict in place; the rewritten schema is also what gets
    rendered into the model prompt, so e.g. a user-supplied ``"varchar"``
    reaches the model as ``"string"``. ``$ref`` values are not resolved;
    callers pass tree-shaped schemas (HTTP JSON input is always a tree).
    """
    if isinstance(schema, list):
        for item in schema:
            normalize_json_schema_types(item)
        return
    if not isinstance(schema, dict):
        return

    if "type" in schema:
        t = schema["type"]
        if isinstance(t, str):
            schema["type"] = _normalize_single_type(t)
        elif isinstance(t, list):
            schema["type"] = _normalize_type_list(t)

    for key in (
        "properties",
        "patternProperties",
        "$defs",
        "definitions",
        "dependentSchemas",
    ):
        nested = schema.get(key)
        if isinstance(nested, dict):
            for v in nested.values():
                normalize_json_schema_types(v)

    for key in ("anyOf", "oneOf", "allOf", "prefixItems"):
        nested = schema.get(key)
        if isinstance(nested, list):
            for v in nested:
                normalize_json_schema_types(v)

    for key in (
        "items",
        "additionalProperties",
        "not",
        "if",
        "then",
        "else",
        "contains",
        "propertyNames",
        "unevaluatedItems",
        "unevaluatedProperties",
    ):
        if key in schema:
            normalize_json_schema_types(schema[key])


def _find_common_prefix(s1: str, s2: str) -> str:
    prefix = ""
    min_length = min(len(s1), len(s2))
    for i in range(0, min_length):
        if s1[i] == s2[i]:
            prefix += s1[i]
        else:
            break
    return prefix


def _partial_json_loads(input_str: str, flags: Allow) -> Tuple[Any, int]:
    """
    Parse incomplete or partial JSON strings commonly encountered during streaming.

    Args:
        input_str (str): The potentially incomplete JSON string to parse.
        flags (Allow): Bitwise flags controlling what types of partial data are allowed.
            Common flags include:
            - Allow.STR: Allow partial strings (e.g., '"hello wo' -> 'hello wo')
            - Allow.OBJ: Allow partial objects (e.g., '{"key":' -> {'key': None})
            - Allow.ARR: Allow partial arrays (e.g., '[1, 2,' -> [1, 2])
            - Allow.ALL: Allow all types of partial data

    Returns:
        Tuple[Any, int]: A tuple containing:
            - parsed_object: The Python object parsed from the JSON
            - consumed_length: Number of characters consumed from input_str
    """
    try:
        return (partial_json_parser.loads(input_str, flags), len(input_str))
    except (JSONDecodeError, IndexError) as e:
        msg = getattr(e, "msg", str(e))
        if "Extra data" in msg or "pop from empty list" in msg:
            start = WHITESPACE.match(input_str, 0).end()
            obj, end = JSONDecoder().raw_decode(input_str, start)
            return obj, end
        raise


def _is_complete_json(input_str: str) -> bool:
    try:
        orjson.loads(input_str)
        return True
    except JSONDecodeError:
        return False


def _get_tool_schema_defs(tools: List[Tool]) -> dict:
    """
    Get consolidated $defs from all tools, validating for conflicts.

    Args:
        tools: List of tools to process

    Returns:
        Dictionary of consolidated $defs from all tools

    Raises:
        ValueError: If conflicting $defs are found
    """
    all_defs = {}
    for tool in tools:
        if tool.function.parameters is None:
            continue
        defs = tool.function.parameters.get("$defs", {})
        for def_name, def_schema in defs.items():
            if def_name in all_defs and all_defs[def_name] != def_schema:
                raise ValueError(
                    f"Tool definition '{def_name}' has "
                    "multiple schemas, which is not "
                    "supported."
                )
            else:
                all_defs[def_name] = def_schema
    return all_defs


def _get_tool_schema(tool: Tool) -> dict:
    return {
        "properties": {
            "name": {"type": "string", "enum": [tool.function.name]},
            "parameters": (
                tool.function.parameters
                if tool.function.parameters
                else {"type": "object", "properties": {}}
            ),
        },
        "required": ["name", "parameters"],
    }


def infer_type_from_json_schema(schema: Dict[str, Any]) -> Optional[str]:
    """
    Infer the primary type of a parameter from JSON Schema.

    Supports complex JSON Schema structures including:
    - Direct type field (including type arrays)
    - anyOf/oneOf: parameter can be any of multiple types
    - enum: parameter must be one of enum values
    - allOf: parameter must satisfy all type definitions
    - properties: inferred as object type
    - items: inferred as array type

    Args:
        schema: JSON Schema definition

    Returns:
        Inferred type ('string', 'number', 'object', 'array', etc.) or None
    """
    if not isinstance(schema, dict):
        return None

    # Priority 1: Direct type field (including type arrays)
    if "type" in schema:
        type_value = schema["type"]
        if isinstance(type_value, str):
            return type_value
        elif isinstance(type_value, list) and type_value:
            # Handle type arrays: return first non-null type
            non_null_types = [t for t in type_value if t != "null"]
            if non_null_types:
                return non_null_types[0]
            return "string"  # If only null, default to string

    # Priority 2: Handle anyOf/oneOf
    if "anyOf" in schema or "oneOf" in schema:
        schemas = schema.get("anyOf") or schema.get("oneOf")
        types = []

        if isinstance(schemas, list):
            for sub_schema in schemas:
                inferred_type = infer_type_from_json_schema(sub_schema)
                if inferred_type:
                    types.append(inferred_type)

            if types:
                # If all types are the same, return unified type
                if len(set(types)) == 1:
                    return types[0]
                # When types differ, prioritize string (safest)
                if "string" in types:
                    return "string"
                # Otherwise return first type
                return types[0]

    # Priority 3: Handle enum (infer type from enum values)
    if "enum" in schema and isinstance(schema["enum"], list):
        if not schema["enum"]:
            return "string"

        # Infer type from enum values
        enum_types = set()
        for value in schema["enum"]:
            if value is None:
                enum_types.add("null")
            elif isinstance(value, bool):
                enum_types.add("boolean")
            elif isinstance(value, int):
                enum_types.add("integer")
            elif isinstance(value, float):
                enum_types.add("number")
            elif isinstance(value, str):
                enum_types.add("string")
            elif isinstance(value, list):
                enum_types.add("array")
            elif isinstance(value, dict):
                enum_types.add("object")

        # If type is uniform, return that type
        if len(enum_types) == 1:
            return enum_types.pop()
        # Mixed types, prioritize string
        return "string"

    # Priority 4: Handle allOf (must satisfy all types)
    if "allOf" in schema and isinstance(schema["allOf"], list):
        schemas = schema["allOf"]
        for sub_schema in schemas:
            inferred_type = infer_type_from_json_schema(sub_schema)
            if inferred_type and inferred_type != "string":
                return inferred_type
        return "string"

    # Priority 5: Infer object type
    if "properties" in schema:
        return "object"

    # Priority 6: Infer array type
    if "items" in schema:
        return "array"

    return None


def get_json_schema_constraint(
    tools: List[Tool],
    tool_choice: Union[ToolChoice, Literal["required"]],
    parallel_tool_calls: bool = True,
) -> Optional[dict]:
    """
    Get the JSON schema constraint for the specified tool choice.

    Args:
        tool_choice: The tool choice specification
        parallel_tool_calls: If False, constrain to exactly one tool call (maxItems=1)

    Returns:
        JSON schema dict, or None if no valid tools found
    """

    if isinstance(tool_choice, ToolChoice):
        # For specific function choice, return the user's parameters schema directly
        fn_name = tool_choice.function.name
        for tool in tools:
            if tool.function.name == fn_name:
                schema = {
                    "type": "array",
                    "minItems": 1,
                    "items": _get_tool_schema(tool),
                }
                if not parallel_tool_calls:
                    schema["maxItems"] = 1
                return schema
        return None
    elif tool_choice == "required":
        json_schema = {
            "type": "array",
            "minItems": 1,
            "items": {
                "type": "object",
                "anyOf": [_get_tool_schema(tool) for tool in tools],
            },
        }
        if not parallel_tool_calls:
            json_schema["maxItems"] = 1
        json_schema_defs = _get_tool_schema_defs(tools)
        if json_schema_defs:
            json_schema["$defs"] = json_schema_defs
        return json_schema

    return None