from typing import TypedDict, Annotated, Literal, List, Optional
from pydantic import BaseModel, Field
from langgraph.graph import StateGraph, END
from langgraph.graph.message import add_messages
from langchain_core.messages import (
    AIMessage, BaseMessage, HumanMessage, SystemMessage, ToolMessage
)
from langchain_core.tools import BaseTool

from langchain_core.utils.function_calling import convert_to_openai_tool
from langchain_core.language_models import BaseChatModel
import json


# State schema
class AgentState(TypedDict):
    """State for agent with tool support"""
    model: BaseModel
    messages: Annotated[List[BaseMessage], add_messages]
    tools: List[BaseTool]  # Available tools
    tool_call_id_counter: int  # For generating unique IDs


# Structured output schema
class ToolCallSchema(BaseModel):
    name: str = Field(description="Exact name of the tool to call")
    arguments: str = Field(
        description=(
            "JSON representation of the mapping from each tool parameter name "
            "to its value. This MUST be valid JSON with double-quoted keys "
            "and string values where appropriate."
        )
    )


class AgentResponseWithTools(BaseModel):
    reasoning: Optional[str] = Field(
        default=None,
        description=(
            "Brief reasoning explaining your plan and why specific tools "
            "are being called. Do not include tool call JSON here."
        ),
    )
    tool_calls: List[ToolCallSchema] = Field(
        default_factory=list,
        description=(
            "List of tools to call in this step. Leave empty when you "
            "are providing a direct answer only."
        ),
    )
    direct_answer: Optional[str] = Field(
        default=None,
        description=(
            "Direct answer to the user when no more tool calls are "
            "required. For the script-generation flow, this MUST be set "
            "only on the final turn and contain exactly the final Python "
            "script with no additional commentary."
        ),
    )

def call_model_node(state: AgentState) -> dict:
    """Node: Send messages from the state to the model"""
    messages = state["messages"]
    tools = state.get("tools", [])
    counter = state.get("tool_call_id_counter", 0)

    # Filter out ToolMessages and convert them to regular messages
    # Also remove tool_calls from AIMessages before sending to endpoint
    cleaned_messages = []
    for msg in messages:
        if isinstance(msg, ToolMessage):
            # Convert ToolMessage to HumanMessage with tool result
            cleaned_messages.append(
                HumanMessage(
                    content=f"Tool {msg.name} returned: {msg.content}"
                )
            )
        elif isinstance(msg, AIMessage) and hasattr(msg, 'tool_calls') and msg.tool_calls:
            # Create AIMessage without tool_calls for the endpoint
            cleaned_messages.append(
                AIMessage(content=msg.content)
            )
        else:
            cleaned_messages.append(msg)

    # Prepare messages with tool descriptions
    enhanced_messages = list(cleaned_messages)

    if tools:
        tools_prompt = f"""You are a tool-using agent. You have access to these tools:

{json.dumps([convert_to_openai_tool(t) for t in tools])}

Your job is to:
1. Explore the given repository using the filesystem and code-analysis tools.
2. Design and iteratively refine a realistic Python example script that
     demonstrates the main functionality of the repository.
3. Use the virtual-environment tools to run the script, fix import or
     dependency issues, and confirm that it executes successfully.

RESPONSE FORMAT (STRUCTURED OUTPUT):
- Always respond as a JSON object that matches this schema:
    {{
        "reasoning": string or null,
        "tool_calls": [
            {{
                "name": string,  // exact tool name
                "arguments": string  // JSON object as a string
            }}
        ],
        "direct_answer": string or null
    }}

RULES FOR TOOL CALLS:
- When you need to use tools, populate `tool_calls` with one or more
    entries and set `direct_answer` to null.
- For each tool call, `arguments` MUST be a valid JSON object encoded
    as a string (use double quotes, no trailing commas).
    Example: "{{\"file_path\": \"src/main.py\", \"start_line\": 1, \"end_line\": 40}}"
- Use your `reasoning` field to briefly explain why you are calling
    specific tools and how the results will help you.

RULES FOR FINAL ANSWER:
- After you have finished exploring and testing, and you are confident
    the example script works, return the final script in `direct_answer`.
- When you set `direct_answer`, you MUST leave `tool_calls` empty.
- The value of `direct_answer` for the final script-generation step
    MUST contain only the complete Python script, with no extra prose,
    markdown formatting, or explanation.

You MUST respond with valid JSON only, with no surrounding commentary."""

        # Add tool prompt to system message or create one
        has_system = any(isinstance(m, SystemMessage) for m in enhanced_messages)
        if has_system:
            for i, msg in enumerate(enhanced_messages):
                if isinstance(msg, SystemMessage):
                    enhanced_messages[i] = SystemMessage(
                        content=f"{msg.content}\n\n{tools_prompt}"
                    )
                    break
        else:
            enhanced_messages.insert(0, SystemMessage(content=tools_prompt))

    # Call model with structured output
    model = state.get("model")  # Assume model is in state or passed via context
    response = model.with_structured_output(AgentResponseWithTools).invoke(enhanced_messages)

    # Convert structured response to AIMessage with tool_calls
    tool_calls = []
    for tc in response.tool_calls:
        call_id = f"call_{counter:04d}"
        counter += 1
        try:
            tool_calls.append({
                "name": tc.name,
                "args": json.loads(tc.arguments) if tc.arguments else {},
                "id": call_id
            })
        except Exception as e:
            tool_calls.append({
                "name": "__invalid_tool_call__",
                "args": {},
                "id": call_id
            })

    # Build content
    content_parts = []
    if response.reasoning:
        content_parts.append(f"Reasoning: {response.reasoning}")
    if response.direct_answer:
        content_parts.append(response.direct_answer)

    msg_kwargs = {'content': "\n".join(content_parts) if content_parts else ""}
    if tool_calls:
        msg_kwargs['tool_calls'] = tool_calls
    ai_message = AIMessage(**msg_kwargs)

    return {
        "messages": [ai_message],
        "tool_call_id_counter": counter
    }


