"""
Agent builder tools for creating and saving new agents.

These tools are used by the agent when it needs to create reusable agent scripts:
- plan_agent: Creates or updates an agent plan
- code_and_save_agent: Generates and saves agent code from a plan
- create_or_update_plan: Implementation function for plan creation/update
- build_or_patch_code: Implementation function for code generation/patching
"""

import ast
import asyncio
import os
import uuid
from typing import TYPE_CHECKING, Any, cast

from langchain_core.messages import AIMessage, HumanMessage
from langchain_core.tools import tool as langchain_tool
from langchain_google_genai import ChatGoogleGenerativeAI

from universal_mcp.agents.codeact0.prompts import (
    AGENT_BUILDER_CODE_PATCH_PROMPT,
    AGENT_BUILDER_GENERATING_PROMPT,
    AGENT_BUILDER_HTML_FORM_PROMPT,
    AGENT_BUILDER_META_PROMPT,
    AGENT_BUILDER_PLAN_PATCH_PROMPT,
    AGENT_BUILDER_PLANNING_PROMPT,
)
from universal_mcp.agents.codeact0.state import (
    AgentBuilderCode,
    AgentBuilderMeta,
    AgentBuilderPatch,
    AgentBuilderPlan,
    CodeActState,
)
from universal_mcp.agents.codeact0.utils import (
    apply_patch_or_use_proposed,
    extract_code_tools,
    sanitize_messages,
)
from universal_mcp.agents.llm import load_chat_model
from universal_mcp.agents.utils import convert_tool_ids_to_dict

if TYPE_CHECKING:
    from universal_mcp.agents.codeact0.agent import CodeActPlaybookAgent


def create_agent_builder_tools() -> dict[str, Any]:
    """Create tools for agent plan and code creation, saving, modifying"""

    @langchain_tool
    async def plan_agent():
        """
        Creates a new agent plan if none exists, or update the existing plan using a minimal patch.
        Non-interactive; relies only on conversation/code history. Missing details must be external variables.
        No arguments are required; all context is implicit.
        """
        return "ok"

    @langchain_tool
    async def schedule_code_and_save_agent(
        cron: str | None = None,
        update_code: bool = True,
        update_general_info: bool = True,
        generate_html: bool = False,
    ):
        """
        Generates/Updates Python code and metadata for an agent from a confirmed plan, AND/OR schedules it to run automatically.
        Non-interactive; main agent must ensure capabilities are loaded beforehand.

        Args:
            cron: Cron expression (e.g. "0 9 * * *") for the agent to run. If provided, updates/sets the schedule. This sets the agent to run automatically for the user.
            update_code: Whether to regenerate/patch the Python code. Set to False if only updating schedule/metadata/HTML form.
            update_general_info: Whether to regenerate/patch name and description. Set to False if only updating schedule/code/HTML form.
            generate_html: Whether to generate an HTML form for the agent. Set to True to generate a form, ONLY when asked for by the user. Note that you can call this tool with just this argument as True if other details need not be updated.
        """
        return "ok"

    return {
        "plan_agent": plan_agent,
        "schedule_code_and_save_agent": schedule_code_and_save_agent,
    }


