import dspy
import json
from typing import List, Dict, Any

class QuestionDecider(dspy.Signature):
    """
    Based on the user's question, chat history, and the content of the retrieved files, decide whether you have enough information to answer the question or if you need to request more files.
    If you have enough information, your decision should be 'answer' and you should provide a detailed answer.
    If you need more information, your decision should be 'files' and you should provide a list of file paths that are needed.
    """
    question = dspy.InputField(desc="The user's current question.")
    chat_history = dspy.InputField(desc="The ongoing dialogue between the user and the AI.")
    retrieved_data = dspy.InputField(desc="A JSON object containing the content of the retrieved files relevant to the question.")
    decision = dspy.OutputField(
        desc="Your decision, which must be either 'answer' or 'files'."
    )
    answer = dspy.OutputField(
        desc="If the decision is 'answer', this is your final, complete answer. If the decision is 'files', this is a JSON-formatted list of file paths needed to answer the question."
    )
    
class CodeRagQuestionDecider(dspy.Module):
    """
    A pipeline to decide whether to answer a question or request more files.
    """
    def __init__(self):
        super().__init__()
        # Assign the system prompt directly to the dspy.Predict instance
        system_prompt = "You are a helpful AI assistant. Your job is to decide whether to answer the user's question based on the provided context or to request more files. If you have enough information, respond with 'answer' and provide a detailed answer. If not, respond with 'files' and provide a JSON-formatted list of file paths needed."
        self.decider = dspy.Predict(QuestionDecider, system_prompt=system_prompt)

    async def forward(self, question: str, history: List[str], retrieved_data: Dict[str, Any]) -> str | List[str]:
        # Format the chat history for the signature
        history_text = "\n".join(history)
        
        # Convert the retrieved_data to a JSON string
        retrieved_data_json_string = json.dumps(retrieved_data, indent=2)

        # Use the decider to determine the next action
        prediction = await self.decider.acall(
            question=question,
            chat_history=history_text,
            retrieved_data=retrieved_data_json_string
        )
        
        # Check the decision and return the appropriate output
        if prediction.decision.lower() == 'answer':
            # The AI has enough information, so return the final answer.
            return prediction.answer
        elif prediction.decision.lower() == 'files':
            # The AI needs more files, so parse and return the list of file paths.
            try:
                # The LLM's output for files should be a JSON array string
                return json.loads(prediction.answer)
            except (json.JSONDecodeError, TypeError) as e:
                # Fallback for malformed JSON output from the LLM
                print(f"Warning: Failed to parse files list from LLM: {e}")
                print(f"Raw LLM output: {prediction.answer}")
                # Return an empty list or a list parsed from lines, depending on how you want to handle bad output.
                # Assuming the model returns a string like 'file1.py,file2.py' or '["file1.py", "file2.py"]'
                if isinstance(prediction.answer, str):
                    cleaned_string = prediction.answer.strip("[]'\" \n\t")
                    return [item.strip() for item in cleaned_string.split(',') if item.strip()]
                return []
        else:
            # Handle unexpected decisions from the LLM
            print(f"Warning: Unexpected decision from LLM: {prediction.decision}")
            return []