import json
import re
from datetime import date, datetime
from typing import Any


def parse_json_markdown(json_string: str) -> dict:
    # Remove the triple backticks if present
    json_string = json_string.strip()
    start_index = json_string.find("```json")
    end_index = json_string.find("```", start_index + len("```json"))

    if start_index != -1 and end_index != -1:
        extracted_content = json_string[start_index + len("```json") : end_index].strip()

        # Parse the JSON string into a Python dictionary
        parsed = json.loads(extracted_content)
    elif start_index != -1 and end_index == -1 and json_string.endswith("``"):
        end_index = json_string.find("``", start_index + len("```json"))
        extracted_content = json_string[start_index + len("```json") : end_index].strip()

        # Parse the JSON string into a Python dictionary
        parsed = json.loads(extracted_content)
    elif json_string.startswith("{"):
        # Parse the JSON string into a Python dictionary
        parsed = json.loads(json_string)
    else:
        raise Exception("Could not find JSON block in the output.")

    return parsed


def parse_and_check_json_markdown(text: str, expected_keys: list[str]) -> dict[str, Any]:
    try:
        json_obj = parse_json_markdown(text)
    except json.JSONDecodeError as e:
        raise ValueError(f"Got invalid JSON object. Error: {e}") from e

    for key in expected_keys:
        if key not in json_obj:
            raise ValueError(
                f"Got invalid return object. Expected key `{key}` to be present, but got {json_obj}"
            ) from None
    return json_obj


def json_serial(obj):
    """JSON serializer for objects not serializable by default json code."""
    if isinstance(obj, (datetime, date)):
        return obj.isoformat()
    raise TypeError("Type %s not serializable" % type(obj))


def to_json(obj):
    return json.dumps(obj, default=json_serial)


def extract_json_str(text: str) -> str:
    """Extract JSON string from text."""
    # NOTE: this regex parsing is taken from langchain.output_parsers.pydantic
    match = re.search(r"\{.*\}", text.strip(), re.MULTILINE | re.IGNORECASE | re.DOTALL)
    if not match:
        match = re.search(r"\[.*\]", text.strip(), re.MULTILINE | re.IGNORECASE | re.DOTALL)
        if not match:
            raise ValueError(f"Could not extract json string from output: {text}")
    return match.group()


def make_pydantic_schema_strict_json(schema: Any) -> Any:
    """Used to verifiy if a pydantic class host others pydantic class."""
    if isinstance(schema, dict):
        if schema.get("type") == "object" or "properties" in schema or "anyOf" in schema:
            schema["additionalProperties"] = False
        if "properties" in schema:
            schema["required"] = list(schema["properties"].keys())
        for key, value in schema.items():
            schema[key] = make_pydantic_schema_strict_json(value)
    elif isinstance(schema, list):
        return [make_pydantic_schema_strict_json(item) for item in schema]
    return schema