async def create_or_update_plan(self: "CodeActPlaybookAgent", state: CodeActState, plan: list[str] | None):
    """Sub-agent helper: create or patch-update the agent plan and emit UI updates.
    Returns: (plan: list[str], tool_result: str)
    """
    str(uuid.uuid4())

    # Determine existing plan (prefer persisted agent's plan) and base messages
    existing_plan_steps = plan
    base = sanitize_messages(state["messages"])

    def with_sys(text: str):
        return [{"role": "system", "content": text}] + base

    current_code = (
        self.agent.instructions.get("script") if self.agent and getattr(self.agent, "instructions", None) else None
    )

    if existing_plan_steps:
        current = "\n".join(map(str, existing_plan_steps or []))
        sys_prompt = self.instructions + "\n" + AGENT_BUILDER_PLAN_PATCH_PROMPT + "\n".join(self.default_tool_defs)
        msgs = with_sys(sys_prompt) + [
            HumanMessage(content=f"Current plan (one step per line):\n{current} \n Current code:\n{current_code}")
        ]
        if isinstance(self.model_instance, ChatGoogleGenerativeAI):
            patch_model = self.agent_builder_model_instance.with_structured_output(
                AgentBuilderPatch.model_json_schema(), method="json_schema"
            )
            proposed = cast(AgentBuilderPatch, await patch_model.with_config({"tags": ["nostream"]}).ainvoke(msgs)).get(
                "patch"
            )
        else:
            patch_model = self.agent_builder_model_instance.with_structured_output(AgentBuilderPatch)
            proposed = cast(
                AgentBuilderPatch, await patch_model.with_config({"tags": ["nostream"]}).ainvoke(msgs)
            ).patch
        updated = apply_patch_or_use_proposed(current, proposed)
        plan = [line for line in updated.splitlines() if line.strip()]
    else:
        sys_prompt = self.instructions + AGENT_BUILDER_PLANNING_PROMPT + "\n".join(self.default_tool_defs)
        if isinstance(self.model_instance, ChatGoogleGenerativeAI):
            plan_model = self.agent_builder_model_instance.with_structured_output(
                AgentBuilderPlan.model_json_schema(), method="json_schema"
            )
            plan = cast(dict, await plan_model.with_config({"tags": ["nostream"]}).ainvoke(with_sys(sys_prompt))).get(
                "steps"
            )
        else:
            plan_model = self.agent_builder_model_instance.with_structured_output(AgentBuilderPlan)
            plan = cast(
                AgentBuilderPlan, await plan_model.with_config({"tags": ["nostream"]}).ainvoke(with_sys(sys_prompt))
            ).steps

    tool_result = {
        "plan": plan,
        "update": bool(plan),
        "message": "Successfully generated the agent plan. However, you must call schedule_code_and_save_agent to save the plan (as well as the code, regardless of any changes in the code)",
    }
    return plan, tool_result


async def _generate_meta(
    self: "CodeActPlaybookAgent",
    state: "CodeActState",
    update_general_info: bool,
    base_messages: list,
) -> tuple[str | None, str | None, str | None]:
    agent_name = state.get("agent_name")
    agent_description = state.get("agent_description")
    long_description = None

    current_name = None
    current_description = None
    current_long_description = None

    if self.agent:
        current_name = getattr(self.agent, "name", None)
        current_description = getattr(self.agent, "description", None)
        if self.agent.instructions:
            current_long_description = self.agent.instructions.get("display")

    if not update_general_info:
        if not agent_name and current_name:
            agent_name = current_name
        if not agent_description and current_description:
            agent_description = current_description
        long_description = current_long_description
        return agent_name, agent_description, long_description

    meta_prompt = self.instructions + AGENT_BUILDER_META_PROMPT
    if current_name:
        meta_prompt += f"\nExisting Name: {current_name}"
    if current_description:
        meta_prompt += f"\nExisting Description: {current_description}"
    if current_long_description:
        meta_prompt += f"\nExisting Long Description: {current_long_description}"

    messages = [{"role": "system", "content": meta_prompt}] + base_messages

    if isinstance(self.model_instance, ChatGoogleGenerativeAI):
        meta_model = self.agent_builder_model_instance.with_structured_output(
            AgentBuilderMeta.model_json_schema(), method="json_schema"
        )
        meta = cast(
            dict,
            await meta_model.with_config({"tags": ["nostream"]}).ainvoke(messages),
        )
        agent_name, agent_description, long_description = (
            meta.get("name"),
            meta.get("description"),
            meta.get("long_description"),
        )
    else:
        meta_model = self.agent_builder_model_instance.with_structured_output(AgentBuilderMeta)
        meta = cast(
            AgentBuilderMeta,
            await meta_model.with_config({"tags": ["nostream"]}).ainvoke(messages),
        )
        agent_name, agent_description, long_description = meta.name, meta.description, meta.long_description

    return agent_name, agent_description, long_description


