"""Mavi Firecrawl MCP Server - Web Search and Scraping using Firecrawl API"""
import asyncio
import os
import requests
from mcp.server import Server
from mcp.server.stdio import stdio_server
from mcp.types import Tool, TextContent

# Create MCP server
app = Server("mavi-firecrawl-server")

# Firecrawl API configuration
FIRECRAWL_API_KEY = os.getenv("FIRECRAWL_API_KEY")
FIRECRAWL_BASE_URL = "https://api.firecrawl.dev/v1"

# Define the tools
@app.list_tools()
async def list_tools() -> list[Tool]:
    return [
        Tool(
            name="search",
            description="Search the web using Firecrawl API",
            inputSchema={
                "type": "object",
                "properties": {
                    "query": {
                        "type": "string",
                        "description": "Search query"
                    },
                    "limit": {
                        "type": "number",
                        "description": "Maximum number of results (default: 5)",
                        "default": 5
                    }
                },
                "required": ["query"]
            }
        ),
        Tool(
            name="scrape",
            description="Scrape content from a URL using Firecrawl API",
            inputSchema={
                "type": "object",
                "properties": {
                    "url": {
                        "type": "string",
                        "description": "URL to scrape"
                    },
                    "formats": {
                        "type": "array",
                        "items": {"type": "string"},
                        "description": "Output formats (e.g., ['markdown', 'html'])",
                        "default": ["markdown"]
                    }
                },
                "required": ["url"]
            }
        )
    ]

@app.call_tool()
async def call_tool(name: str, arguments: dict) -> list[TextContent]:
    """Handle tool calls"""
    
    # Check if API key is set
    if not FIRECRAWL_API_KEY:
        return [TextContent(
            type="text",
            text="Error: FIRECRAWL_API_KEY environment variable not set. Please set your Firecrawl API key."
        )]
    
    try:
        if name == "search":
            return await handle_search(arguments)
        elif name == "scrape":
            return await handle_scrape(arguments)
        else:
            return [TextContent(
                type="text",
                text=f"Error: Unknown tool '{name}'"
            )]
    except Exception as e:
        return [TextContent(
            type="text",
            text=f"Error: {str(e)}"
        )]

async def handle_search(arguments: dict) -> list[TextContent]:
    """Handle web search using Firecrawl API"""
    query = arguments.get("query")
    limit = arguments.get("limit", 5)
    
    if not query:
        return [TextContent(type="text", text="Error: Query is required")]
    
    try:
        headers = {
            "Authorization": f"Bearer {FIRECRAWL_API_KEY}",
            "Content-Type": "application/json"
        }
        
        payload = {
            "query": query,
            "limit": limit
        }
        
        response = requests.post(
            f"{FIRECRAWL_BASE_URL}/search",
            headers=headers,
            json=payload,
            timeout=30
        )
        
        response.raise_for_status()
        data = response.json()
        
        # Format the results
        if "data" in data and data["data"]:
            results = []
            for item in data["data"]:
                title = item.get("title", "No title")
                url = item.get("url", "No URL")
                description = item.get("description", "No description")
                
                results.append(f"**{title}**\nURL: {url}\nDescription: {description}\n")
            
            result_text = f"Search results for '{query}':\n\n" + "\n---\n".join(results)
            return [TextContent(type="text", text=result_text)]
        else:
            return [TextContent(
                type="text",
                text=f"No results found for query: '{query}'"
            )]
        
    except requests.RequestException as e:
        return [TextContent(
            type="text",
            text=f"Error calling Firecrawl API: {str(e)}"
        )]

async def handle_scrape(arguments: dict) -> list[TextContent]:
    """Handle web scraping using Firecrawl API"""
    url = arguments.get("url")
    formats = arguments.get("formats", ["markdown"])
    
    if not url:
        return [TextContent(type="text", text="Error: URL is required")]
    
    try:
        headers = {
            "Authorization": f"Bearer {FIRECRAWL_API_KEY}",
            "Content-Type": "application/json"
        }
        
        payload = {
            "url": url,
            "formats": formats
        }
        
        response = requests.post(
            f"{FIRECRAWL_BASE_URL}/scrape",
            headers=headers,
            json=payload,
            timeout=30
        )
        
        response.raise_for_status()
        data = response.json()
        
        # Extract content based on format
        if "data" in data:
            content = data["data"]
            
            # Get markdown or HTML content
            text_content = content.get("markdown") or content.get("html") or content.get("content", "")
            
            if text_content:
                # Limit content length
                if len(text_content) > 5000:
                    text_content = text_content[:5000] + "\n\n... (content truncated)"
                
                result = f"Content from {url}:\n\n{text_content}"
                return [TextContent(type="text", text=result)]
            else:
                return [TextContent(
                    type="text",
                    text=f"No content extracted from {url}"
                )]
        else:
            return [TextContent(
                type="text",
                text=f"Error: Unexpected response format from Firecrawl API"
            )]
        
    except requests.RequestException as e:
        return [TextContent(
            type="text",
            text=f"Error calling Firecrawl API: {str(e)}"
        )]

# Run the MCP server
async def main():
    """Run the MCP server"""
    async with stdio_server() as (read_stream, write_stream):
        await app.run(read_stream, write_stream, app.create_initialization_options())

def run():
    """Entry point for the installed package"""
    asyncio.run(main())

if __name__ == "__main__":
    run()
