# import json
# import ujson
# import logging
# import inspect
# import typing
# import pydantic

# from functools import lru_cache
# from pydantic.fields import FieldInfo
# from typing import Annotated, Callable, List, Tuple, Union  # noqa: UP035

# import dspy
# from dspy.dsp.adapters import passages2text
# from dspy.primitives.prediction import Prediction
# from dspy.signatures.signature import ensure_signature, make_signature

# @lru_cache(maxsize=None)
# def warn_once(msg: str):
#     logging.warning(msg)


# def predictor(*args: tuple, **kwargs) -> Callable[..., dspy.Module]:
#     def _predictor(func) -> dspy.Module:
#         """Decorator that creates a predictor module based on the provided function."""
#         signature = _func_to_signature(func)
#         *_, output_key = signature.output_fields.keys()
#         return _StripOutput(TypedPredictor(signature, **kwargs), output_key)

#     # if we have only a single callable argument, the decorator was invoked with no key word arguments
#     #  so we just return the wrapped function
#     if len(args) == 1 and callable(args[0]) and len(kwargs) == 0:
#         return _predictor(args[0])
#     return _predictor


# def cot(*args: tuple, **kwargs) -> Callable[..., dspy.Module]:
#     def _cot(func) -> dspy.Module:
#         """Decorator that creates a chain of thought module based on the provided function."""
#         signature = _func_to_signature(func)
#         *_, output_key = signature.output_fields.keys()
#         return _StripOutput(TypedChainOfThought(signature, **kwargs), output_key)

#     # if we have only a single callable argument, the decorator was invoked with no key word arguments
#     #  so we just return the wrapped function
#     if len(args) == 1 and callable(args[0]) and len(kwargs) == 0:
#         return _cot(args[0])
#     return _cot


# class _StripOutput(dspy.Module):
#     def __init__(self, predictor, output_key):
#         super().__init__()
#         self.predictor = predictor
#         self.output_key = output_key

#     def copy(self):
#         return _StripOutput(self.predictor.copy(), self.output_key)

#     def forward(self, **kwargs):
#         prediction = self.predictor(**kwargs)
#         return prediction[self.output_key]


# class FunctionalModule(dspy.Module):
#     """To use the @cot and @predictor decorators, your module needs to inherit form this class."""

#     def __init__(self):
#         super().__init__()
#         for name in dir(self):
#             attr = getattr(self, name)
#             if isinstance(attr, dspy.Module):
#                 self.__dict__[name] = attr.copy()


# def TypedChainOfThought(signature, instructions=None, reasoning=None, *, max_retries=3) -> dspy.Module:  # noqa: N802
#     """Just like TypedPredictor, but adds a ChainOfThought OutputField."""
#     signature = ensure_signature(signature, instructions)
#     output_keys = ", ".join(signature.output_fields.keys())

#     default_rationale = dspy.OutputField(
#         prefix="Reasoning: Let's think step by step in order to",
#         desc="${produce the " + output_keys + "}. We ...",
#     )
#     reasoning = reasoning or default_rationale

#     return TypedPredictor(
#         signature.prepend(
#             "reasoning",
#             reasoning,
#         ),
#         max_retries=max_retries,
#     )


# class TypedPredictor(dspy.Module):
#     def __init__(self, signature, instructions=None, *, max_retries=3, wrap_json=False, explain_errors=False):
#         """Like dspy.Predict, but enforces type annotations in the signature.

#         Args:
#             signature: The signature of the module. Can use type annotations.
#             instructions: A description of what the model should do.
#             max_retries: The number of times to retry the prediction if the output is invalid.
#             wrap_json: If True, json objects in the input will be wrapped in ```json ... ```
#             explain_errors: If True, the model will try to explain the errors it encounters.
#         """
#         super().__init__()

#         # Warn: deprecation warning.
#         warn_once(
#                 "\t*** Since DSPy 2.5.16+, TypedPredictors are now deprecated, underperform, and are about to be removed! ***\n"
#                 "Please use standard predictors, e.g. dspy.Predict and dspy.ChainOfThought.\n"
#                 "They now support type annotations and other features of TypedPredictors and "
#                 "tend to work much better out of the box.\n"
#                 "Please let us know if you face any issues: https://github.com/stanfordnlp/dspy/issues"
#             )