async def _generate_code(
    self: "CodeActPlaybookAgent",
    update_code: bool,
    existing_code: str | None,
    plan_text: str | None,
    base_messages: list,
) -> dict[str, Any]:
    if not update_code:
        return {"code": existing_code, "messages": [], "is_patch": False, "skipped": True}

    if existing_code:
        generating_instructions = (
            self.instructions + AGENT_BUILDER_CODE_PATCH_PROMPT + "\n".join(self.default_tool_defs)
        )
        messages = [{"role": "system", "content": generating_instructions}] + base_messages
        if plan_text:
            messages.append(HumanMessage(content=f"Confirmed plan (one step per line):\n{plan_text}"))
        messages.append(HumanMessage(content=f"Current code to update:\n```python\n{existing_code}\n```"))

        patch_returns_dict = False
        if isinstance(self.model_instance, ChatGoogleGenerativeAI):
            patch_returns_dict = True
            patch_model = self.agent_builder_model_instance.with_structured_output(
                AgentBuilderPatch.model_json_schema(), method="json_schema"
            )
        else:
            patch_model = self.agent_builder_model_instance.with_structured_output(AgentBuilderPatch)

        patch_response = await patch_model.with_config({"tags": ["nostream"]}).ainvoke(messages)
        proposed = (
            cast(dict, patch_response).get("patch")
            if patch_returns_dict
            else cast(AgentBuilderPatch, patch_response).patch
        )
        return {
            "code": apply_patch_or_use_proposed(existing_code, proposed),
            "messages": messages,
            "proposed": proposed,
            "patch_model": patch_model,
            "is_patch": True,
        }
    else:
        messages = [
            {
                "role": "system",
                "content": self.instructions + AGENT_BUILDER_GENERATING_PROMPT + "\n".join(self.default_tool_defs),
            }
        ] + base_messages
        if isinstance(self.model_instance, ChatGoogleGenerativeAI):
            code_returns_dict = True
            code_model = self.agent_builder_model_instance.with_structured_output(
                AgentBuilderCode.model_json_schema(), method="json_schema"
            )
        else:
            code_returns_dict = False
            code_model = self.agent_builder_model_instance.with_structured_output(AgentBuilderCode)

        code_response = await code_model.with_config({"tags": ["nostream"]}).ainvoke(messages)
        python_code = (
            cast(dict, code_response).get("code") if code_returns_dict else cast(AgentBuilderCode, code_response).code
        )
        return {"code": python_code, "messages": messages, "is_patch": False}


