import dspy
import json
import os
from typing import List, Dict, Any, Tuple
from datetime import datetime


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 the request is to modify code, suggest changes using git diff syntax.

    - If you have enough information, your decision must be 'answer' and you must provide a long, complete, and well-explained answer.
    - If you need more information, your decision must be 'files' and you must provide a JSON-formatted list of file paths that are needed.
    - If the request is to modify code, your decision should be 'code_change', and you must provide both a high-level suggestion and code changes in git diff format.
    """

    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 string containing the content of retrieved files relevant to the question.")

    decision = dspy.OutputField(
        desc="Your decision: either 'answer', 'files', or 'code_change'."
    )

    answer = dspy.OutputField(
        desc=(
            "If decision is 'answer', provide a detailed, comprehensive, and well-structured response. "
            "If decision is 'files', return a JSON-formatted list of file paths required to answer the question. "
            "If decision is 'code_change', provide a high-level suggestion or explanation for the proposed changes. "
            "For code changes, this summary is optional and can be brief; the git diff is the priority."
        )
    )
    
    code_diff = dspy.OutputField(
        desc=(
            "If decision is 'code_change', provide the code modifications using a clear git diff syntax. "
            "This output must be complete and not truncated. "
            "For 'answer' or 'files' decisions, this field should be empty."
        )
    )

class CodeRagQuestionDecider(dspy.Module):
    """
    A pipeline that uses DSPy to decide whether to answer a user's question,
    request additional files, or suggest a code change based on the available data.
    """

    def __init__(self):
        super().__init__()

        # Modified system prompt to prioritize git diff output
        # Revised system prompt for CodeRagQuestionDecider
        system_prompt = """
        You are a highly specialized AI assistant for a software engineering workspace. Your task is to accurately and efficiently handle user requests based on provided file information.

        **Core Directives:**
        1.  **Analyze the User's Question First:** Read the user's question and identify the specific file(s) they are asking about.
        2.  **Filter Aggressively:** The provided `retrieved_data` is a JSON string that may contain a large number of files. **You must strictly ignore all irrelevant files.**
            * **Irrelevant files include, but are not limited to:**
                * Compiled Python files (`.pyc`).
                * Documentation files (`.md`).
                * Shell scripts (`.sh`).
                * Files in `__pycache__` directories.
        3.  **Prioritize Relevant Files:**
            * Find the exact file path mentioned in the question (e.g., `app.core.services.workspace.py`).
            * Identify any related test files (e.g., `test_app_core_services_workspace.py` or similar).
            * **Only process the content of these highly relevant files.** The content for other files is not provided and should not be discussed.
        4.  **Decide Your Action:**
            * If the user's question is "polish the X.py and its test code" and you have the code for `X.py` in your context, your decision is always **'answer'**. You have enough information.
            * If you do not have the content for the specifically requested file(s), your decision is **'files'**, and you must provide a list of the required file paths (e.g., `["/path/to/X.py", "/path/to/test_X.py"]`).
        5.  **Format the Output Strictly:**
            * If the decision is 'answer', the `answer` field must be a detailed response in markdown.
            * If the decision is 'files', the `answer` field must be a **valid JSON array of strings**, representing the file paths you need.
            * The `decision` field must be either 'answer' or 'files'.
            
        **Your goal is to provide a correct, concise, and focused response based ONLY on the provided relevant file content.** Do not hallucinate or discuss files for which you do not have content.
        """
        self.decider = dspy.Predict(QuestionDecider, system_prompt=system_prompt)

    async def forward(
        self,
        question: str,
        history: List[str],
        retrieved_data: Dict[str, Any]
    ) -> Tuple[str, str, str]:
        """
        Decide whether to answer, request more files, or suggest a code change.

        Args:
            question (str): The user's question.
            history (List[str]): The chat history.
            retrieved_data (Dict[str, Any]): The content of files relevant to the question.

        Returns:
            Tuple[str, str, str]: The model's answer, decision ('answer', 'files', or 'code_change'), and code diff.
        """
        history_text = "\n".join(history)
        retrieved_data_json = json.dumps(retrieved_data, indent=0)

        input_payload = {
            "question": question,
            "chat_history": history_text,
            "retrieved_data": retrieved_data_json
        }

        prediction = await self.decider.acall(**input_payload)

        self._log_to_file(input_payload, prediction)

        return prediction.answer, prediction.decision.lower(), prediction.code_diff

    def _log_to_file(self, request_payload: Dict[str, Any], prediction: Any) -> None:
        """
        Saves the input and output of the AI call to a JSON file.

        Args:
            request_payload (Dict[str, Any]): The input sent to the AI.
            prediction (Any): The AI's response.
        """
        log_dir = "ai_payloads"
        os.makedirs(log_dir, exist_ok=True)

        timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
        filename = os.path.join(log_dir, f"ai_payload_{timestamp}.json")

        log_data = {
            "timestamp": timestamp,
            "request_payload": request_payload,
            "response_payload": {
                "decision": prediction.decision,
                "answer": prediction.answer,
                "code_diff": prediction.code_diff
            }
        }

        with open(filename, "w", encoding="utf-8") as f:
            json.dump(log_data, f, indent=4)

        print(f"[LOG] AI payload and response saved to {filename}")