#         signature = ensure_signature(signature, instructions)
#         self.predictor = dspy.Predict(signature, _parse_values=False)
#         self.max_retries = max_retries
#         self.wrap_json = wrap_json
#         self.explain_errors = explain_errors

#     @property
#     def signature(self) -> dspy.Signature:
#         return self.predictor.signature

#     @signature.setter
#     def signature(self, value: dspy.Signature):
#         self.predictor.signature = value

#     def copy(self) -> "TypedPredictor":
#         return TypedPredictor(
#             self.signature,
#             max_retries=self.max_retries,
#             wrap_json=self.wrap_json,
#             explain_errors=self.explain_errors,
#         )

#     def __repr__(self):
#         """Return a string representation of the TypedPredictor object."""
#         return f"TypedPredictor({self.signature})"

#     def _make_example(self, field) -> str:
#         # Note: DSPy will cache this call so we only pay the first time TypedPredictor is called.
#         if hasattr(field, "model_json_schema"):
#             pass
#         schema = field.json_schema_extra["schema"]
#         parser = field.json_schema_extra["parser"]
#         if self.wrap_json:
#             schema = "```json\n" + schema + "\n```\n"
#         json_object = dspy.Predict(
#             make_signature(
#                 "json_schema -> json_object",
#                 "Make a very succinct json object that validates with the following schema",
#             ),
#             _parse_values=False,
#         )(json_schema=schema).json_object
#         # We use the parser to make sure the json object is valid.
#         try:
#             parser(_unwrap_json(json_object, parser))
#         except (pydantic.ValidationError, ValueError):
#             return ""  # Unable to make an example
#         return json_object
#         # TODO: Another fun idea is to only (but automatically) do this if the output fails.
#         # We could also have a more general "suggest solution" prompt that tries to fix the output
#         # More directly.
#         # TODO: Instead of using a language model to create the example, we can also just use a
#         # library like https://pypi.org/project/polyfactory/ that's made exactly to do this.

#     def _format_error(
#         self,
#         error: Exception,
#         task_description: Union[str, FieldInfo],
#         model_output: str,
#         lm_explain: bool,
#     ) -> str:
#         if isinstance(error, pydantic.ValidationError):
#             errors = []
#             for e in error.errors():
#                 fields = ", ".join(map(str, e["loc"]))
#                 errors.append(f"{e['msg']}: {fields} (error type: {e['type']})")
#             error_text = "; ".join(errors)
#         else:
#             error_text = repr(error)

#         if self.explain_errors and lm_explain:
#             if isinstance(task_description, FieldInfo):
#                 args = task_description.json_schema_extra
#                 task_description = args["prefix"] + " " + args["desc"]
#             return (
#                 error_text
#                 + "\n"
#                 + self._make_explanation(
#                     task_description=task_description,
#                     model_output=model_output,
#                     error=error_text,
#                 )
#             )

#         return error_text

#     def _make_explanation(self, task_description: str, model_output: str, error: str) -> str:
#         class Signature(dspy.Signature):
#             """I gave my language model a task, but it failed.

#             Figure out what went wrong, and write instructions to help it avoid the error next time.
#             """

#             task_description: str = dspy.InputField(desc="What I asked the model to do")
#             language_model_output: str = dspy.InputField(desc="The output of the model")
#             error: str = dspy.InputField(desc="The validation error triggered by the models output")
#             explanation: str = dspy.OutputField(desc="Explain what the model did wrong")
#             advice: str = dspy.OutputField(
#                 desc="Instructions for the model to do better next time. A single paragraph.",
#             )

#         # TODO: We could also try repair the output here. For example, if the output is a float, but the
#         # model returned a "float + explanation", the repair could be to remove the explanation.

#         return dspy.Predict(Signature)(
#             task_description=task_description,
#             language_model_output=model_output,
#             error=error,
#             _parse_values=False,
#         ).advice

#     def _prepare_signature(self) -> dspy.Signature:
#         """Add formats and parsers to the signature fields, based on the type annotations of the fields."""
#         signature = self.signature
#         for name, field in self.signature.fields.items():
#             is_output = field.json_schema_extra["__dspy_field_type"] == "output"
#             type_ = field.annotation
#             if is_output:
#                 if type_ is bool:

#                     def parse(x):
#                         x = x.strip().lower()
#                         if x not in ("true", "false"):
#                             raise ValueError("Respond with true or false")
#                         return x == "true"

