tools/data_fetcher.py

"""
CFA AI Agent - Real-time Financial Data Fetcher
This module handles fetching real-time financial data using yfinance.
"""

import yfinance as yf
import pandas as pd
import numpy as np
from typing import Dict, List, Optional, Union
from datetime import datetime, timedelta
from langchain.tools import tool


@tool
def get_stock_price(ticker: str) -> Dict[str, Union[float, str]]:
    """
    Get current stock price and basic information.

    Args:
        ticker: Stock ticker symbol

    Returns:
        Dictionary with current price and market data
    """
    try:
        stock = yf.Ticker(ticker.upper())
        info = stock.info

        # Get latest price data
        hist = stock.history(period="1d")
        if hist.empty:
            raise ValueError(f"No data available for ticker {ticker}")

        current_price = hist['Close'].iloc[-1]
        previous_close = info.get('previousClose', current_price)
        change = current_price - previous_close
        change_percent = (change / previous_close) * 100 if previous_close != 0 else 0

        return {
            "ticker": ticker.upper(),
            "company_name": info.get('longName', 'Unknown'),
            "current_price": round(current_price, 2),
            "previous_close": round(previous_close, 2),
            "change": round(change, 2),
            "change_percent": round(change_percent, 2),
            "volume": hist['Volume'].iloc[-1],
            "market_cap": info.get('marketCap'),
            "currency": info.get('currency', 'USD'),
            "last_updated": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        }
    except Exception as e:
        return {"error": f"Failed to fetch stock price for {ticker}: {str(e)}"}


@tool
def get_historical_data(
    ticker: str,
    period: str = "1y",
    interval: str = "1d"
) -> Dict[str, Union[List, str]]:
    """
    Get historical stock data for analysis.

    Args:
        ticker: Stock ticker symbol
        period: Time period (1d, 5d, 1mo, 3mo, 6mo, 1y, 2y, 5y, 10y, ytd, max)
        interval: Data interval (1m, 2m, 5m, 15m, 30m, 60m, 90m, 1h, 1d, 5d, 1wk, 1mo, 3mo)

    Returns:
        Dictionary with historical price data and statistics
    """
    try:
        stock = yf.Ticker(ticker.upper())
        hist = stock.history(period=period, interval=interval)

        if hist.empty:
            raise ValueError(f"No historical data available for {ticker}")

        # Calculate basic statistics
        returns = hist['Close'].pct_change().dropna()

        stats = {
            "ticker": ticker.upper(),
            "period": period,
            "interval": interval,
            "data_points": len(hist),
            "start_date": hist.index[0].strftime("%Y-%m-%d"),
            "end_date": hist.index[-1].strftime("%Y-%m-%d"),

            # Price statistics
            "highest_price": round(hist['High'].max(), 2),
            "lowest_price": round(hist['Low'].min(), 2),
            "avg_price": round(hist['Close'].mean(), 2),
            "current_price": round(hist['Close'].iloc[-1], 2),

            # Return statistics
            "total_return": round(((hist['Close'].iloc[-1] / hist['Close'].iloc[0]) - 1) * 100, 2),
            "volatility": round(returns.std() * np.sqrt(252) * 100, 2),  # Annualized volatility
            "avg_daily_return": round(returns.mean() * 100, 4),
            "max_daily_gain": round(returns.max() * 100, 2),
            "max_daily_loss": round(returns.min() * 100, 2),

            # Volume statistics
            "avg_volume": int(hist['Volume'].mean()),
            "max_volume": int(hist['Volume'].max()),
            "min_volume": int(hist['Volume'].min()),

            # Recent data (last 5 days)
            "recent_prices": hist['Close'].tail(5).round(2).tolist(),
            "recent_dates": [date.strftime("%Y-%m-%d") for date in hist.index[-5:]],
            "recent_volumes": hist['Volume'].tail(5).tolist()
        }

        return stats
    except Exception as e:
        return {"error": f"Failed to fetch historical data for {ticker}: {str(e)}"}


