from dotenv import load_dotenv

load_dotenv()

import os, asyncio, json

from langchain.chat_models import init_chat_model
from langchain.agents import create_agent, AgentState
from langchain.messages import HumanMessage, AIMessage, ToolMessage
from langchain.tools import tool, ToolRuntime

from langchain_mcp_adapters.client import MultiServerMCPClient

from langchain_community.utilities import SQLDatabase
from langchain_community.document_loaders import PyPDFLoader

from langgraph.checkpoint.memory import InMemorySaver
from langgraph.types import Command

from mcp.shared.exceptions import McpError
from mcp.types import CallToolResult, TextContent

from typing import Dict, Any
from tavily import TavilyClient
from pprint import pprint
from dataclasses import dataclass
from datetime import date



DEEPSEEK_MODEL = "deepseek-chat"
DEEPSEEK_BASE_URL = "https://api.deepseek.com"

def deepseek_model(model: str = DEEPSEEK_MODEL, max_tokens=1000, **kwargs):
    return init_chat_model(
        model=model,
        # DeepSeek uses LangChain's OpenAI-compatible transport.
        model_provider="openai",
        api_key=os.environ["DEEPSEEK_API_KEY"],
        base_url=DEEPSEEK_BASE_URL,
        max_tokens=max_tokens,
        **kwargs,
    )

tavily_client = TavilyClient()
db = SQLDatabase.from_uri("sqlite:///resources/Chinook.db")


@tool
def web_search(query: str) -> Dict[str, Any]:
    """Search the web for information"""
    return tavily_client.search(query)

@tool
def sql_query(query: str) -> str:
    """Obtain information from the database using SQL queries"""
    try:
        return db.run(query)
    except Exception as e:
        return f"Error: {e}"

from dataclasses import dataclass

@dataclass
class UserRole:
    user_role: str = "external"

from langchain.agents.middleware import wrap_model_call, ModelRequest, ModelResponse
from typing import Callable

@wrap_model_call
def dynamic_tool_call(request: ModelRequest, 
handler: Callable[[ModelRequest], ModelResponse]) -> ModelResponse:

    """Dynamically call tools based on the runtime context"""

    user_role = request.runtime.context.user_role
    
    if user_role == "internal":
        pass # internal users get access to all tools
    else:
        tools = [web_search] # external users only get access to web search
        request = request.override(tools=tools) 

    return handler(request)

from langchain.agents import create_agent

agent = create_agent(
    model=deepseek_model(),
    tools=[web_search, sql_query],
    middleware=[dynamic_tool_call],
    context_schema=UserRole
)

from langchain.messages import HumanMessage

response = agent.invoke(
    {"messages": [HumanMessage(content="How many artists are in the database?")]},
    context={"user_role": "external"}
)

print(response["messages"][-1].content)