"""Search and lookup MCP tools for YFinance.

This module provides tools for searching stocks and looking up financial instruments.
"""

import yfinance as yf
from typing import Literal, Optional
from mcp.server.fastmcp import FastMCP

from ..utils.serializers import df_to_json


def register(mcp: FastMCP):
    """Register all search-related tools with the MCP server."""

    @mcp.tool()
    def search(
        query: str,
        max_results: int = 8,
        news_count: int = 5
    ) -> dict:
        """Search for stocks, ETFs, and related news on Yahoo Finance.
        
        Args:
            query: Search query (company name, ticker symbol, or keyword)
            max_results: Maximum number of stock quotes to return (default 8)
            news_count: Number of news articles to include (default 5)
            
        Returns:
            Search results including matching quotes and related news articles
        """
        search_result = yf.Search(
            query=query,
            max_results=max_results,
            news_count=news_count
        )
        
        return {
            "query": query,
            "quotes": search_result.quotes,
            "news": search_result.news
        }

    @mcp.tool()
    def lookup(
        query: str,
        instrument_type: Literal[
            "all", "stock", "etf", "mutualfund", 
            "index", "future", "currency", "crypto"
        ] = "all",
        count: int = 25
    ) -> dict:
        """Look up financial instruments by type.
        
        This is more specific than search() as it filters by instrument type.
        
        Args:
            query: Search query (company name or ticker symbol)
            instrument_type: Type of financial instrument to search for
                - all: All types
                - stock: Stocks/equities
                - etf: Exchange-traded funds
                - mutualfund: Mutual funds
                - index: Market indices
                - future: Futures contracts
                - currency: Forex currencies
                - crypto: Cryptocurrencies
            count: Number of results to return (default 25)
            
        Returns:
            List of matching financial instruments with details
        """
        lookup_result = yf.Lookup(query=query)
        
        # Map user-friendly names to API names
        type_map = {
            "all": "all",
            "stock": "equity",
            "etf": "etf",
            "mutualfund": "mutualfund",
            "index": "index",
            "future": "future",
            "currency": "currency",
            "crypto": "cryptocurrency"
        }
        
        api_type = type_map[instrument_type]
        
        # Get results based on type
        if api_type == "all":
            df = lookup_result.get_all(count=count)
        elif api_type == "equity":
            df = lookup_result.get_stock(count=count)
        elif api_type == "etf":
            df = lookup_result.get_etf(count=count)
        elif api_type == "mutualfund":
            df = lookup_result.get_mutualfund(count=count)
        elif api_type == "index":
            df = lookup_result.get_index(count=count)
        elif api_type == "future":
            df = lookup_result.get_future(count=count)
        elif api_type == "currency":
            df = lookup_result.get_currency(count=count)
        else:  # cryptocurrency
            df = lookup_result.get_cryptocurrency(count=count)
        
        return df_to_json(df)