@tool
def get_company_info(ticker: str) -> Dict[str, Union[str, float, int]]:
    """
    Get comprehensive company information and fundamentals.

    Args:
        ticker: Stock ticker symbol

    Returns:
        Dictionary with company information and key metrics
    """
    try:
        stock = yf.Ticker(ticker.upper())
        info = stock.info

        company_data = {
            "ticker": ticker.upper(),
            "company_name": info.get('longName', 'Unknown'),
            "sector": info.get('sector', 'Unknown'),
            "industry": info.get('industry', 'Unknown'),
            "country": info.get('country', 'Unknown'),
            "website": info.get('website', 'N/A'),
            "business_summary": info.get('longBusinessSummary', 'N/A'),

            # Key executives
            "ceo": info.get('companyOfficers', [{}])[0].get('name', 'N/A') if info.get('companyOfficers') else 'N/A',

            # Financial metrics
            "market_cap": info.get('marketCap'),
            "enterprise_value": info.get('enterpriseValue'),
            "shares_outstanding": info.get('sharesOutstanding'),
            "float_shares": info.get('floatShares'),

            # Employee info
            "full_time_employees": info.get('fullTimeEmployees'),

            # Exchange info
            "exchange": info.get('exchange', 'Unknown'),
            "quote_type": info.get('quoteType', 'Unknown'),
            "currency": info.get('currency', 'USD'),

            # ESG scores (if available)
            "esg_scores": info.get('esgScores'),
            "sustainability_score": info.get('sustainabilityScore'),

            # Analyst recommendations
            "recommendation": info.get('recommendationKey', 'N/A'),
            "target_high_price": info.get('targetHighPrice'),
            "target_low_price": info.get('targetLowPrice'),
            "target_mean_price": info.get('targetMeanPrice'),
            "number_of_analyst_opinions": info.get('numberOfAnalystOpinions'),

            # Risk metrics
            "audit_risk": info.get('auditRisk'),
            "board_risk": info.get('boardRisk'),
            "compensation_risk": info.get('compensationRisk'),
            "shareholder_rights_risk": info.get('shareHolderRightsRisk'),
            "overall_risk": info.get('overallRisk')
        }

        # Remove None values
        company_data = {k: v for k, v in company_data.items() if v is not None}

        return company_data
    except Exception as e:
        return {"error": f"Failed to fetch company info for {ticker}: {str(e)}"}


@tool
def get_financial_statements(ticker: str) -> Dict[str, Union[pd.DataFrame, str]]:
    """
    Get financial statements (income statement, balance sheet, cash flow).

    Args:
        ticker: Stock ticker symbol

    Returns:
        Dictionary with financial statement data
    """
    try:
        stock = yf.Ticker(ticker.upper())

        # Fetch financial statements
        income_stmt = stock.financials
        balance_sheet = stock.balance_sheet
        cash_flow = stock.cashflow

        result = {
            "ticker": ticker.upper(),
            "has_income_statement": not income_stmt.empty,
            "has_balance_sheet": not balance_sheet.empty,
            "has_cash_flow": not cash_flow.empty,
        }

        # Convert to dictionaries for easier handling
        if not income_stmt.empty:
            result["income_statement_years"] = [str(col.year) for col in income_stmt.columns]
            result["total_revenue"] = income_stmt.loc['Total Revenue'].to_dict() if 'Total Revenue' in income_stmt.index else {}
            result["net_income"] = income_stmt.loc['Net Income'].to_dict() if 'Net Income' in income_stmt.index else {}

        if not balance_sheet.empty:
            result["balance_sheet_years"] = [str(col.year) for col in balance_sheet.columns]
            result["total_assets"] = balance_sheet.loc['Total Assets'].to_dict() if 'Total Assets' in balance_sheet.index else {}
            result["total_debt"] = balance_sheet.loc['Total Debt'].to_dict() if 'Total Debt' in balance_sheet.index else {}

        if not cash_flow.empty:
            result["cash_flow_years"] = [str(col.year) for col in cash_flow.columns]
            result["operating_cash_flow"] = cash_flow.loc['Operating Cash Flow'].to_dict() if 'Operating Cash Flow' in cash_flow.index else {}
            result["free_cash_flow"] = cash_flow.loc['Free Cash Flow'].to_dict() if 'Free Cash Flow' in cash_flow.index else {}

        return result
    except Exception as e:
        return {"error": f"Failed to fetch financial statements for {ticker}: {str(e)}"}


