from pydantic import BaseModel
from typing import Optional, List, Annotated, TypedDict
from langchain_core.messages import BaseMessage, HumanMessage
from langgraph.graph.message import add_messages
from langgraph.graph import StateGraph, START, END
from langgraph.checkpoint.memory import MemorySaver
from langgraph.prebuilt import ToolNode
# tool_node = ToolNode(tools) 

from services.llms.LLM import model_4o_2
from services.agents.LLMAgent import LLMAgent

PROMPT_PATH = 'src/prompts/LearningRoadmapPrompt.md'
with open(PROMPT_PATH, encoding='utf-8') as f:
    prompt_template = f.read()


class Competency(BaseModel):
    skill: str
    description: str
    resources: List[str]

class Roadmap(BaseModel):
    roadmap: List[Competency]
    insight: str

class Roadmaps(BaseModel):
    roadmaps: List[Roadmap]

class AgentState(TypedDict):
    """State for our agent"""
    messages: Annotated[list[BaseMessage], add_messages]
    roadmaps: Roadmaps




def gap_analysis(state: AgentState) -> AgentState:
    """Gatekeeper: answer directly"""
    roadmap_llm = model_4o_2.with_structured_output(Roadmaps)

    gap_analysis_prompt = """\
    {user_profile_and_target}

    Task: Create 3 options of learning roadmap for the user to achieve their target. 
    Each roadmap has 2 keys, 'roadmap' and 'insight'. In the 'roadmap' key provide list of 'skill', 'description', and 'resources' for each competency.
    The 'insight' key should provide a brief summary of the roadmap and how it relates to the user's profile and target.
    """

    user_profile_and_target = state['messages'][-1].content
    if not user_profile_and_target:
        return {"error": "user_profile_and_target must be provided."}
    
    input_gap_analysis_prompt = gap_analysis_prompt.format(
        user_profile_and_target = user_profile_and_target
    )

    result = roadmap_llm.invoke(input_gap_analysis_prompt)
    recommendations = {}
    for i, a_roadmap in enumerate(result.roadmaps):
        recommendations[f'roadmap_{i+1}'] = {}
        recommendations[f'roadmap_{i+1}']['competencies'] = []
        recommendations[f'roadmap_{i+1}']['insight'] = a_roadmap.insight
        for a_competency in a_roadmap.roadmap:
            recommendations[f'roadmap_{i+1}']['competencies'].append({"skill":a_competency.skill, "description": a_competency.description, "resources": a_competency.resources})
    return {**state, "roadmaps": recommendations}


class LearningRoadmapAgent(LLMAgent):
    def __init__(self):
        super().__init__()
        self.user_profile: Optional[str] = None
        self.user_target: Optional[str] = None
        self.thread_id = None
        self.session_id = None
        self.roadmap_agent_graph = StateGraph(AgentState)
        self.roadmap_agent_graph.add_node("gap_analysis", gap_analysis)
        self.roadmap_agent_graph.add_edge(START, "gap_analysis")
        self.roadmap_agent_graph.add_edge("gap_analysis", END)
        self.roadmap_agent = self.roadmap_agent_graph.compile(checkpointer=MemorySaver())


    def generate_roadmap(self, user_profile: str = None, user_target: str = None) -> list:
        """
        Generate a learning roadmap based on user_profile and user_target.
        """

        if user_profile is None or user_target is None:
            return [{"error": "user_profile and user_target must be provided."}]
        
        merged_user_profile_target = """
        Please create roadmap from following input:

        User Profile: {user_profile}

        User Target: {user_target}
        """.format(user_profile=user_profile, user_target=user_target).replace('{','{{').replace('}','}}')

        config = {"configurable": {"thread_id": self.thread_id, "session_id": self.session_id}}
        result = self.roadmap_agent.invoke({"messages": [HumanMessage(content=merged_user_profile_target)]}, config)
        if result and 'roadmaps' in result:
            return result['roadmaps']
        else:
            return [{"error": "Failed to generate roadmap."}]