import logging
from typing import TypeVar

from dspy import InputField, OutputField, Predict, Signature
from pydantic import BaseModel, ValidationError

from dslmodel.template import render

logger = logging.getLogger(__name__)
logger.setLevel(logging.ERROR)


def eval_dict_str(dict_str: str) -> dict:
    """Safely convert str to dict"""
    return ast.literal_eval(dict_str)


class PromptToPydanticInstanceSignature(Signature):
    """Synthesize the prompt into the kwargs to fit the model.
    Do not duplicate the field descriptions
    """

    root_pydantic_model_class_name = InputField(
        desc="The class name of the pydantic model to receive the kwargs"
    )
    pydantic_model_definitions = InputField(desc="Pydantic model class definitions as a string")
    prompt = InputField(
        desc="The prompt to be synthesized into data. Do not duplicate descriptions"
    )
    root_model_kwargs_dict = OutputField(
        prefix="```python\nkwargs_dict: dict = ",
        desc="Generate a Python dictionary. IT WILL HAVE TO VALIDATE WITHIN PYDANTIC",
    )


class PromptToPydanticInstanceErrorSignature(Signature):
    """Synthesize the prompt into the kwargs fit the model"""

    error = InputField(desc="Error message to fix the kwargs")
    generated_kwargs = InputField(desc="The kwargs generated by the LLM")
    root_pydantic_model_class_name = InputField(
        desc="The class name of the pydantic model to receive the kwargs"
    )
    pydantic_model_definitions = InputField(desc="Pydantic model class definitions as a string")
    prompt = InputField(desc="The prompt to be synthesized into data")
    what_went_wrong = OutputField(desc="What went wrong in the kwargs generation")
    how_to_fix = OutputField(desc="How to fix the kwargs generation")
    root_model_kwargs_dict = OutputField(
        prefix="```python\nkwargs_dict: dict = ",
        desc="Generate a Python dictionary. IT WILL HAVE TO VALIDATE WITHIN PYDANTIC",
    )


class DiagnosisSignature(Signature):
    """Diagnose why the LLM couldn't create the dictionary and suggest potential changes to the source."""

    error_message = InputField(
        desc="The error message indicating why the dictionary creation failed."
    )
    root_pydantic_model_class_name = InputField(
        desc="The class name of the Pydantic model that was expected."
    )
    prompt = InputField(desc="The original prompt that was used to generate the dictionary.")
    suggested_changes = OutputField(
        prefix="```python\nsuggested_changes: List[str] = ",
        desc="A list of suggested changes to the source code or prompt to improve dictionary generation.",
    )


T = TypeVar("T", bound=BaseModel)


import ast
import logging
from typing import TypeVar

import dspy
from dspy import Assert, ChainOfThought, InputField, OutputField, Signature
from pydantic import BaseModel

from dslmodel.utils.source_tools import collect_all_sources_as_string

logger = logging.getLogger(__name__)
logger.setLevel(logging.ERROR)


# Helper function to evaluate dictionary from string
def eval_dict_str(dict_str: str) -> dict:
    """Safely convert str to dict"""
    return ast.literal_eval(dict_str)


# Signatures and other components remain unchanged

T = TypeVar("T", bound=BaseModel)


class GenPydanticInstance(dspy.Module):
    """A module for generating and validating Pydantic model instances based on prompts."""

    def __init__(
        self,
        model: type[T],
        generate_sig=PromptToPydanticInstanceSignature,
        correct_generate_sig=PromptToPydanticInstanceErrorSignature,
        diagnosis_sig=DiagnosisSignature,
        verbose=False,
    ):
        super().__init__()
        self.output_key = "root_model_kwargs_dict"
        self.model = model
        self.verbose = verbose

        # Collect source code for model validation and correction logic
        self.model_sources = collect_all_sources_as_string(model)

        # Initialize DSPy ChainOfThought dspy_modules for generation, correction, and diagnosis
        self.generate = Predict(generate_sig)
        self.correct_generate = ChainOfThought(correct_generate_sig)
        self.diagnosis_generate = ChainOfThought(diagnosis_sig)
        self.validation_error = None

    def generate_output(self, prompt: str) -> str:
        """Generates output from the prompt."""
        output = self.generate(
            prompt=prompt,
            root_pydantic_model_class_name=self.model.__name__,
            pydantic_model_definitions=self.model_sources,
        )
        return output[self.output_key]

    def validate_root_model(self, output: str) -> bool:
        """Validates the generated output against the root Pydantic model."""
        try:
            model_inst = self.model.model_validate(eval_dict_str(output))
            return isinstance(model_inst, self.model)
        except (ValidationError, ValueError, TypeError, SyntaxError) as error:
            self.validation_error = error
            logger.debug(f"Validation error: {error}")
            return False

    def validate_output(self, output: str) -> T:
        """Validates the generated output and returns the model instance if successful."""
        Assert(
            self.validate_root_model(output),
            f"You need to create a kwargs dict for {self.model.__name__}\nValidation error:\n{self.validation_error}",
        )
        return self.model.model_validate(eval_dict_str(output))

    def handle_correction(self, prompt: str, output: str) -> T:
        """Attempts to correct the generated output based on errors."""
        try:
            corrected_output = self.correct_generate(
                prompt=prompt,
                root_pydantic_model_class_name=self.model.__name__,
                pydantic_model_definitions=self.model_sources,
                generated_kwargs=output,
                error=f"Error: {self.validation_error}",
            )[self.output_key]
            return self.validate_output(corrected_output)
        except (AssertionError, ValueError, TypeError) as error:
            logger.error(f"Correction failed: {error}")
            raise

    def diagnose_issue(self, prompt: str) -> None:
        """Diagnoses the error when both generation and correction steps fail."""
        diagnosis_output = self.diagnosis_generate(
            error_message=str(self.validation_error),
            root_pydantic_model_class_name=self.model.__name__,
            prompt=prompt,
        )
        suggested_changes = diagnosis_output["suggested_changes"]
        logger.error(f"Diagnosis suggestions: {suggested_changes}")
        raise ValueError(f"Model generation failed. Suggested changes: {suggested_changes}")

    def forward(self, prompt: str) -> T:
        """The main function that handles generation, validation, correction, and diagnosis."""
        # Step 1: Generate initial output
        prompt = render(prompt)

        output = self.generate_output(prompt)

        # Step 2: Attempt to validate the generated output
        try:
            return self.validate_output(output)
        except (AssertionError, ValueError, TypeError) as error:
            logger.error(f"Error during validation: {error}\nOutput:\n{output}")

            # Step 3: Try to correct the output
            try:
                return self.handle_correction(prompt, output)
            except (AssertionError, ValueError, TypeError) as error2:
                # Step 4: If correction fails, perform diagnosis
                print(self.diagnose_issue(prompt))
                raise error2

    def __call__(self, prompt: str):
        return self.forward(prompt)


def gen_instance(model, prompt, verbose=False):
    """Helper function to instantiate and use GenPydanticInstance."""
    model_module = GenPydanticInstance(model, verbose=verbose)
    return model_module(prompt)


def main():
    from dslmodel import init_instant

    init_instant()
    # Example usage would go here


if __name__ == "__main__":
    main()