async def _execute_and_retry_code(
    self: "CodeActPlaybookAgent",
    code_result: dict[str, Any],
    existing_code: str | None,
    state: "CodeActState",
) -> tuple[str | None, dict, str | None]:
    python_code = code_result["code"]
    messages = code_result["messages"]
    is_patch = code_result.get("is_patch", False)
    proposed = code_result.get("proposed")
    code_skipped = code_result.get("skipped", False)
    new_context = {}
    execution_error = None

    # Constants
    RETRY_MODEL_NAME = "anthropic:claude-opus-4-5-20251101"
    retry_model_instance = None

    if not code_skipped and python_code:
        retries = 0
        while retries < 3:
            output, new_context = await self.sandbox.handle_execute_python_code(python_code)
            if "Error" not in output:
                execution_error = None
                break
            else:
                execution_error = output
                retries += 1
                if retries >= 3:
                    break

                # Switch to fallback model for retries
                if not retry_model_instance:
                    retry_model_instance = load_chat_model(
                        RETRY_MODEL_NAME,
                        temperature=0.3,
                        thinking=False,
                        disable_streaming=True,
                        tags=("quiet",),
                    )

                if existing_code and is_patch:  # Logic for retry patch
                    messages.extend(
                        [
                            AIMessage(content=f"Generated patch:{proposed}"),
                            HumanMessage(
                                content=f"The following error was encountered in the generated patch: {output}\nRegenerate a patch for the same existing code."
                            ),
                        ]
                    )

                    # Determine model and method
                    current_model = retry_model_instance
                    patch_returns_dict = isinstance(current_model, ChatGoogleGenerativeAI)

                    if patch_returns_dict:
                        patch_model = current_model.with_structured_output(
                            AgentBuilderPatch.model_json_schema(), method="json_schema"
                        )
                    else:
                        patch_model = current_model.with_structured_output(AgentBuilderPatch)

                    patch_response = await patch_model.with_config({"tags": ["nostream"]}).ainvoke(messages)
                    proposed = (
                        cast(dict, patch_response).get("patch")
                        if patch_returns_dict
                        else cast(AgentBuilderPatch, patch_response).patch
                    )
                    python_code = apply_patch_or_use_proposed(existing_code, proposed)

                else:  # Logic for retry code
                    messages.extend(
                        [
                            AIMessage(content=f"Generated code:{python_code}"),
                            HumanMessage(
                                content=f"The following error was encountered in the generated code: {output}\nRewrite the entire script without the error."
                            ),
                        ]
                    )

                    # Determine model and method
                    current_model = retry_model_instance
                    code_returns_dict = isinstance(current_model, ChatGoogleGenerativeAI)

                    if code_returns_dict:
                        code_model = current_model.with_structured_output(
                            AgentBuilderCode.model_json_schema(), method="json_schema"
                        )
                    else:
                        code_model = current_model.with_structured_output(AgentBuilderCode)

                    code_response = await code_model.with_config({"tags": ["nostream"]}).ainvoke(messages)
                    python_code = (
                        cast(dict, code_response).get("code")
                        if code_returns_dict
                        else cast(AgentBuilderCode, code_response).code
                    )

    return python_code, new_context, execution_error


async def _generate_html_form(
    self: "CodeActPlaybookAgent",
    python_code: str | None,
    plan_text: str | None,
) -> str | None:
    if not python_code:
        return None

    html_form = None
    try:
        # Attempt 1: Try with main function and plan only
        main_func_source = None
        tree = ast.parse(python_code)
        for node in tree.body:
            if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and not node.name.startswith("_"):
                main_func_source = ast.get_source_segment(python_code, node)
                break

        if not main_func_source:
            # Fallback to entire code if main function not found
            main_func_source = python_code

        frontend_url = os.getenv("FRONTEND_URL", "https://ruzo.ai")
        # We use .replace() here because .format() would require escaping all other {} which are used heavily in the prompt for JSON/CSS
        prompt_with_url = AGENT_BUILDER_HTML_FORM_PROMPT.replace("{FRONTEND_URL}", frontend_url)

        html_prompt = prompt_with_url + f"\n\nPlan:\n{plan_text}\n\nFunction Definition:\n{main_func_source}"

        html_response = await self.agent_builder_model_instance.with_config({"tags": ["nostream"]}).ainvoke(
            [{"role": "user", "content": html_prompt}]
        )
        html_form = html_response.content
    except Exception:
        pass

    # Cleanup markdown if present
    if html_form:
        if html_form.startswith("```html"):
            html_form = html_form.replace("```html", "").replace("```", "").strip()
        elif html_form.startswith("```"):
            html_form = html_form.replace("```", "").strip()

    return html_form


