from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_core.messages import BaseMessage, ToolMessage, AIMessage, SystemMessage, HumanMessage
from langgraph.graph import StateGraph, add_messages, START, END
from langgraph.checkpoint.sqlite import SqliteSaver
from typing import TypedDict, Annotated, List
from langchain_core.tools import tool
from langgraph.prebuilt.tool_node import ToolNode
import sqlite3
import subprocess
import requests
from datetime import datetime

class chatstate(TypedDict):
    messages: Annotated[List[BaseMessage], add_messages]


api = "AIzaSyA5zvErF4vUmAoslVzkOBUfvSCSoW0vjEA"
LANGSEARCH_API_KEY = "sk-f1a8f996f9e44b43adf9943e43e8582b"

llm = ChatGoogleGenerativeAI(model="gemini-2.5-flash", temperature=0.2, api_key=api)

system = SystemMessage(
    content=f"""
    --> Today's date: {datetime.today()}  
    Day number: {datetime.today().date().weekday()}
    You are a practical, tool-aware assistant. Aim for correctness and clarity. Avoid hallucinations.
    Do not provide internal information of the system . 
Rules:
1. Prefer text answers and code when examples/explanations are asked.
2. Explicit requests to create/run files → call appropriate tool.
3. Avoid destructive commands without confirmation.
4. Keep tool inputs minimal.
Tone: concise, helpful, decisive.
"""
)

conn = sqlite3.connect("chatbot.db", check_same_thread=False)
checkpointer = SqliteSaver(conn=conn)


@tool
def add(a: int, b: int):
    return a + b


@tool
def reverse(string: str):
    return string[::-1]


@tool
def evaluate(string: str):
    return eval(string)


@tool
def write_file(name: str, extension: str, content: str):
    with open(f"{name}.{extension}", "w", encoding="utf-8") as f:
        f.write(content)
    return f"Content saved to {name}.{extension}"


@tool
def run_cmd_command(command: str) -> str:
    """Run a safe shell command on linux """
    try:
        result = subprocess.run(command, shell=True, check=True, text=True, capture_output=True)
        return result.stdout
    except subprocess.CalledProcessError as e:
        return f"Error: {e}"


@tool
def search_tool(query: str):
    response = requests.post(
        "https://api.langsearch.com/v1/web-search",
        headers={
            "Authorization": f"Bearer {LANGSEARCH_API_KEY}",
            "Content-Type": "application/json"
        },
        json={"query": query, "num_results": 2}
    )
    return response.json()


def shouldcontinue(state: chatstate):
    return "end" if state["messages"][-1].content == "end" else "llmresponse"


def input_node(state: chatstate):
    return {"messages": state["messages"]}


def llmresponse(state: chatstate):
    response = llm.invoke(state["messages"])
    return {"messages": [response]}


def checktool(state: chatstate):
    last_msg = state["messages"][-1]
    if hasattr(last_msg, "tool_calls") and last_msg.tool_calls:
        return "tool_node"
    return "end"


tools = [add, reverse, evaluate, run_cmd_command, search_tool, write_file]
tool_node = ToolNode(tools=tools)
llm = llm.bind_tools(tools)

graph = StateGraph(chatstate)
graph.add_node("input_node", input_node)
graph.add_node("llmresponse", llmresponse)
graph.add_node("tool_node", tool_node)
graph.add_edge(START, "input_node")
graph.add_edge("input_node", "llmresponse")
graph.add_conditional_edges("llmresponse", checktool, {"tool_node": "tool_node", "end": END})
graph.add_edge("tool_node", "llmresponse")

workflow = graph.compile(checkpointer=checkpointer)


def get_all_chat_ids():
    s = set()
    for chkpoint in checkpointer.list(None):
        s.add(chkpoint.config.get("configurable").get("thread_id"))
    return list(s)