# Copyright 2025 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.

import json
import uuid

from IPython import display as ipd
from langgraph.pregel import remote

agent_name = "task-planner"
agent_url = f"http://127.0.0.1:3000/{agent_name}"
id_token = None

# Configure remote agent pointing to local development server
graph = remote.RemoteGraph(
    agent_name,
    url=agent_url,
    headers={"Authorization": f"Bearer {id_token}"} if id_token else {},
)

test_thread = f"test-{uuid.uuid4().hex}"

ipd.Image(graph.get_graph().draw_mermaid_png())

def handle_chunk(chunk: dict, task_idx: int = 0) -> tuple[str, str]:
    if "text" in chunk:
        text = chunk["text"]
        return "text", text

    if "response" in chunk:
        text = chunk["response"]
        return "response", text

    if "guardrail_classification" in chunk:
        is_blocked = chunk["guardrail_classification"]["blocked"]
        classification_emoji = "❌" if is_blocked else "✅"
        reason = chunk["guardrail_classification"]["reason"]

        text = f"Guardrail classification: {classification_emoji}\n\nReason: {reason}"

        return "guardrail_classification", text

    if "router_classification" in chunk:
        target = chunk["router_classification"]["target"]
        reason = chunk["router_classification"]["reason"]

        text = f"Agent Classification: {target}\n\nReason: {reason}"

        return "router_classification", text

    if "function_call" in chunk:
        function_call_dict = chunk["function_call"]

        fn_name = function_call_dict.get("name") or "unknown"
        fn_args = function_call_dict.get("args") or {}

        fn_args_string = ", ".join(f"{k}={v}" for k, v in fn_args.items())
        fn_string = f"**{fn_name}**({fn_args_string})"

        text = f"Calling function... {fn_string}"

        return "fn_call", text

    if "function_response" in chunk:
        function_response_dict = chunk["function_response"]

        fn_name = function_response_dict.get("name") or "unknown"

        if function_response_dict.get("response") is None:
            text = f"Received empty function response (name={fn_name})."

        elif "result" in function_response_dict.get("response"):
            fn_result = function_response_dict["response"]["result"]
            text = "\n\n".join(
                [
                    f"Function result for **{fn_name}**...",
                    "```json",
                    json.dumps(fn_result, indent=2),
                    "```",
                ]
            )

        elif "error" in function_response_dict.get("response"):
            fn_result = function_response_dict["response"]["error"]
            text = f"Function error (name={fn_name})... {fn_result}"

        return "fn_response", text

    if "plan" in chunk:
        plan_dict = chunk["plan"]
        plan_string = _stringify_plan(plan=plan_dict, include_results=False)
        text = f"### Generated execution plan...\n\n{plan_string}"

        return "plan", text

    if "executed_task" in chunk:
        task_idx += 1
        task_dict = chunk["executed_task"]
        task_string = _stringify_task(task=task_dict, include_results=True)
        text = f"### Executed task #{task_idx}...\n\n{task_string}"

        return f"executed_task_{task_idx}", text

    if "error" in chunk:
        text = chunk["error"]

        return "error", text

    text = f"Unhandled chunk. keys={set(chunk.keys())}"

    return "unhandled", text


def _stringify_plan(plan: dict, include_results: bool = True) -> str:
    """Formats an execution plan dictionary into a human-readable string."""
    tasks_str = "\n\n".join(
        f"**Task #{idx + 1}**\n\n"
        + _stringify_task(task, include_results=include_results)
        for idx, task in enumerate(plan["tasks"])
    )

    response = f"**Plan**: {plan['goal']}\n\n{tasks_str}"

    return response


def _stringify_task(task: dict, include_results: bool = True) -> str:
    """Formats a task dictionary into a human-readable string."""
    output = f"**Goal**: {task['goal']}"

    if include_results:
        output += f"\n\n**Result**: {task.get('result') or 'incomplete'}"

    return output

for chunk in graph.stream(
    input={"current_turn": {"user_input": "hi"}},
    config={"configurable": {"thread_id": test_thread}},
    stream_mode="updates",
):
    print(chunk)

task_idx = 0
response_text = ""
current_source = last_source = None
for stream_mode, chunk in graph.stream(
    input={"current_turn": {"user_input": "what products does Cymbal Retail sell?"}},
    config={"configurable": {"thread_id": test_thread}},
    stream_mode=["updates", "custom"],
):
    if stream_mode == "custom":
        assert isinstance(chunk, dict), "Expected dictionary data"

        current_source, text = handle_chunk(chunk, task_idx)

        if "executed_task" in current_source:
            task_idx += 1

        if last_source is not None and last_source != current_source:
            text = "\n\n---\n\n" + text

        last_source = current_source

        response_text += text

        display(ipd.Markdown(response_text), clear=True)

snapshot = graph.get_state(config={"configurable": {"thread_id": test_thread}})

snapshot_list = list(
    graph.get_state_history(config={"configurable": {"thread_id": test_thread}})
)