# Copyright 2026 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Async Agent Designer for generating blueprints, instructions, and tools."""

import json
import logging
import textwrap
from typing import Any

from cxas_scrapi.migration.data_models import IRTool, MigrationIR
from cxas_scrapi.migration.prompts import Prompts

logger = logging.getLogger(__name__)


class AsyncAgentDesigner:
    """Handles Step 2: Architecture Planning, Tool Generation, and
    Instruction Formatting."""

    def __init__(self, gemini_client: Any) -> None:
        self.gemini = gemini_client

    @staticmethod
    def _validate_tree_view(tree_view: str) -> None:
        """Validates that tree_view is provided."""
        if not tree_view or not tree_view.strip():
            raise ValueError(
                "tree_view is required for generative synthesis. "
                "Please use the 'flow_visualizer.py' module to generate the "
                "tree view string for this flow before calling the designer."
            )

    @staticmethod
    def _get_available_toolsets_context(
        ir_tools_dict: dict[str, IRTool],
    ) -> str:
        """Formats the loaded OpenAPI toolsets into a clean string for the LLM
        context."""
        toolset_summaries = []
        for t_id, t_data in ir_tools_dict.items():
            if t_data.type == "TOOLSET":
                ops = t_data.operation_ids
                name = t_data.payload.get("display_name", t_id)
                # Webhook meta might be in payload or extra field in IRTool if
                # we add it. For now, let's assume it might be in the payload
                # under 'webhook_meta' or similar.
                meta = t_data.payload.get("webhook_meta", {})

                summary = (
                    f"- Toolset: '{name}' | OpenAPI operation_id: "
                    f"'{ops[0] if ops else 'unknown'}'"
                )

                # Append DFCX specific metadata if it's a webhook
                if meta:
                    summary += f"\n  Webhook Type: {meta.get('webhook_type')}"
                    summary += f"\n  Original URI: {meta.get('original_uri')}"
                    if meta.get("request_body_template"):
                        indent_template = textwrap.indent(
                            meta.get("request_body_template"), "    "
                        )
                        summary += (
                            "\n  Request Payload Template (DFCX Format):\n"
                            f"{indent_template}"
                        )
                    if meta.get("parameter_mapping"):
                        indent_mapping = textwrap.indent(
                            json.dumps(meta.get("parameter_mapping"), indent=2),
                            "    ",
                        )
                        summary += (
                            "\n  Response Parameter Mapping (JSONPath -> "
                            f"Agent Variable):\n{indent_mapping}"
                        )

                toolset_summaries.append(summary)

        return (
            "\n\n".join(toolset_summaries)
            if toolset_summaries
            else "None available."
        )

    @staticmethod
    def build_2a_shared_context(target_ir: MigrationIR) -> tuple[str, str]:
        """Return (system_prompt, shared_content) for a 2A Gemini context cache.

        The shared content covers inputs 2-4 (global vars, toolsets, tools),
        which are identical for every group in a single synthesis run.
        """
        global_vars = json.dumps(
            {
                name: data.get("schema", {}).get("type", "UNKNOWN")
                for name, data in target_ir.parameters.items()
            },
            indent=2,
        )
        toolsets = AsyncAgentDesigner._get_available_toolsets_context(
            target_ir.tools
        )
        tools = AsyncAgentDesigner._get_available_tools_context(target_ir.tools)
        system_prompt = Prompts.STEP_3A_CONSOLIDATION_ARCHITECTURE["system"]
        shared_content = Prompts.STEP_3A_CONSOLIDATION_ARCHITECTURE[
            "cache_shared_template"
        ].format(
            global_variables=global_vars,
            available_backend_toolsets=toolsets,
            available_tools=tools,
        )
        return system_prompt, shared_content

    @staticmethod
    def build_2b_shared_context(target_ir: MigrationIR) -> tuple[str, str]:
        """Return (system_prompt, shared_content) for a 2B Gemini context cache.

        The shared content covers input 3 (available tools), which is the same
        across all groups in a single synthesis run.
        """
        tools = AsyncAgentDesigner._get_available_tools_context(target_ir.tools)
        system_prompt = Prompts.STEP_3B_CONSOLIDATION_INSTRUCTIONS["system"]
        shared_content = Prompts.STEP_3B_CONSOLIDATION_INSTRUCTIONS[
            "cache_shared_template"
        ].format(available_tools=tools)
        return system_prompt, shared_content

    @staticmethod
    def _get_available_tools_context(
        ir_tools_dict: dict[str, IRTool],
    ) -> str:
        """Render the FULL tool registry (TOOLSET + PYTHON + TOOL) as a
        flat list of exact tool IDs the LLM may reference in
        ``{@TOOL: …}`` and ``{@AGENT: …}`` directives.

        Unlike :meth:`_get_available_toolsets_context` (which is OpenAPI-
        only and lists tools by their human display_name), this helper
        emits the *exact key* under which the tool lives in
        ``MigrationIR.tools`` — the key Gemini must paste verbatim
        without re-suffixing (``_wrapper``, ``_tool``) or paraphrasing.

        Also includes the ``end_session`` sentinel that's always
        auto-registered at deploy time.
        """
        by_type: dict[str, list[str]] = {
            "TOOLSET": [],
            "PYTHON": [],
            "TOOL": [],
        }
        for tool_id, tool in ir_tools_dict.items():
            by_type.setdefault(tool.type, []).append(tool_id)

        lines: list[str] = []
        for kind in ("TOOLSET", "PYTHON", "TOOL"):
            tids = sorted(by_type.get(kind, []))
            if not tids:
                continue
            lines.append(f"### {kind} tools ({len(tids)}):")
            for tid in tids:
                lines.append(f"- {tid}")
            lines.append("")

        # SYSTEM-injected sentinels (always available at runtime).
        lines.append("### SYSTEM tools (always available):")
        lines.append("- end_session")
        return "\n".join(lines).rstrip()

    async def run_step_2a(
        self,
        flow_name: str,
        tree_view: str,
        target_ir: MigrationIR,
        available_groups: str | None = None,
        self_group: str | None = None,
        cached_content_name: str | None = None,
    ) -> dict[str, Any]:
        """Runs the Principal Architect prompt to generate the JSON Blueprint.

        ``available_groups`` is an optional pre-rendered string listing
        the other consolidated group names + their absorbed members.
        When supplied (consolidation flow), the prompt includes it so
        Gemini's blueprint references real group names in
        ``exit_routes`` / transition targets. ``self_group`` is the
        current group being synthesized, surfaced in the prompt so
        Gemini doesn't recommend transferring to itself.

        ``cached_content_name`` is the resource name of a Gemini context
        cache holding inputs 2-4 (global vars, toolsets, tools). When
        provided the per-call prompt contains only per-group content
        (input 1 + input 5 + output format) and the cache name is
        forwarded to generate_async.
        """
        AsyncAgentDesigner._validate_tree_view(tree_view)
        logger.info(
            f"[{flow_name}] Starting 2A: Architecture Expert Blueprinting"
        )

        if available_groups is not None:
            system_prompt = Prompts.STEP_3A_CONSOLIDATION_ARCHITECTURE["system"]
            if cached_content_name:
                prompt_2a = Prompts.STEP_3A_CONSOLIDATION_ARCHITECTURE[
                    "cache_per_group_template"
                ].format(
                    flow_name=flow_name,
                    resource_visualization=tree_view,
                    available_groups=available_groups,
                    self_group=self_group or flow_name,
                )
            else:
                global_vars_context = json.dumps(
                    {
                        param_name: param_data.get("schema", {}).get(
                            "type", "UNKNOWN"
                        )
                        for param_name, param_data in (
                            target_ir.parameters.items()
                        )
                    },
                    indent=2,
                )
                toolset_context = (
                    AsyncAgentDesigner._get_available_toolsets_context(
                        target_ir.tools
                    )
                )
                available_tools_context = (
                    AsyncAgentDesigner._get_available_tools_context(
                        target_ir.tools
                    )
                )
                prompt_2a = Prompts.STEP_3A_CONSOLIDATION_ARCHITECTURE[
                    "template"
                ].format(
                    flow_name=flow_name,
                    resource_visualization=tree_view,
                    global_variables=global_vars_context,
                    available_backend_toolsets=toolset_context,
                    available_tools=available_tools_context,
                    available_groups=available_groups,
                    self_group=self_group or flow_name,
                )
        else:
            global_vars_context = json.dumps(
                {
                    param_name: param_data.get("schema", {}).get(
                        "type", "UNKNOWN"
                    )
                    for param_name, param_data in target_ir.parameters.items()
                },
                indent=2,
            )
            toolset_context = (
                AsyncAgentDesigner._get_available_toolsets_context(
                    target_ir.tools
                )
            )
            system_prompt = Prompts.STEP_2A_ARCHITECTURE_EXPERT["system"]
            prompt_2a = Prompts.STEP_2A_ARCHITECTURE_EXPERT["template"].format(
                flow_name=flow_name,
                resource_visualization=tree_view,
                global_variables=global_vars_context,
                available_backend_toolsets=toolset_context,
            )

        response_raw = await self.gemini.generate_async(
            prompt=prompt_2a,
            system_prompt=system_prompt,
            cached_content_name=cached_content_name,
        )

        blueprint = {}
        if response_raw:
            try:
                json_str = (
                    response_raw.replace("```json", "")
                    .replace("```", "")
                    .strip()
                )
                json_start = json_str.find("{")
                if json_start != -1:
                    json_str = json_str[json_start:]
                blueprint = json.loads(json_str)
                logger.info(
                    f"[{flow_name}] ✅ 2A: Architecture Blueprint "
                    "Generated Successfully"
                )
            except Exception as e:
                logger.warning(
                    f"[{flow_name}] ⚠️ Error parsing 2A Blueprint JSON: {e}"
                )
                blueprint = {
                    "error": "JSON Parse Failure",
                    "raw_response": response_raw,
                }
        return blueprint

    async def run_step_2b_instructions(
        self,
        flow_name: str,
        blueprint: dict[str, Any],
        tree_view: str,
        target_ir: MigrationIR | None = None,
        available_groups: str | None = None,
        self_group: str | None = None,
        cached_content_name: str | None = None,
        feedback: str | None = None,
    ) -> str:
        """Runs the Instructions Expert prompt to generate the PIF XML.

        ``target_ir`` is used to render the exact tool registry into the
        prompt so Gemini can only reference real tool IDs. Optional for
        back-compat with callers that haven't been updated; when omitted,
        the prompt falls back to a generic "use only blueprint tools"
        directive (the historical behavior).

        ``available_groups`` / ``self_group``: optional consolidated-
        group context. When supplied, the prompt includes a list of
        valid ``{@AGENT: …}`` transfer targets so Gemini stops inventing
        agent names. Used by the StructuralConsolidator path; ignored
        by 1:1 migration.

        ``cached_content_name`` is the resource name of a Gemini context
        cache holding input 3 (available tools). When provided the
        per-call prompt contains only per-group content (inputs 1, 2, 4
        + output format) and available_tools_context is not recomputed.

        ``feedback`` is an optional addendum appended to the prompt
        body. Used by :meth:`StructuralConsolidator.synthesize_instructions`
        when a first-pass response fails canonical-XML validation and we
        re-prompt Gemini with the lint diagnostics so it can correct the
        schema violations.
        """
        AsyncAgentDesigner._validate_tree_view(tree_view)
        logger.info(
            f"[{flow_name}] Starting 2B: Instructions Expert (XML Generation)"
        )

        blueprint_json_str = json.dumps(blueprint, indent=2)

        if available_groups is not None:
            system_prompt = Prompts.STEP_3B_CONSOLIDATION_INSTRUCTIONS["system"]
            if cached_content_name:
                prompt_2b = Prompts.STEP_3B_CONSOLIDATION_INSTRUCTIONS[
                    "cache_per_group_template"
                ].format(
                    agent_name=flow_name,
                    architecture_blueprint=blueprint_json_str,
                    resource_visualization=tree_view,
                    available_groups=available_groups,
                    self_group=self_group or flow_name,
                )
            else:
                available_tools_context = (
                    AsyncAgentDesigner._get_available_tools_context(
                        target_ir.tools
                    )
                    if target_ir is not None
                    else (
                        "(not provided - use tools from the "
                        "Architecture Blueprint only)"
                    )
                )
                prompt_2b = Prompts.STEP_3B_CONSOLIDATION_INSTRUCTIONS[
                    "template"
                ].format(
                    agent_name=flow_name,
                    architecture_blueprint=blueprint_json_str,
                    resource_visualization=tree_view,
                    available_tools=available_tools_context,
                    available_groups=available_groups,
                    self_group=self_group or flow_name,
                )
        else:
            system_prompt = Prompts.STEP_2B_INSTRUCTIONS_EXPERT["system"]
            prompt_2b = Prompts.STEP_2B_INSTRUCTIONS_EXPERT["template"].format(
                agent_name=flow_name,
                architecture_blueprint=blueprint_json_str,
                resource_visualization=tree_view,
            )

        if feedback:
            prompt_2b = f"{prompt_2b}\n\n{feedback}"

        response_raw = await self.gemini.generate_async(
            prompt=prompt_2b,
            system_prompt=system_prompt,
            cached_content_name=cached_content_name,
        )

        xml_instructions = ""
        if response_raw:
            xml_instructions = (
                response_raw.replace("```xml", "").replace("```", "").strip()
            )
            logger.info(
                f"[{flow_name}] ✅ 2B: XML Instructions Generated Successfully"
            )
        else:
            logger.error(
                f"[{flow_name}] ❌ 2B: LLM returned empty response "
                "for instructions."
            )

        return xml_instructions

    async def run_step_2c_tools_and_callbacks(
        self,
        flow_name: str,
        blueprint: dict[str, Any],
        tree_view: str,
        target_ir: MigrationIR,
    ) -> dict[str, Any]:
        """Runs the Tools & Callbacks Expert prompt to generate Python Code."""
        AsyncAgentDesigner._validate_tree_view(tree_view)
        logger.info(
            f"[{flow_name}] Starting 2C: Tools & Callbacks Expert "
            "(Python Generation)"
        )

        blueprint_json_str = json.dumps(blueprint, indent=2)
        global_vars_context = json.dumps(
            {
                param_name: param_data.get("schema", {}).get("type", "UNKNOWN")
                for param_name, param_data in target_ir.parameters.items()
            },
            indent=2,
        )
        toolset_context = AsyncAgentDesigner._get_available_toolsets_context(
            target_ir.tools
        )

        prompt_2c = Prompts.STEP_2C_TOOLS_AND_CALLBACKS_EXPERT[
            "template"
        ].format(
            agent_name=flow_name,
            architecture_blueprint=blueprint_json_str,
            resource_visualization=tree_view,
            global_variables=global_vars_context,
            available_backend_toolsets=toolset_context,
        )

        response_raw = await self.gemini.generate_async(
            prompt=prompt_2c,
            system_prompt=Prompts.STEP_2C_TOOLS_AND_CALLBACKS_EXPERT["system"],
        )

        tools_and_callbacks = {"tools": [], "callbacks": {}}
        if response_raw:
            try:
                json_str = (
                    response_raw.replace("```json", "")
                    .replace("```", "")
                    .strip()
                )
                json_start = json_str.find("{")
                if json_start != -1:
                    json_str = json_str[json_start:]
                tools_and_callbacks = json.loads(json_str)
                logger.info(
                    f"[{flow_name}] ✅ 2C: Python Tools & Callbacks "
                    "Generated Successfully"
                )
            except Exception as e:
                logger.warning(
                    f"[{flow_name}] ⚠️ Error parsing 2C Tools JSON: {e}"
                )
                tools_and_callbacks = {
                    "error": "JSON Parse Failure",
                    "raw_response": response_raw,
                }
        return tools_and_callbacks