#                     signature = signature.with_updated_fields(
#                         name,
#                         desc=field.json_schema_extra.get("desc", "")
#                         + (" (Respond with true or false)" if type_ != str else ""),
#                         format=lambda x: x if isinstance(x, str) else str(x),
#                         parser=parse,
#                     )
#                 elif type_ in (str, int, float):
#                     signature = signature.with_updated_fields(
#                         name,
#                         desc=field.json_schema_extra.get("desc", "")
#                         + (f" (Respond with a single {type_.__name__} value)" if type_ != str else ""),
#                         format=lambda x: x if isinstance(x, str) else str(x),
#                         parser=type_,
#                     )
#                 else:
#                     # Anything else we wrap in a pydantic object
#                     if (
#                         inspect.isclass(type_)
#                         and typing.get_origin(type_) not in (list, tuple)  # To support Python 3.9
#                         and issubclass(type_, pydantic.BaseModel)
#                     ):
#                         def to_json(x):
#                             return x.model_dump_json()
#                         def from_json(x, type_=type_):
#                             return type_.model_validate_json(x)
#                         schema = json.dumps(type_.model_json_schema())
#                     else:
#                         adapter = pydantic.TypeAdapter(type_)
#                         def to_json(x):
#                             return adapter.serializer.to_json(x)
#                         def from_json(x, type_=adapter):
#                             return type_.validate_json(x)
#                         schema = json.dumps(adapter.json_schema())
#                     if self.wrap_json:
#                         def to_json(x, inner=to_json):
#                             return "```json\n" + inner(x) + "\n```\n"
#                         schema = "```json\n" + schema + "\n```"
#                     signature = signature.with_updated_fields(
#                         name,
#                         desc=field.json_schema_extra.get("desc", "")
#                         + (". Respond with a single JSON object. JSON Schema: " + schema),
#                         format=lambda x, to_json=to_json: (x if isinstance(x, str) else to_json(x)),
#                         parser=lambda x, from_json=from_json: from_json(_unwrap_json(x, from_json)),
#                         schema=schema,
#                         type_=type_,
#                     )
#             else:  # If input field
#                 is_json = False
#                 def format_(x):
#                     return x if isinstance(x, str) else str(x)
#                 if type_ in (List[str], list[str], Tuple[str], tuple[str]):
#                     format_ = passages2text
#                 # Special formatting for lists of known types. Maybe the output fields sohuld have this too?
#                 elif typing.get_origin(type_) in (List, list, Tuple, tuple):
#                     (inner_type,) = typing.get_args(type_)
#                     if inspect.isclass(inner_type) and issubclass(inner_type, pydantic.BaseModel):
#                         def format_(x):
#                             return x if isinstance(x, str) else "[" + ",".join(i.model_dump_json() for i in x) + "]"
#                     else:
#                         def format_(x):
#                             return x if isinstance(x, str) else json.dumps(x)
#                     is_json = True
#                 elif inspect.isclass(type_) and issubclass(type_, pydantic.BaseModel):
#                     def format_(x):
#                         return x if isinstance(x, str) else x.model_dump_json()
#                     is_json = True
#                 if self.wrap_json and is_json:
#                     def format_(x, inner=format_):
#                         return x if isinstance(x, str) else "```json\n" + inner(x) + "\n```\n"
#                 signature = signature.with_updated_fields(name, format=format_)

#         return signature

#     def forward(self, **kwargs) -> dspy.Prediction:
#         modified_kwargs = kwargs.copy()
#         # We have to re-prepare the signature on every forward call, because the base
#         # signature might have been modified by an optimizer or something like that.
#         signature = self._prepare_signature()
#         for try_i in range(self.max_retries):
#             result = self.predictor(**modified_kwargs, new_signature=signature)
#             errors = {}
#             parsed_results = []
#             # Parse the outputs
#             for completion in result.completions:
#                 parsed = {}
#                 for name, field in signature.output_fields.items():
#                     try:
#                         value = completion[name]
#                         parser = field.json_schema_extra.get("parser", lambda x: x)
#                         parsed[name] = parser(value)
#                     except (pydantic.ValidationError, ValueError) as e:
#                         errors[name] = self._format_error(
#                             e,
#                             signature.fields[name],
#                             value,
#                             lm_explain=try_i + 1 < self.max_retries,
#                         )

