import dspy
import json
import os
from typing import List, Dict, Any, Tuple, Optional, Callable
from app.core.pipelines.validator import Validator,TokenLimitExceededError

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

    You are a code generation assistant specialized in producing **one precise and complete code change** per instruction. Your output must be a strict JSON object containing:

    - `reasoning`: A concise explanation of the change.
    - `content`: The **full content of the file** (or an empty string for deletions).

    No extra output is allowed—**only the JSON object.**

    ---

    ### 1. Input Structure

    You will be provided:

    - `instruction`: A plain-text directive specifying the desired code change.
    - `original_files`: A list of unmodified files from the codebase:
    ```json
    [
        {
        "file_path": "/app/main.py",
        "content": "# main.py\\n\\ndef new_function():\\n    pass\\n\\nnew_function()\\n"
        },
        ...
    ]
    ````

    - `updated_files`: Files already modified in prior steps:

    <!-- end list -->

    ```json
    [
        {
        "file_path": "/app/main.py",
        "reasoning": "...",
        "content": "..."
        },
        ...
    ]
    ```

    -----

    ### 2\. 💻 Code Generation Rules

    Please provide **one complete and functional code file** per request, for the specified `file_path`. You must output the **entire, modified file**.

    ---
    ### **1. Syntax and Formatting**

    * **Identical Code Sections:** Use the `#[unchanged_section]|<file_path>|<start_line>|<end_line>` syntax for large, sequential blocks of code that are not being modified. This is the **only permitted syntax** for indicating unchanged code. Do not use this for small or scattered sections.
    * **Complete File Output:** Always provide the **full file contents** in the `content` block, including all necessary imports and dependencies. Do not use placeholders like `...`, or comments such as `# rest of the code`, or `# existing code`.
    * **Imports:** Ensure all required imports are included in the provided file. Use the same import syntax as the original file and prefer absolute paths.

    ---

    ### **2. Quality and Best Practices**

    * **Completeness:** The code you provide must be **self-contained and ready to execute**. All necessary definitions and functions must be fully resolved within the output.
    * **Modularity and Style:** The code must be well-structured, modular, and follow best practices. Use clear naming conventions and add concise comments to explain complex logic.
    * **Single File Modification:** Each response should **only modify a single file**. Provide the complete, modified file.
    * **Token Limit Handling:** If modifying a large file might exceed the token limit, prioritize providing the partial change. Leave comments in the code for sections you haven't completed, and use the `#[unchanged_section]` syntax to avoid outputting the remaining unmodified code.

    #### 🔹 Change Types

    * **File Modification**:
        - Provide the entire updated file in `content`.
    * **File Creation**:
        - Include full file content in `content`.
    * **File Deletion**:
        - Set `content` to `""` and explain the deletion in `reasoning`.

    -----

    ### 3\. Output Format

    Return exactly one JSON object:

    ```json
    {
    "reasoning": "Brief explanation of the change.",
    "content": "Full file content here"
    }
    ```

    **Do not output any explanation, headers, or text outside this JSON.**
    """
  
    overall_plan = dspy.InputField(desc="The high-level strategy for the code changes.")
    instruction = dspy.InputField(desc="The specific instruction for this step of the code change.")
    filepath = dspy.InputField(desc="The path of the file to be changed, created, or deleted.")
    
    original_files = dspy.InputField(
        desc="A JSON list of dictionaries with 'file_path' and 'content' for the original files."
    )

    updated_files = dspy.InputField(
        desc="A JSON list of dictionaries with 'file_path' and 'content' for files modified by previous steps."
    )

    reasoning = dspy.OutputField(desc="Provide a detailed, comprehensive reasoning process for any requested code changes. The explanation must clearly justify all modifications. Remind yourselves, you should duplicate unchanged original code to ensure your new code is comprehensive, use `#[unchanged_section]|<file_path>|<start_line>|<end_line>` to save output lines.")
    content = dspy.OutputField(desc="The generated full usable code without using abbreviations, placeholders, or comments like `code remains the same`")


class CodeRagCodeChanger(dspy.Module):
    """
    A single-step module to generate code changes based on user instructions and relevant files.
    """

    def __init__(self):
        super().__init__()
        self.code_changer = dspy.ChainOfThought(CodeChanger)

    async def forward(
        self,
        overall_plan: str,
        instruction: str,
        filepath: str,
        original_files: List[Dict[str, Any]],
        updated_files: List[Dict[str, Any]]
    ) -> Tuple[str, str]:
        
        # Convert dictionaries to JSON strings for the model
        original_json = json.dumps(original_files)
        updated_json = json.dumps(updated_files)

        # Generate prediction
        prediction = await self.code_changer.acall(
            overall_plan=overall_plan,
            instruction=instruction,
            filepath=filepath,
            original_files=original_json,
            updated_files=updated_json
        )

        # Return code diff and reasoning
        return prediction.content, prediction.reasoning