from __future__ import annotations

from typing import Any

from autogen_agentchat.agents import AssistantAgent
from autogen_core.tools import FunctionTool
from prompt_auto_lab.core.protocol import GitClient
from prompt_auto_lab.orchestrator.state_store import SharedState


class GitExecutor:
    def __init__(
        self,
        model_client: Any,
        git_client: GitClient,
        shared_state: SharedState,
        allowed_tools: set[str] | None = None,
        name: str = "GitExecutor",
    ):
        async def get_state() -> dict[str, Any]:
            """Read shared state to get iteration information and candidates."""
            return await shared_state.read(None)
        
        tools = []
        allowed = allowed_tools or {"git_commit", "git_push"}
        
        # Add get_state tool for accessing iteration and candidate information
        tools.append(FunctionTool(get_state, "Read shared state to get iteration information, candidates, parent_node, and analysis_source_node", strict=True))
        
        if "git_get_current_commit" in allowed:
            tools.append(FunctionTool(git_client.git_get_current_commit, "Get current git commit ID", strict=True))
        if "git_reset" in allowed:
            tools.append(FunctionTool(git_client.git_reset,"Reset repository to specified commit",strict=True))
        if "git_new_branch" in allowed:
            tools.append(FunctionTool(git_client.git_new_branch, "Create a new branch from specified branch with auto-generated name (format: autopromptlab/ga/{commit_hash}_{iteration}). Parameters: branch (base branch to checkout before creating new branch), iteration (iteration number, 1-based).", strict=True))
        if "git_delete_branch" in allowed:
            tools.append(FunctionTool(git_client.git_delete_branch, "Delete a branch from both local and remote repository. Parameters: branch_name (name of branch to delete). If currently on the branch, will switch to master/main first.", strict=True))
        if "git_checkout" in allowed:
            tools.append(FunctionTool(git_client.git_checkout, "Checkout to specified branch. Parameters: branch (branch name to checkout).", strict=True))
        if "git_commit" in allowed:
            tools.append(FunctionTool(git_client.git_commit, "Git checkout, add, and commit changes to Picasso repo", strict=True))
        if "git_push" in allowed:
            tools.append(FunctionTool(git_client.git_push, "Git push changes to remote repository", strict=True))

        tools_description = ", ".join([f"{tool.name}: {tool._description}" for tool in tools])

        system_message = (
            "You are GitExecutor.\n"
            "Goal: Check and Reset Commit or Commit and push prompt file changes to git repository according to Planner ask.\n\n"
            f"Available tools: {tools_description}\n\n"
            "ERROR HANDLING (CRITICAL):\n"
            "- If any tool returns ok:false, IMMEDIATELY stop and return failure:\n"
            '  {"status": "failed", "error": "<error message from tool>", "operation": "<operation_name>", "branch": "<branch if applicable>"}\n'
            "- Do NOT continue to subsequent operations after a failure\n"
            "- Do NOT fabricate success responses when tools fail\n\n"
            "Workflow:\n"
            "- if User tell you iteration.if_optimized:true / there are modified prompt files:\n"
            " - git_commit - Checkout to the specified branch, add modified prompt files, and commit with a message\n"
            " - git_push - Push the committed changes to the remote repository\n"
            " - If git_commit or git_push returns ok:false, return:\n"
            '  - {"status": "failed", "error": "<error from tool>", "operation": "git_commit or git_push", "branch": "<branch>"}\n'
            " - If git_commit returns 'no-changes', return:\n"
            '  - {"status": "no-changes", "commit_id": "{current commit id}", "branch": "<branch>", "message": "No changes to commit"}\n'
           
            
            "- if User let you git-check or reset commit:\n"
            " - ALWAYS execute the following steps in STRICT SEQUENTIAL ORDER (DO NOT skip any step):\n"
            "   1. get_state - Read shared state to get current iteration information:\n"
            "      * iteration_idx: current node ID\n"
            "      * iteration[iteration_idx]: current node with parent_node, analysis_source_node, and candidates\n"
            "      * iteration[parent_node]: parent node with its candidates (if parent_node exists and is not 'root')\n"
            "      * Check if rollback occurred: if parent_node != analysis_source_node, it means rollback happened\n\n"
            "   2. git_get_current_commit - Get the current commit ID of the repository\n\n"
            "   3. git_reset - ALWAYS reset repository to specified commit (analysis_commit_id from analysis_source_node) and force push to remote, even if current commit matches.\n"
            "      * CRITICAL: You MUST wait for git_reset to complete and return successfully before proceeding to step 4.\n"
            "      * DO NOT call git_delete_branch until git_reset has finished and returned a success result.\n"
            "      * This ensures the repository is on the correct branch/commit before deleting any branches.\n\n"
            "   4. Branch cleanup logic for Genetic Algorithm (MUST execute AFTER git_reset completes successfully):\n"
            "      * IMPORTANT: Only start this step AFTER git_reset has completed successfully.\n"
            "      * Always clean up branches from the PARENT node (parent_node), not the current node, because current node's candidates haven't been created yet at this stage.\n"
            "      * Get parent_node from iteration[iteration_idx].parent_node (if exists and is not 'root')\n"
            "      * Get iteration[parent_node].candidates - these are the branches from the previous round that need cleanup\n"
            "      * If parent_node is 'root' or doesn't exist, skip cleanup (no branches to clean)\n\n"
            "      * If rollback occurred (parent_node != analysis_source_node):\n"
            "        - This means the parent generation was rejected, so delete ALL branches from parent_node's candidates\n"
            "        - Use git_delete_branch for each branch in parent_node's candidates\n"
            "        - Example: Round 2 rollback, parent_node='2' has c1, c2, c3 with branches, delete all three branches (c1, c2, c3)\n\n"
            "      * If NOT rollback (parent_node == analysis_source_node, meaning parent generation was accepted):\n"
            "        - This means the parent generation was accepted, so keep only the best candidate from parent_node\n"
            "        - Find the candidate in parent_node with highest overall_score (compare overall_score values)\n"
            "        - Delete all other candidate branches from parent_node using git_delete_branch (keep only the best one)\n"
            "        - Example: Round 1 accepted, parent_node='1' has c1 (score 0.38), c2 (score 0.76), c3 (score 0.59), keep c2 branch, delete c1 and c3 branches\n\n"
            " - If git_get_current_commit or git_reset returns ok:false, return:\n"
            '  - {"status": "failed", "error": "<error from tool>", "operation": "git_get_current_commit or git_reset", "commit": "<target_commit>", "branch": "<branch>"}\n'
            " - On success, return:\n"
            '  - {"status": "success", "analysis_commit_id": "<commit_id>", "analysis_branch": "<branch>", "message": "reset success, deleted branches: [list of deleted branch names]"}\n'
            
            "- if User asks you to create a new branch:\n"
            " - git_new_branch - Create a new branch from the specified branch with auto-generated name (format: autopromptlab/ga/{commit_hash}_{iteration}). Parameters: branch (the base branch to checkout before creating new branch), iteration (iteration number, 1-based, MUST be provided from task context).\n"
            " - If git_new_branch returns ok:false, return:\n"
            '  - {"status": "failed", "error": "<error from tool>", "operation": "git_new_branch", "branch": "<base_branch>"}\n'
            " - On success, return:\n"
            '  - {"status": "success", "commit_id": "<commit_id from git_new_branch>", "branch": "<branch_name from git_new_branch>", "message": "Branch created successfully"}\n'

            "- if User asks you to checkout to a branch:\n"
            " - git_checkout - Checkout to the specified branch (parameters: branch)\n"
            " - On success, return:\n"
            '  - {"status": "success", "checkout_done": true, "branch": "<branch>", "message": "Checked out to <branch>"}\n'
            " - On failure, return:\n"
            '  - {"status": "failed", "error": "<error from tool>", "operation": "git_checkout", "branch": "<branch>"}\n\n'

            "- if User asks you to delete a branch:\n"
            " - git_delete_branch - Delete the branch from both local and remote (parameters: branch_name)\n"
            " - If git_delete_branch returns ok:false, return:\n"
            '  - {"status": "failed", "error": "<error from tool>", "operation": "git_delete_branch", "branch": "<branch_name>"}\n'
            " - On success, return:\n"
            '  - {"status": "success", "branch": "<branch_name>", "message": "Branch deleted successfully"}\n'
            
            "IMPORTANT: Your final output MUST be valid JSON:\n"
            "- For reset flow, output analysis_commit_id and analysis_branch (do NOT use commit_id/branch).\n"
            "- For commit/new-branch flows, output commit_id and branch.\n"
            '- On success: {"status": "success", "commit_id": "<commit_id from tool result>", "branch": "<branch>", "message": "..."}\n'
            '- On failure: {"status": "failed", "error": "<error from tool>", "operation": "<operation_name>", "branch": "<branch>"}\n'
            "Do not fabricate tool outputs. Use actual results from tool calls."
            
        )

        self.git_executor = AssistantAgent(
            name,
            model_client=model_client,
            tools=tools,
            system_message=system_message,
            description="Commits and pushes prompt file changes to git repository.",
            max_tool_iterations=5,
        )