async def _save_agent(
    self: "CodeActPlaybookAgent",
    agent_name: str | None,
    agent_description: str | None,
    long_description: str | None,
    python_code: str | None,
    plan: list[str] | None,
    cron: str | None,
    html_form: str | None,
    code_skipped: bool,
) -> dict | str:
    try:
        if not self.agent_builder_registry:
            raise ValueError("AgentBuilder registry is not configured")

        instructions_payload = {
            "plan": plan,
            "script": python_code,
            "display": long_description,
        }
        extracted_tools = extract_code_tools(python_code) if python_code else []
        tool_dict = convert_tool_ids_to_dict(extracted_tools)
        if code_skipped and not html_form:
            if cron:
                await self.agent_builder_registry.schedule_agent(cron=cron)
            res = await self.agent_builder_registry.get_agent()
            did_update = False
        else:
            res, did_update = await self.agent_builder_registry.upsert_agent(
                name=agent_name,
                description=agent_description,
                instructions=instructions_payload,
                tools=tool_dict,
                cron=cron,
                html=html_form,
                details=long_description,
            )

        message = "Successfully saved the agent code and plan. Ask the user if they wish to test the saved agent. If you test and it does not work as expected, diagnose and correct the saved agent. Note that it has been added to the context, and you do not need to redefine it in execute_python_code."
        if not html_form and not code_skipped:
            message += " Also, ask the user if they would like to generate a customized, sharable webpage for using this agent."

        tool_result = {
            "id": str(res.id) if res else "unknown",
            "update": did_update,
            "name": agent_name,
            "description": agent_description,
            "message": message,
        }
    except Exception as e:
        extracted_tools = extract_code_tools(python_code) if python_code else []
        tool_result = f"Displaying the final code :\n\n{python_code}\nFinal Name: {agent_name}\nDescription: {agent_description} Saved Tools: {extracted_tools}. Inform the user that the code was not saved due to an error{e}, but do not attempt to save it again. Note that it has been added to the context, and you do not need to redefine it in execute_python_code."

    return tool_result


async def build_or_patch_code(
    self: "CodeActPlaybookAgent",
    state: "CodeActState",
    cron: str | None = None,
    update_code: bool = True,
    update_general_info: bool = True,
    generate_html: bool = False,
):
    """Sub-agent helper: generate new code or patch existing code, save, and emit UI updates.
    Returns: (tool_result: str, new_context: base64 encoded context, agent_name: str | None, agent_description: str | None)
    """
    str(uuid.uuid4())

    base_messages = sanitize_messages(state["messages"])
    plan = state.get("plan")
    plan_text = "\n".join(map(str, plan)) if plan else None
    existing_code = (
        self.agent.instructions.get("script") if self.agent and getattr(self.agent, "instructions", None) else None
    )
    if not plan and self.agent and getattr(self.agent, "instructions", None):
        plan = self.agent.instructions.get("plan")
        plan_text = "\n".join(map(str, plan)) if plan else None

    # 1. Execute parallel tasks
    results = await asyncio.gather(
        _generate_meta(self, state, update_general_info, base_messages),
        _generate_code(self, update_code, existing_code, plan_text, base_messages),
    )

    (agent_name, agent_description, long_description), code_result = results
    code_skipped = code_result.get("skipped", False)

    # 2. Code Execution and Retry Logic (Sandbox)
    python_code, new_context, execution_error = await _execute_and_retry_code(self, code_result, existing_code, state)

    if execution_error:
        tool_result = f"Failed to save the agent. The code contains errors after 3 attempts. Try writing a smaller code snippet if the code size is too large. The error is- \n{execution_error}. The code is- {python_code}."
        return tool_result, new_context, agent_name, agent_description

    # 3. Generate HTML Form (Best Effort)
    html_form = await _generate_html_form(self, python_code, plan_text) if generate_html else None

    # 4. Save Agent
    tool_result = await _save_agent(
        self,
        agent_name,
        agent_description,
        long_description,
        python_code,
        plan,
        cron,
        html_form,
        code_skipped,
    )

    return tool_result, new_context, agent_name, agent_description
