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,
    )

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

large_model = deepseek_model("deepseek-v4-pro")
standard_model = deepseek_model()


@wrap_model_call
def state_based_model(  request: ModelRequest, 
                        handler: Callable[[ModelRequest], ModelResponse]
                     ) -> ModelResponse:
    """Select model based on State conversation length."""
    # request.messages is a shortcut for request.state["messages"]
    message_count = len(request.messages)  

    if message_count > 10:
        # Long conversation - use model with larger context window
        model = large_model
    else:
        # Short conversation - use efficient model
        model = standard_model

    request = request.override(model=model)  

    return handler(request)

agent = create_agent(
    model=deepseek_model(),
    middleware=[state_based_model],
    system_prompt="You are roleplaying a real life helpful office intern."
)

response = agent.invoke(
    {"messages": [
        HumanMessage(content="Did you water the office plant today?")
        ]}
)

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

print(response["messages"][-1].response_metadata["model_name"])

response = agent.invoke(
    {"messages": [
        HumanMessage(content="Did you water the office plant today?"),
        AIMessage(content="Yes, I gave it a light watering this morning."),
        HumanMessage(content="Has it grown much this week?"),
        AIMessage(content="It's sprouted two new leaves since Monday."),
        HumanMessage(content="Are the leaves still turning yellow on the edges?"),
        AIMessage(content="A little, but it's looking healthier overall."),
        HumanMessage(content="Did you remember to rotate the pot toward the window?"),
        AIMessage(content="I rotated it a quarter turn so it gets more even light."),
        HumanMessage(content="How often should we be fertilizing this plant?"),
        AIMessage(content="About once every two weeks with a diluted liquid fertilizer."),
        HumanMessage(content="When should we expect to have to replace the pot?")
        ]}
)

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

print(response["messages"][-1].response_metadata["model_name"])