@tool
def get_market_indices() -> Dict[str, Dict[str, Union[float, str]]]:
    """
    Get current prices and performance of major market indices.

    Returns:
        Dictionary with major market index data
    """
    try:
        indices = {
            "S&P 500": "^GSPC",
            "Dow Jones": "^DJI",
            "NASDAQ": "^IXIC",
            "Russell 2000": "^RUT",
            "VIX": "^VIX",
            "10-Year Treasury": "^TNX"
        }

        results = {}
        for name, ticker in indices.items():
            try:
                index = yf.Ticker(ticker)
                hist = index.history(period="2d")
                if not hist.empty:
                    current_price = hist['Close'].iloc[-1]
                    previous_close = hist['Close'].iloc[-2] if len(hist) > 1 else current_price
                    change = current_price - previous_close
                    change_percent = (change / previous_close) * 100 if previous_close != 0 else 0

                    results[name] = {
                        "ticker": ticker,
                        "current_value": round(current_price, 2),
                        "change": round(change, 2),
                        "change_percent": round(change_percent, 2),
                        "last_updated": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
                    }
            except Exception:
                results[name] = {"error": f"Failed to fetch data for {name}"}

        return results
    except Exception as e:
        return {"error": f"Failed to fetch market indices: {str(e)}"}


@tool
def compare_stocks(tickers: List[str], metric: str = "performance") -> Dict[str, Union[List, str]]:
    """
    Compare multiple stocks on various metrics.

    Args:
        tickers: List of stock ticker symbols
        metric: Comparison metric ('performance', 'valuation', 'volatility')

    Returns:
        Dictionary with comparison results
    """
    try:
        if len(tickers) < 2:
            raise ValueError("Need at least 2 tickers for comparison")

        results = {
            "tickers": [t.upper() for t in tickers],
            "metric": metric,
            "comparison_data": {}
        }

        for ticker in tickers:
            try:
                stock = yf.Ticker(ticker.upper())
                info = stock.info
                hist = stock.history(period="1y")

                if metric == "performance":
                    if not hist.empty:
                        ytd_return = ((hist['Close'].iloc[-1] / hist['Close'].iloc[0]) - 1) * 100
                        results["comparison_data"][ticker.upper()] = {
                            "ytd_return": round(ytd_return, 2),
                            "current_price": round(hist['Close'].iloc[-1], 2),
                            "52_week_high": info.get('fiftyTwoWeekHigh'),
                            "52_week_low": info.get('fiftyTwoWeekLow')
                        }

                elif metric == "valuation":
                    results["comparison_data"][ticker.upper()] = {
                        "pe_ratio": info.get('trailingPE'),
                        "price_to_book": info.get('priceToBook'),
                        "price_to_sales": info.get('priceToSalesTrailing12Months'),
                        "market_cap": info.get('marketCap')
                    }

                elif metric == "volatility":
                    if not hist.empty:
                        returns = hist['Close'].pct_change().dropna()
                        volatility = returns.std() * np.sqrt(252) * 100
                        results["comparison_data"][ticker.upper()] = {
                            "volatility": round(volatility, 2),
                            "beta": info.get('beta'),
                            "max_drawdown": round((hist['Close'].min() / hist['Close'].max() - 1) * 100, 2)
                        }

            except Exception as e:
                results["comparison_data"][ticker.upper()] = {"error": str(e)}

        return results
    except Exception as e:
        return {"error": f"Stock comparison failed: {str(e)}"}