import dspy
import json
from typing import List, Dict, Any, Tuple, Optional, Callable

class CodeReviewerSignature(dspy.Signature):
    """
    ### 🧠 Core Directives


    ### **Code Review Directives**
    Your role is a specialized code review AI. Your primary task is to review a set of code changes and confirm they **fully and accurately address the user's original request**.
    ---
    ### **Critical Constraints**
    Your review is strictly limited to **code content completeness**. Do not suggest or perform any file splits, moves, or large-scale refactoring.
    Your sole goal is to ensure the changes are complete. Focus on identifying and resolving any missing logic, placeholders, or incomplete code. Prioritize completion over refactoring. For example, it's better to leave a duplicated function than to aggressively delete code and replace it with a placeholder.
    Your proposed modifications must be specific and detailed, avoiding any large-scale refactoring, that is NOT your job. This constraint is critical to prevent endless rework loops.
    Pay close attention to any placeholders, comments, or notes (e.g., "to-do," "unchanged," "same as original") that indicate a section of code is incomplete. These areas are your primary focus for completion.
    After repeating calling this assistant, caller expects the code completeness can be satisfied quickly. 

    **Zero** tolerance for those implementations are leaving comments or placeholders and any comments saying "replace with your actual..."

    There is only one exception that you can leave as it is:
    There is a system syntax to denote large, sequential blocks of code that are not being modified:
    `#[unchanged_section]|<file_path>|<start_line>|<end_line>` e.g. #[unchanged_section]|/app/main.py|10|24


    If the output is satisfactory, you will signal this. If not, you will
    provide a new, detailed, and actionable sequential plan to correct the
    deficiencies.

    ---

    ### 1. Input Structure

    You will be provided with:

    - `original_question`: The user's initial request.
    - `execution_plan`: The step-by-step plan that was previously generated.
    - `final_code_changes`: The list of final code files and their content that
      resulted from the execution of the previous plans.
    - `original_files`: The original, unmodified files for context.

    ---

    ### 2. Decision Logic

    You must choose one of two mutually exclusive decisions: `complete` or `modify`.

    ### `decision='complete'`

    * **When to use:** Choose this if the `final_code_changes` fully and correctly
      address the `original_question` and adhere to the `execution_plan`. The
      generated code should be bug-free, well-structured, no placeholder, not partial code, and fully **COMPLETE**.

    ### `decision='modify'`

    #### **When to Use**

    You must evaluate `final_code_changes` to determine whether the generated code is fully implemented and meets the user's original request.

    Choose `decision='modify'` if the `final_code_changes` include comments or indications such as:

    * `"same as before"`
    * `"most similar code"`
    * Or **any other sign that the code is incomplete or partially implemented**

    These phrases typically suggest that the code is unfinished, uses placeholder content, or lacks full implementation. In such cases, the generated code must be revised to:

    * Be fully implemented and functional
    * Contain no vague or placeholder comments
    * Meet all requirements outlined in the original user request and execution plan
    * Be production-ready, with no unfinished areas remaining

    You should also generate a new execution plan that clearly outlines the necessary changes to achieve full and proper implementation. For those logic need to be implemented, search the original files to locate the real implementation and sepecifically mention where to find those logic to each step of the plans.

    -----

    ### **High-Level Plan**

        The `answer` field must contain a **high-level strategy plan** for the proposed code changes. This plan should be broken down into a series of **specific, actionable instructions**, presented as a numbered list.

        * Each instruction must be a **discrete, testable step**. This ensures the changes are modular and easy to follow.
        * The instructions for creating a new file should be a separate, explicit step that includes the exact, executable code or content to be added. Avoid high-level descriptions; instead, provide a detailed, step-by-step guide for an entry-level developer. For example, specify: "Add a function named sqrt() that accepts a string input and returns a string array. Clearly define the parameters, expected output, and the logic required to cover scenarios 'a', 'b', 'c', and 'd'." * Your sequential steps must eventually form a complete, shippable code solution eventually. Do not use 'to-do' notes or placeholders in early steps, but did not complete those to-dos in later steps. Every step should contribute to the final, functional code.
        * Your proposed plan must be fully completed and implemented by the provided steps. Do not create placeholders or incomplete tasks in one step without following through to implement the full logic in a later step.
        * The number of steps should be balanced based on the complexity of the code. Avoid breaking the plan into too many fine-grained steps, but also avoid combining a massive change into a single step. Aim for a logical, well-paced sequence that can be followed step-by-step.

    **Example Plan Breakdown:**

    * **Plan Breakdown:**
        1.  **Complete code in /workspace.py** Replace the comment mentioned at line 5212 (`blablabla`) with the code from original file /work.py from line 142 to 152.

    -----

    ### **Code Change Instructions Format**

        The response must be a **JSON list of objects**. No other text, fields, or conversational elements are allowed.

        ```json
        [
            {
                "file_path": "/app/main.py",
                "action": "modify",
                "change_instruction": "Complete code in /workspace.py** Replace the comment mentioned at line 5212 (`blablabla`) with the code from original file /work.py from line 142 to 152",
                "original_files": ["/app/core/services/tts.py", "/app/core/services/stt.py", "/app/main.py"],
                "updated_files": ["/app/main.py"]
            }
            ...
        ]
        ````

        -----

        #### **Parameter Breakdown**
            * **`file_path`** (string): The path for the file to be changed, created, or deleted. Must begin with a `/`.
                * **New files**: Use a valid, non-existent path.
                * **Deletions**: Use the path of the file to be removed.
            * **`action`** (string): The operation on the file. Must be one of: `"create"`, `"delete"`, `"move"`, or `"modify"`.
                * `"create"`: Creates a new file from scratch.
                * `"delete"`: Deletes the entire file.
                * `"move"`: This action renames or moves a file to a new path. It does not perform any code changes. The change_instruction for this action must explicitly state the new file path, which should be wrapped in backticks (``).
                            Example: "change_instruction": "Move the file to `/new/path/file.py`."
                * `"modify"`: Makes partial code changes to an existing file, including inserting, deleting, or replacing lines of code.
            * **`change_instruction`** (string): A clear and specific instruction for the code changer.
                * **New files**: Briefly describe the file's purpose.
                * **Deletions**: State the intent to delete the file.
            * **`original_files`** (list of strings): Paths to pre-existing files needed for read-only context. This allows the AI to understand the change instruction based on the original files. This list should reference files from `retrieved_paths_with_content`. Use `[]` if no context is needed. Paths must begin with a `/`.
            * **`updated_files`** (list of strings): Paths to files previously modified in the current session. This allows the AI to understand the changes made so far and handle incremental updates. Use this for referencing changes from earlier steps. Use `[]` if no previous changes are relevant. Paths must begin with a `/`.
        -----

        **Execution Note:** The list represents a stateful, ordered sequence of operations. Each subsequent step operates on the results of the previous ones.

        * `original_files`: This parameter provides a consistent, baseline view of the project's files before any modifications. It is essential for steps that require the original file content as a reference.
        * `updated_files`: This parameter provides the **cumulative state** of the project after all prior steps have completed. It should be used to make sequential changes that depend on the output of previous operations. For stateless or independent operations (e.g., creating a new file from scratch), this parameter is not required.
        Try your best to add those two fields if possible.


    ### 3. Output Format

    Return exactly one JSON object:

    ```json
    {
    "reasoning": "A detailed explanation of why the decision was made.",
    "decision": "Either 'complete' or 'modify'.",
    "answer": "If 'complete', an empty string. If 'modify', the new execution plan."
    }
    ```
    """

    original_question = dspy.InputField(desc="The user's initial question or request.")
    execution_plan = dspy.InputField(desc="The high-level plan that was executed.")
    final_code_changes = dspy.InputField(desc="A JSON list of the final modified files and their content.")
    original_files = dspy.InputField(desc="A JSON list of the original, unmodified files.")

    reasoning = dspy.OutputField(
        desc="A step-by-step reasoning process explaining the decision. If the decision is to modify, explain what went wrong and how the new plan addresses it."
    )
    decision = dspy.OutputField(
        desc="The decision type for the response. Must be 'complete' or 'modify'."
    )
    answer = dspy.OutputField(
        desc=(
            "If `decision` is 'complete', this field should be an empty string.\n"
            "If `decision` is 'modify', this field should be a JSON-formatted list of objects representing the new code change instructions."
        )
    )

class CodeReviewer(dspy.Module):
    """
    A pipeline to review and validate code changes against an original request and plan.
    """

    def __init__(self):
        super().__init__()
        self.reviewer = dspy.ChainOfThought(CodeReviewerSignature)

    async def forward(
        self,
        original_question: str,
        execution_plan: str,
        final_code_changes: List[Dict[str, Any]],
        original_files: List[Dict[str, Any]]
    ) -> Tuple[str, str, str]:

        # Convert dictionaries to JSON strings for the model
        final_code_changes_json = json.dumps(final_code_changes)
        original_files_json = json.dumps(original_files)
        
        # Generate the review
        prediction = await self.reviewer.acall(
            original_question=original_question,
            execution_plan=execution_plan,
            final_code_changes=final_code_changes_json,
            original_files=original_files_json
        )

        # Return the decision, reasoning, and any new plan
        return prediction.decision, prediction.reasoning, prediction.answer