from typing import Literal, Annotated, List

from langchain.chat_models import init_chat_model

from langgraph.graph import START, END, StateGraph, MessagesState

from dotenv import load_dotenv
load_dotenv()

from model import State
from chains import (
    base_prompt_chain,
    user_prompt_chain,
    prompt_selector_chain,
    post_generetor_chain_without_messages,
    post_generator_chain_with_messages,
    post_eval_chain,
    post_perf_eval_chain
)

llm = init_chat_model("groq:meta-llama/llama-4-maverick-17b-128e-instruct")

performance_dict = {
    "before_post_followers": "1735",
    "after_post_followers": "1756",
    "post_likes": "858",
    "post_comments": "89"
}


def prompt_generator(state: State) -> State:

    response = base_prompt_chain.invoke({
        "domain": state['business_domain']
    })
    return {"base_prompts": response.prompts}


def final_prompts(state: State) -> State:

    response = user_prompt_chain.invoke({
        "prompts": state['base_prompts'],
        "business_name": state['business_name'],
        "platform": state['platform'],
        "user_config": state['user_config']
    })
    return {"configured_prompts": response.prompts}


def prompt_selector(state: State) -> State:

    response = prompt_selector_chain.invoke({
        "prompts": state['configured_prompts'],
        "business_domain": state['business_domain'],
        "business_name": state['business_name'],
        "user_config": state['user_config']
    })
    return {"selected_prompt": response.final_prompt}



def post_generator(state: State) -> State:
    """Agent that generates social media posts based on business details and user configuration."""
    messages = state["messages"]

    if not messages:
        generated_post = post_generetor_chain_without_messages.invoke({
            "prompt": state['selected_prompt']
        })

    else:
        generated_post = post_generator_chain_with_messages.invoke({
            "feedback": state['messages'][-1]
        })

    return {"post_content": generated_post}


def post_evaluator(state: State) -> State:
    """Agent that evaluates generated social media post"""

    response = post_eval_chain.invoke({
        "generated_post": state['post_content'],
        "business_domain": state['business_domain'],
        "business_name": state['business_name'],
        "platform": state['platform'],
        "user_config": state['user_config']
    })
    round = state['current_round'] + 1

    return {"messages": response.feedback, "next_step": response.next_step, "current_round": round}


def performance_evaluator(state: State) -> State:
    """Agent that evaluates generated social media post's performance on social media platform"""

    response = post_perf_eval_chain.invoke({
        "actual_post": state['post_content'],
        "platform_name": state['platform'],
        "performance": performance_dict,
        "selected_prompt": state['selected_prompt']
    })

    return {"prompt_accepted": response.accepted_or_not, "post_remarks": response.remarks, "post_rating": response.rating}


def should_continue(state: State) -> State:
    """Determine whether to continue the debate or end it."""
    if state["current_round"] >= 3:
        next_step_value = "END"
    else:
        next_step_value = "CONTINUE"

    return {"next_step": next_step_value}

def go_to_top(state: State) -> State:
    """Determine whether to accept or reject prompt based on post performance"""
    if state['prompt_accepted'] == "NOT_ACCEPTED" or state["post_rating"] <= 5:
        where_to_go_value = "CONTINUE"
    else: 
        where_to_go_value = "END"
    
    return {"where_to_go": where_to_go_value}


graph = StateGraph(State)

graph.add_node("prompt_gen_agent", prompt_generator)
graph.add_node("prompt_config_agent", final_prompts)
graph.add_node("prompt_selector_agent", prompt_selector)
graph.add_node("post_gen_agent", post_generator)
graph.add_node("post_eval_agent", post_evaluator)
graph.add_node("post_perf_eval_agent", performance_evaluator)
graph.add_node("router", should_continue)
graph.add_node("router_2", go_to_top)

graph.add_edge(START, "prompt_gen_agent")
graph.add_edge("prompt_gen_agent", "prompt_config_agent")
graph.add_edge("prompt_config_agent", "prompt_selector_agent")
graph.add_edge("prompt_selector_agent", "post_gen_agent")
graph.add_edge("post_gen_agent", "post_eval_agent")
graph.add_edge("post_eval_agent", "router")
graph.add_conditional_edges(
    source="router",
    path=lambda x: x["next_step"],
    path_map={
        "CONTINUE": "post_gen_agent",
        "END": "post_perf_eval_agent"
    }
)
graph.add_edge("post_perf_eval_agent", "router_2")
graph.add_conditional_edges(
    source="router_2",
    path=lambda x: x["where_to_go"],
    path_map={
        "CONTINUE": "post_gen_agent",
        "END": END
    }
)


builder = graph.compile()

result = builder.invoke({
    "business_domain": "Restraurant",
    "business_name": "Dosa Charcoal",
    "platform": "Instagram",
    "user_config": {
        "tone": "casual and energetic",
        "target_audience": "young professionals",
        "content_preferences": ["Authentic south indian dosa", "Good taste", "restaurant ambiance", "No mention of non-veg"]
    },
    "current_round": 0
})

print("Generated Post     :", result['post_content'].post_content)
print("Total Rounds       :", result['current_round'])
print("Base Prompts       :", result['base_prompts'])
print("Configured Prompts :", result['configured_prompts'])
print("Selected Prompt    :", result['selected_prompt'])
print("Post Rating        :", result['post_rating'])


tell me in this code 
how can improve functionality and performance without adding any extra nodes

just add some logic or any extra fields etc.