#                         # If we can, we add an example to the error message
#                         current_desc = field.json_schema_extra.get("desc", "")
#                         i = current_desc.find("JSON Schema: ")
#                         if i == -1:
#                             continue  # Only add examples to JSON objects
#                         suffix, current_desc = current_desc[i:], current_desc[:i]
#                         prefix = "You MUST use this format: "
#                         if (
#                             try_i + 1 < self.max_retries
#                             and prefix not in current_desc
#                             and (example := self._make_example(field))
#                         ):
#                             signature = signature.with_updated_fields(
#                                 name,
#                                 desc=current_desc + "\n" + prefix + example + "\n" + suffix,
#                             )
#                 # No reason trying to parse the general signature, or run more completions, if we already have errors
#                 if errors:
#                     break
#                 # Instantiate the actual signature with the parsed values.
#                 # This allow pydantic to validate the fields defined in the signature.
#                 try:
#                     _ = self.signature(**kwargs, **parsed)
#                     parsed_results.append(parsed)
#                 except pydantic.ValidationError as e:
#                     errors["general"] = self._format_error(
#                         e,
#                         signature.instructions,
#                         "\n\n".join(
#                             "> " + field.json_schema_extra["prefix"] + " " + completion[name]
#                             for name, field in signature.output_fields.items()
#                         ),
#                         lm_explain=try_i + 1 < self.max_retries,
#                     )
#             if errors:
#                 # Add new fields for each error
#                 for name, error in errors.items():
#                     modified_kwargs[f"error_{name}_{try_i}"] = error
#                     if name == "general":
#                         error_prefix = "General:"
#                     else:
#                         error_prefix = signature.output_fields[name].json_schema_extra["prefix"]
#                     number = "" if try_i == 0 else f" ({try_i+1})"
#                     signature = signature.append(
#                         f"error_{name}_{try_i}",
#                         dspy.InputField(
#                             prefix=f"Past Error{number} in {error_prefix}",
#                             desc="An error to avoid in the future",
#                         ),
#                     )
#             else:
#                 # If there are no errors, we return the parsed results
#                 return Prediction.from_completions(
#                     {key: [r[key] for r in parsed_results] for key in signature.output_fields},
#                 )
#         raise ValueError(
#             "Too many retries trying to get the correct output format. " + "Try simplifying the requirements.",
#             errors,
#         )


# def _func_to_signature(func):
#     """Make a dspy.Signature based on a function definition."""
#     sig = inspect.signature(func)
#     annotations = typing.get_type_hints(func, include_extras=True)
#     output_key = func.__name__
#     instructions = func.__doc__
#     fields = {}

#     # Input fields
#     for param in sig.parameters.values():
#         if param.name == "self":
#             continue
#         # We default to str as the type of the input
#         annotation = annotations.get(param.name, str)
#         kwargs = {}
#         if typing.get_origin(annotation) is Annotated:
#             desc = next((arg for arg in typing.get_args(annotation) if isinstance(arg, str)), None)
#             if desc is not None:
#                 kwargs["desc"] = desc
#         fields[param.name] = (annotation, dspy.InputField(**kwargs))

#     # Output field
#     kwargs = {}
#     annotation = annotations.get("return", str)
#     if typing.get_origin(annotation) is Annotated:
#         desc = next((arg for arg in typing.get_args(annotation) if isinstance(arg, str)), None)
#         if desc is not None:
#             kwargs["desc"] = desc
#     fields[output_key] = (annotation, dspy.OutputField(**kwargs))

#     return dspy.Signature(fields, instructions)


# def _unwrap_json(output, from_json: Callable[[str], Union[pydantic.BaseModel, str, None]]):
#     try:
#         parsing_result = from_json(output)
#         if isinstance(parsing_result, pydantic.BaseModel):
#             return parsing_result.model_dump_json()
#         else:
#             return output
#     except (ValueError, pydantic.ValidationError, AttributeError):
#         output = output.strip()
#         if output.startswith("```"):
#             if not output.startswith("```json"):
#                 raise ValueError("json output should start with ```json") from None
#             if not output.endswith("```"):
#                 raise ValueError("Don't write anything after the final json ```") from None
#             output = output[7:-3].strip()
#         return ujson.dumps(ujson.loads(output))  # ujson is a bit more robust than the standard json