import logging

from langgraph.graph import StateGraph, START, END
from langchain_core.messages import HumanMessage, AIMessage, SystemMessage

from src.checkpointer import get_checkpointer, _checkpointer
from src.nodes.interview_nodes import generate_questions, chat, performance_generation_node
from src.models.interview_model import InterviewState


builder = StateGraph(InterviewState)
builder.add_node("chat", chat)
builder.add_node("question_generator", generate_questions)
builder.add_node("performance", performance_generation_node)

builder.add_edge(START, "question_generator")
builder.add_conditional_edges(
    "question_generator",
    lambda x: "performance" if x.time_remaining <= 0 else "chat",
    {"performance": "performance", "chat": "chat"}
)
builder.add_edge("performance", "chat")

builder.add_edge("chat", END)


_compiled_graph = None


async def get_graph():
    global _compiled_graph
    if _compiled_graph is None:
        cp = await get_checkpointer()
        _compiled_graph = builder.compile(checkpointer=cp)
        with open("interview_graph.png", "wb") as f:
            f.write(_compiled_graph.get_graph().draw_mermaid_png())
    return _compiled_graph


async def chat_interviewer(thread_id: str, time_remain: int, topic: str, user_input: str):
    graph = await get_graph()
    result = await graph.ainvoke(
        {
            "messages": [HumanMessage(content=user_input)],
            "topic": topic,
            "time_remaining": time_remain
        },
        config={"configurable": {"thread_id": thread_id}}
    )
    return result


async def chat_interviewer_stream(thread_id: str, time_remain: int, topic: str, user_input: str):
    graph = await get_graph()
    final_state = {}
    async for event in graph.astream_events(
        {
            "messages": [HumanMessage(content=user_input)],
            "topic": topic,
            "time_remaining": time_remain
        },
        config={"configurable": {"thread_id": thread_id}},
        version="v2"
    ):
        kind = event.get("event")
        if kind == "on_chat_model_stream":
            if "chat_token" in event.get("tags", []):
                chunk = event["data"]["chunk"]
                token = chunk.content if hasattr(chunk, "content") else ""
                if token:
                    yield {"type": "token", "content": token}
        elif kind == "on_chain_end" and event.get("name") == "LangGraph":
            final_state = event["data"].get("output", {})

    performance = final_state.get("performance")
    if performance:
        perf_dict = performance.model_dump() if hasattr(performance, "model_dump") else performance
        logging.info(f"Performance generated for thread_id: {thread_id}")
        yield {"type": "performance", "content": perf_dict}
    else:
        yield {"type": "done", "content": None}


async def load_conversation(thread_id: str):
    graph = await get_graph()
    state = await graph.aget_state(config={"configurable": {"thread_id": thread_id}})
    return state.values.get("messages", [])


async def deleteThread(thread_id: str):
    try:
        cp = await get_checkpointer()
        state = await cp.aget_tuple(config={"configurable": {"thread_id": thread_id}})
        if state is None:
            logging.info(f"Thread {thread_id} not found, nothing to delete.")
            return False

        await cp.adelete_thread(thread_id=thread_id)
        logging.info(f"Thread {thread_id} deleted successfully.")
        return True
    except Exception as e:
        logging.error(f"Error deleting thread {thread_id}: {e}")
        return False

async def close_checkpointer():
    global _checkpointer
    if _checkpointer is not None:
        try:
            logging.info("Closing checkpointer database connection...")
            # AsyncSqliteSaver has a 'conn' attribute which is an aiosqlite connection
            if hasattr(_checkpointer, 'conn') and _checkpointer.conn:
                await _checkpointer.conn.close()
            _checkpointer = None
            logging.info("Checkpointer connection closed.")
        except Exception as e:
            logging.error(f"Error closing checkpointer: {e}")