import argparse
import asyncio
from datetime import timedelta
import json

from mcp import ClientSession
from mcp.client.stdio import stdio_client

from inspect_library_mcp import library_server
from library_mcp import GUIDE_URI, HOURS
from semantic_notes import api

TOOL = "lookup_hours"


def approved_arguments(calls, selected_branch):
    if not isinstance(calls, list) or len(calls) != 1 or not isinstance(calls[0], dict):
        return None
    function = calls[0].get("function")
    if not isinstance(function, dict) or function.get("name") != TOOL:
        return None
    arguments = function.get("arguments")
    if arguments == {"branch": selected_branch} and selected_branch in HOURS:
        return arguments
    return None


def chat(messages, tools=None):
    payload = {
        "model": "qwen2.5:1.5b",
        "messages": messages,
        "options": {"temperature": 0, "num_ctx": 2048, "num_predict": 140},
        "stream": False,
    }
    if tools:
        payload["tools"] = tools
    return api("/api/chat", payload)["message"]


async def get_answer(branch):
    async with stdio_client(library_server()) as (read, write):
        async with ClientSession(read, write, read_timeout_seconds=timedelta(seconds=15)) as session:
            await session.initialize()
            tools = (await session.list_tools()).tools
            tool = next((item for item in tools if item.name == TOOL), None)
            if tool is None:
                raise ValueError("The library server does not provide lookup_hours.")
            tool_schema = {
                "type": "function",
                "function": {"name": tool.name, "description": tool.description,
                             "parameters": tool.inputSchema},
            }
            messages = [{
                "role": "user",
                "content": f"Use the lookup_hours tool for the {branch} branch.",
            }]
            reply = await asyncio.to_thread(chat, messages, [tool_schema])
            arguments = approved_arguments(reply.get("tool_calls"), branch)
            if arguments is None:
                print(f"Please confirm the branch. No lookup ran for {branch}.")
                return

            result = await session.call_tool(TOOL, arguments)
            data = result.structuredContent
            if result.isError or not isinstance(data, dict):
                raise ValueError("The library server could not complete the lookup.")
            hours = data.get("hours")
            if data.get("branch") != branch or not isinstance(hours, str) or not hours:
                raise ValueError("The tool result did not match the selected branch.")
            print(f"Verified lookup: {branch} {hours}")
            guide = (await session.read_resource(GUIDE_URI)).contents[0].text
            messages.extend([
                reply,
                {"role": "tool", "tool_name": TOOL, "content": json.dumps(data)},
                {"role": "user", "content": (
                    f"State the hours from the tool and one visitor tip from this reference: {guide}"
                )},
            ])
            final = await asyncio.to_thread(chat, messages)
            print("Model reply:", final.get("content") or "(empty)")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Ask a local model to look up one branch.")
    parser.add_argument("branch", choices=HOURS)
    args = parser.parse_args()
    try:
        asyncio.run(get_answer(args.branch))
    except (OSError, KeyError, ValueError) as error:
        raise SystemExit(str(error)) from None