from typing import Annotated, List, Dict
from typing_extensions import TypedDict
from langgraph.graph import END, START, StateGraph
from langgraph.graph.message import add_messages
from langgraph.prebuilt import ToolNode, tools_condition
from langchain_core.messages import SystemMessage, HumanMessage, AIMessage, BaseMessage
import json

from config import llm, logger
from tools import fitness_analysis_tool


class State(TypedDict):
    messages     : Annotated[list, add_messages]
    youtube_url  : str
    user_video   : str
    groq_api_key : str
    result       : dict


SYSTEM_PROMPT = """
You are an AI fitness coach. You have one tool: fitness_analysis_tool.

When given a YouTube URL and a user video path, IMMEDIATELY call 
fitness_analysis_tool with:
  - youtube_url     = the YouTube URL provided
  - user_video_path = the user video path provided
  - groq_api_key    = the groq api key provided

Do NOT ask questions. Do NOT ask for context.
Just call the tool immediately with the provided parameters.
"""


def make_tool_graph():
    tools          = [fitness_analysis_tool]
    tool_node      = ToolNode(tools)
    llm_with_tools = llm.bind_tools(tools)

    def call_llm(state: State):
        messages = state["messages"]

        # ── Normalize: accept both dicts and LangChain message objects ──
        normalized = []
        for m in messages:
            if isinstance(m, BaseMessage):
                # Already a LangChain message object — use as-is
                normalized.append(m)
            elif isinstance(m, dict):
                role    = m.get("role", "user")
                content = m.get("content", "")
                if role == "system":
                    normalized.append(SystemMessage(content=content))
                elif role == "assistant":
                    normalized.append(AIMessage(content=content))
                else:
                    normalized.append(HumanMessage(content=content))
            else:
                # Fallback — convert to string
                normalized.append(HumanMessage(content=str(m)))

        # ── Prepend system prompt if not present ──
        has_system = any(isinstance(m, SystemMessage) for m in normalized)
        if not has_system:
            normalized = [SystemMessage(content=SYSTEM_PROMPT)] + normalized

        # ── Inject context into last HumanMessage ──
        for i in range(len(normalized) - 1, -1, -1):
            if isinstance(normalized[i], HumanMessage):
                ctx  = ""
                if state.get("youtube_url"):
                    ctx += f"\nYouTube URL: {state['youtube_url']}"
                if state.get("user_video"):
                    ctx += f"\nUser video path: {state['user_video']}"
                if state.get("groq_api_key"):
                    ctx += f"\nGroq API key: {state['groq_api_key']}"
                if ctx:
                    normalized[i] = HumanMessage(content=normalized[i].content + ctx)
                break

        try:
            response = llm_with_tools.invoke(normalized)
            return {"messages": [response]}
        except Exception as e:
            logger.error(f"LLM call error: {e}")
            return {
                "messages": [AIMessage(content=f"Error: {str(e)}")]
            }

    builder = StateGraph(State)
    builder.add_node("llm",   call_llm)
    builder.add_node("tools", tool_node)
    builder.add_edge(START, "llm")
    builder.add_conditional_edges("llm", tools_condition)
    builder.add_edge("tools", "llm")

    return builder.compile()


tool_agent = make_tool_graph()