def call_tool_node(state: AgentState) -> dict:
    """Node: Execute tool calls from the last AI message."""
    messages = state["messages"]
    tools = state.get("tools", [])
    last_message = messages[-1]

    # Create tool registry
    tool_registry = {tool.name: tool for tool in tools}

    # Execute each tool call
    tool_messages = []
    for tool_call in last_message.tool_calls:
        if (tool_call["name"] == "__invalid_tool_call__"):
            error_msg = "Error: Invalid tool call format. Please follow the specified schema."
            tool_messages.append(
                ToolMessage(
                    content=error_msg,
                    tool_call_id=tool_call["id"],
                    name="__invalid_tool_call__"
                )
            )
            continue
        tool_name = tool_call["name"]
        tool_args = tool_call["args"]
        call_id = tool_call["id"]

        if tool_name not in tool_registry:
            error_msg = f"Tool '{tool_name}' not found"
            tool_messages.append(
                ToolMessage(
                    content=error_msg,
                    tool_call_id=call_id,
                    name=tool_name
                )
            )
            continue
        try:
            tool = tool_registry[tool_name]
            result = tool.invoke(tool_args)
            tool_messages.append(
                ToolMessage(
                    content=str(result),
                    tool_call_id=call_id,
                    name=tool_name
                )
            )
        except Exception as e:
            error_msg = f"Error: Internal error when executing the tool: {str(e)}"
            tool_messages.append(
                ToolMessage(
                    content=error_msg,
                    tool_call_id=call_id,
                    name=tool_name
                )
            )

    return {"messages": tool_messages}


def should_continue(state: AgentState) -> Literal["tools", "end"]:
    """Conditional edge: Route based on whether tools were called."""
    last_message = state["messages"][-1]

    if isinstance(last_message, AIMessage) and last_message.tool_calls:
        return "tools"
    return "end"


def create_tool_aware_agent(
    model: BaseChatModel,
    tools: List[BaseTool]
) -> StateGraph:
    """Create an agent with tool support using structured output instead of real tool calls in API."""
    # Build graph
    workflow = StateGraph(AgentState)

    # Add nodes
    workflow.add_node("call_model", call_model_node)
    workflow.add_node("call_tool", call_tool_node)

    # Set entry point
    workflow.set_entry_point("call_model")

    # Add conditional edge: model -> tools or end
    workflow.add_conditional_edges(
        "call_model",
        should_continue,
        {
            "tools": "call_tool",
            "end": END
        }
    )

    # Tool execution loops back to model
    workflow.add_edge("call_tool", "call_model")

    # Compile with initial state setup
    def setup_state(messages, **kwargs):
        return {
            "messages": messages if isinstance(messages, list) else [HumanMessage(content=str(messages))],
            "tools": tools,
            "model": model,
            "tool_call_id_counter": 0,
            **kwargs
        }

    compiled = workflow.compile()
    compiled.setup_state = setup_state

    return compiled