import json
from re import search

import yfinance as yf
from mcp.server import FastMCP

mcp = FastMCP("yfinance-server")


@mcp.tool()
def search_from_yahoo_finance(
        query: str,
        max_results: int = 10,
        news_count: int = 10,
        lists_count: int = 10,
        include_news: bool = True,
        include_lists: bool = True,
        include_quotes: bool = True,
        include_research: bool = True,
        enable_fuzzy_query: bool = True
) -> dict:
    """
    Search from Yahoo Finance with customizable parameters
    
    Args:
        query: The search query (ticker symbol or company name)
        max_results: Maximum number of results to return (default: 10)
        news_count: Number of news items to return (default: 10)
        lists_count: Number of lists to return (default: 10)
        include_news: Whether to include news in results (default: True)
        include_lists: Whether to include lists in results (default: True)
        include_quotes: Whether to include quotes in results (default: True)
        include_research: Whether to include research in results (default: True)
        enable_fuzzy_query: Whether to enable fuzzy matching (default: True)
    
    Returns:
        dict: A dictionary containing the search results with the following possible keys:
            - news: List of news articles
            - lists: List of related lists
            - quotes: List of stock quotes
            - research: List of research items
    """
    searcher = yf.Search(
        query,
        max_results=max_results,
        news_count=news_count if include_news else 0,
        lists_count=lists_count if include_lists else 0,
        include_nav_links=True,
        include_research=include_research,
        include_cultural_assets=True,
        recommended=max_results,
        enable_fuzzy_query=enable_fuzzy_query
    )

    result = searcher.search()

    # 构建返回结果
    response = {}
    if include_news and hasattr(result, 'news') and result.news:
        response['news'] = [{
            'title': item['title'],
            'link': item['link'],
            'publisher': item.get('publisher', ''),
            'published_time': item.get('published_time', '')
        } for item in result.news]

    if include_lists and hasattr(result, 'lists') and result.lists:
        response['lists'] = result.lists

    if include_quotes and hasattr(result, 'quotes') and result.quotes:
        response['quotes'] = result.quotes

    if include_research and hasattr(result, 'research') and result.research:
        response['research'] = result.research

    return response


@mcp.tool()
def get_sector_info(
    sector_key: str,
    include_top_companies: bool = True,
    include_top_etfs: bool = True,
    include_top_mutual_funds: bool = True,
    include_industries: bool = True,
) -> dict:
    """
    Get information about a specific market sector
    
    Args:
        sector_key: The sector key. Valid keys include:
            - basic-materials
            - communication-services  
            - consumer-cyclical
            - consumer-defensive
            - energy
            - financial-services
            - healthcare
            - industrials
            - real-estate
            - technology
            - utilities
        include_top_companies: Whether to include top companies in the sector (default: True)
        include_top_etfs: Whether to include top ETFs in the sector (default: True)
        include_top_mutual_funds: Whether to include top mutual funds (default: True)
        include_industries: Whether to include industries in this sector (default: True)
    
    Returns:
        dict: A dictionary containing the sector information with the following possible keys:
            - name: Sector name
            - overview: Sector overview information
            - top_companies: DataFrame of top companies in the sector
            - top_etfs: Dictionary of top ETF symbols and names
            - top_mutual_funds: Dictionary of top mutual fund symbols and names
            - industries: DataFrame with industries' key, name, symbol, and market weight
            - research_reports: List of research reports related to the sector
    """
    try:
        sector = yf.Sector(sector_key)
        response = {
            'name': sector.name,
            'overview': sector.overview
        }

        if include_top_companies:
            response['top_companies'] = sector.top_companies.to_dict('records')
            
        if include_top_etfs:
            response['top_etfs'] = sector.top_etfs
            
        if include_top_mutual_funds:
            response['top_mutual_funds'] = sector.top_mutual_funds
            
        if include_industries:
            response['industries'] = sector.industries.to_dict('records')
            
        if hasattr(sector, 'research_reports'):
            response['research_reports'] = sector.research_reports
            
        return response
        
    except Exception as e:
        return {
            'error': f"Failed to get sector information: {str(e)}",
            'valid_sectors': [
                'basic-materials', 'communication-services', 'consumer-cyclical',
                'consumer-defensive', 'energy', 'financial-services', 'healthcare',
                'industrials', 'real-estate', 'technology', 'utilities'
            ]
        }


@mcp.tool()
def get_industry_info(
    industry_key: str,
    include_top_companies: bool = True,
    include_top_growth: bool = True,
    include_top_performing: bool = True,
) -> dict:
    """
    Get detailed information about a specific industry
    
    Args:
        industry_key: The key identifier for the industry
        include_top_companies: Whether to include top companies in the industry (default: True)
        include_top_growth: Whether to include top growth companies (default: True)
        include_top_performing: Whether to include top performing companies (default: True)
    
    Returns:
        dict: A dictionary containing the industry information with the following possible keys:
            - key: Industry key identifier
            - name: Industry name
            - symbol: Industry symbol
            - overview: Industry overview information
            - sector_key: Parent sector key
            - sector_name: Parent sector name
            - top_companies: List of top companies in the industry
            - top_growth_companies: List of top growth companies
            - top_performing_companies: List of top performing companies
            - research_reports: List of research reports related to the industry
    """
    try:
        industry = yf.Industry(industry_key)
        
        # 基础信息总是包含
        response = {
            'key': industry.key,
            'name': industry.name,
            'symbol': industry.symbol,
            'overview': industry.overview,
            'sector_key': industry.sector_key,
            'sector_name': industry.sector_name
        }
        
        # 可选信息根据参数决定是否包含
        if include_top_companies and hasattr(industry, 'top_companies'):
            response['top_companies'] = industry.top_companies.to_dict('records')
            
        if include_top_growth and hasattr(industry, 'top_growth_companies'):
            growth_companies = industry.top_growth_companies
            if growth_companies is not None:
                response['top_growth_companies'] = growth_companies.to_dict('records')
            
        if include_top_performing and hasattr(industry, 'top_performing_companies'):
            performing_companies = industry.top_performing_companies
            if performing_companies is not None:
                response['top_performing_companies'] = performing_companies.to_dict('records')
        
        # 研究报告如果存在也包含
        if hasattr(industry, 'research_reports'):
            response['research_reports'] = industry.research_reports
            
        return response
        
    except Exception as e:
        return {
            'error': f"Failed to get industry information: {str(e)}"
        }


if __name__ == "__main__":
    mcp.run(transport="stdio")
