"""Market-related MCP tools for YFinance.

This module provides tools for market calendars, stock screening,
and sector/industry data.
"""

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

from ..utils.serializers import df_to_json, dict_to_json


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

    @mcp.tool()
    def get_market_calendar(
        calendar_type: Literal["earnings", "ipo", "economic", "splits"] = "earnings",
        start: Optional[str] = None,
        end: Optional[str] = None,
        limit: int = 25
    ) -> dict:
        """Get market calendar events (earnings, IPOs, economic events, splits).
        
        Args:
            calendar_type: Type of calendar to retrieve
                - earnings: Upcoming earnings announcements
                - ipo: Upcoming IPO dates
                - economic: Economic events (FOMC, jobs report, etc.)
                - splits: Stock split dates
            start: Start date in YYYY-MM-DD format (default: today)
            end: End date in YYYY-MM-DD format (default: 7 days from start)
            limit: Maximum number of results (default 25)
            
        Returns:
            List of calendar events with dates and details
        """
        calendars = yf.Calendars(start=start, end=end)
        
        if calendar_type == "earnings":
            df = calendars.get_earnings_calendar(limit=limit)
        elif calendar_type == "ipo":
            df = calendars.get_ipo_info_calendar(limit=limit)
        elif calendar_type == "economic":
            df = calendars.get_economic_events_calendar(limit=limit)
        else:  # splits
            df = calendars.get_splits_calendar(limit=limit)
        
        return df_to_json(df)

    @mcp.tool()
    def screen_stocks(
        query_type: Literal[
            "most_actives", "day_gainers", "day_losers",
            "undervalued_growth", "aggressive_small_caps", "growth_tech",
            "undervalued_large_caps", "most_shorted_stocks"
        ] = "most_actives",
        count: int = 25
    ) -> dict:
        """Screen stocks using predefined filters.
        
        Args:
            query_type: Predefined screening query
                - most_actives: Most actively traded stocks today
                - day_gainers: Biggest gainers today
                - day_losers: Biggest losers today
                - undervalued_growth: Undervalued stocks with growth potential
                - aggressive_small_caps: Small cap growth stocks
                - growth_tech: Technology growth stocks
                - undervalued_large_caps: Undervalued large cap stocks
                - most_shorted_stocks: Most heavily shorted stocks
            count: Number of results to return (default 25)
            
        Returns:
            List of stocks matching the screening criteria
        """
        # Use predefined screener queries
        result = yf.screen(query=query_type, count=count)
        return df_to_json(result)

    @mcp.tool()
    def get_sector_data(
        sector_key: str,
        data_type: Literal["industries", "top_etfs", "top_funds"] = "industries"
    ) -> dict:
        """Get data for a market sector.
        
        Args:
            sector_key: Sector identifier (e.g., "technology", "healthcare", 
                "financial-services", "consumer-cyclical", "industrials",
                "communication-services", "consumer-defensive", "energy",
                "basic-materials", "real-estate", "utilities")
            data_type: Type of data to retrieve
                - industries: List of industries in this sector
                - top_etfs: Top ETFs for this sector
                - top_funds: Top mutual funds for this sector
                
        Returns:
            Sector data based on the selected type
        """
        sector = yf.Sector(sector_key)
        
        if data_type == "industries":
            return df_to_json(sector.industries)
        elif data_type == "top_etfs":
            return {"top_etfs": sector.top_etfs}
        else:  # top_funds
            return {"top_mutual_funds": sector.top_mutual_funds}

    @mcp.tool()
    def get_industry_data(
        industry_key: str,
        data_type: Literal["top_companies", "top_etfs", "top_funds"] = "top_companies"
    ) -> dict:
        """Get data for a specific industry.
        
        Args:
            industry_key: Industry identifier (e.g., "software-infrastructure",
                "semiconductors", "internet-retail", "banks-regional")
            data_type: Type of data to retrieve
                - top_companies: Top companies in this industry
                - top_etfs: Top ETFs for this industry
                - top_funds: Top mutual funds for this industry
                
        Returns:
            Industry data based on the selected type
        """
        industry = yf.Industry(industry_key)
        
        if data_type == "top_companies":
            return df_to_json(industry.top_companies)
        elif data_type == "top_etfs":
            return {"top_etfs": industry.top_etfs}
        else:  # top_funds
            return {"top_mutual_funds": industry.top_mutual_funds}
