# server.py
import asyncio
import sys
from bs4 import BeautifulSoup
from mcp.server import Server, NotificationOptions
from mcp.server.models import InitializationOptions
import mcp.types as types
import mcp.server.stdio
import os
from firecrawl import FirecrawlApp
import argparse

from dotenv import load_dotenv
load_dotenv()

# Initialize the MCP server
server = Server("usman_firecrawl")
# Configuration 


# Tool implementations
async def search_tools(query: str) -> str:
    """Search using FireCrawl API"""
    print(f"[DEBUG] Searching FireCrawl for: {query}", file=sys.stderr)
    try:
        app = FirecrawlApp(api_key=os.getenv("FIRECRAWL_API_KEY"))
        # Perform search using the SDK
        search_result = app.search(query)
        print(f"[DEBUG] Raw search result type: {type(search_result)}", file=sys.stderr)
        
        # Handle search_result as object or dict
        if isinstance(search_result, dict):
            web_results = search_result.get('web', [])
        else:
            web_results = getattr(search_result, 'web', [])
        
        if not web_results:
            print(f"[DEBUG] No web results found in: {search_result}", file=sys.stderr)
            return f"No results found for '{query}'"
            
        results = web_results[:3]  # Get top 3
        output = []
        for i, res in enumerate(results):
            # The results might also be objects or dicts
            if isinstance(res, dict):
                title = res.get('title', 'No Title')
                url = res.get('url', 'No URL')
                snippet = res.get('description', 'No description')
            else:
                title = getattr(res, 'title', 'No Title')
                url = getattr(res, 'url', 'No URL')
                snippet = getattr(res, 'description', 'No description')
            output.append(f"{i+1}. {title}\n   URL: {url}\n   Snippet: {snippet}")
            
        return "\n\n".join(output)
    except Exception as e:
        print(f"[ERROR] FireCrawl Search Error: {e}", file=sys.stderr)
        raise e

def scrape(url: str) -> str:
    """Scrape website content using BeautifulSoup"""
    print(f"[DEBUG] Scraping URL: {url}", file=sys.stderr)
    import requests

    response = requests.get(url)
    soup = BeautifulSoup(response.text, "html.parser")
    # Return visible text only
    text = soup.get_text(separator="\n", strip=True)
    return text[:1000] + ("..." if len(text) > 1000 else "")  # limit output

# List available tools
@server.list_tools()
async def list_tools() -> list[types.Tool]:
    return [
        types.Tool(
            name="search_tools",
            description="Search using FireCrawl internally",
            inputSchema={
                "type": "object",
                "properties": {"query": {"type": "string", "description": "Search query"}},
                "required": ["query"],
            },
        ),
        types.Tool(
            name="scrape",
            description="Scrape website content from URL using BeautifulSoup",
            inputSchema={
                "type": "object",
                "properties": {"url": {"type": "string", "description": "URL to scrape"}},
                "required": ["url"],
            },
        ),
    ]

# Handle tool calls
@server.call_tool()
async def call_tool(name: str, arguments: dict) -> list[types.TextContent]:
    print(f"[DEBUG] Tool called: {name} with args: {arguments}", file=sys.stderr)
    try:
        if name == "search_tools":
            query = arguments.get("query", "")
            result = await search_tools(query)
            return [types.TextContent(type="text", text=result)]
        elif name == "scrape":
            url = arguments.get("url", "")
            result = scrape(url)
            return [types.TextContent(type="text", text=result)]
        else:
            return [types.TextContent(type="text", text=f"Error: Unknown tool '{name}'")]
    except Exception as e:
        print(f"[ERROR] {e}", file=sys.stderr)
        return [types.TextContent(type="text", text=f"Error executing tool '{name}': {str(e)}")]

# Run MCP server
async def main():
    print("Starting usman_firecrawl MCP server...", file=sys.stderr)
    async with mcp.server.stdio.stdio_server() as (read_stream, write_stream):
        await server.run(
            read_stream,
            write_stream,
            InitializationOptions(
                server_name="usman_firecrawl",
                server_version="1.0.0",
                capabilities=server.get_capabilities(
                    notification_options=NotificationOptions(),
                    experimental_capabilities={},
                ),
            ),
        )


def entry_point():
    #using argparser to get the api key
    parser = argparse.ArgumentParser()
    parser.add_argument("--FIRECRAWL_API_KEY",required=True, help="Firecrawl API key")
    args = parser.parse_args()
    os.environ["FIRECRAWL_API_KEY"] = args.FIRECRAWL_API_KEY
    asyncio.run(main())

if __name__ == "__main__":
    entry_point()
