#!/usr/bin/env python3
"""
01_data_acquisition.py
----------------------
Step 1: Data Acquisition & Market Diagnostics for Chegg (CHGG)

This script fetches:
- Historical stock price and volume data (last 5 years)
- Annual and quarterly financial statements
- Key statistics (Short Ratio, Beta, Market Cap, etc.)

Outputs:
- data/chgg_market_data.csv
- data/chgg_income_statement.csv
- data/chgg_balance_sheet.csv
- data/chgg_cash_flow.csv
- data/chgg_key_stats.json
- results/data_inventory.txt
"""

import json
import os
from datetime import datetime, timedelta
from pathlib import Path

import pandas as pd
import yfinance as yf

# Set up paths
SESSION_DIR = Path("/app/sandbox/session_20260203_085912_030e6e80a419")
DATA_DIR = SESSION_DIR / "data"
RESULTS_DIR = SESSION_DIR / "results"

# Ensure directories exist
DATA_DIR.mkdir(exist_ok=True)
RESULTS_DIR.mkdir(exist_ok=True)

def fetch_market_data(ticker_symbol: str, years: int = 5) -> pd.DataFrame:
    """Fetch historical stock price and volume data."""
    print(f"\n[1/4] Fetching historical market data for {ticker_symbol}...")

    end_date = datetime.now()
    start_date = end_date - timedelta(days=years * 365)

    ticker = yf.Ticker(ticker_symbol)
    df = ticker.history(start=start_date.strftime('%Y-%m-%d'),
                        end=end_date.strftime('%Y-%m-%d'))

    if df.empty:
        print(f"  WARNING: No market data retrieved for {ticker_symbol}")
        return pd.DataFrame()

    # Reset index to make Date a column
    df = df.reset_index()

    print(f"  Retrieved {len(df)} trading days")
    print(f"  Date range: {df['Date'].min().strftime('%Y-%m-%d')} to {df['Date'].max().strftime('%Y-%m-%d')}")
    print(f"  Columns: {list(df.columns)}")

    return df

def fetch_financial_statements(ticker_symbol: str) -> dict:
    """Fetch annual and quarterly financial statements."""
    print(f"\n[2/4] Fetching financial statements for {ticker_symbol}...")

    ticker = yf.Ticker(ticker_symbol)

    financials = {}

    # Income Statement
    print("  Fetching Income Statement...")
    try:
        income_stmt = ticker.financials.T  # Annual
        income_stmt_q = ticker.quarterly_financials.T  # Quarterly

        if not income_stmt.empty:
            income_stmt.index.name = 'Date'
            income_stmt = income_stmt.reset_index()
            financials['income_statement_annual'] = income_stmt
            print(f"    Annual: {len(income_stmt)} periods")

        if not income_stmt_q.empty:
            income_stmt_q.index.name = 'Date'
            income_stmt_q = income_stmt_q.reset_index()
            financials['income_statement_quarterly'] = income_stmt_q
            print(f"    Quarterly: {len(income_stmt_q)} periods")
    except Exception as e:
        print(f"    WARNING: Could not fetch income statement: {e}")

    # Balance Sheet
    print("  Fetching Balance Sheet...")
    try:
        balance = ticker.balance_sheet.T  # Annual
        balance_q = ticker.quarterly_balance_sheet.T  # Quarterly

        if not balance.empty:
            balance.index.name = 'Date'
            balance = balance.reset_index()
            financials['balance_sheet_annual'] = balance
            print(f"    Annual: {len(balance)} periods")

        if not balance_q.empty:
            balance_q.index.name = 'Date'
            balance_q = balance_q.reset_index()
            financials['balance_sheet_quarterly'] = balance_q
            print(f"    Quarterly: {len(balance_q)} periods")
    except Exception as e:
        print(f"    WARNING: Could not fetch balance sheet: {e}")

    # Cash Flow
    print("  Fetching Cash Flow Statement...")
    try:
        cashflow = ticker.cashflow.T  # Annual
        cashflow_q = ticker.quarterly_cashflow.T  # Quarterly

        if not cashflow.empty:
            cashflow.index.name = 'Date'
            cashflow = cashflow.reset_index()
            financials['cash_flow_annual'] = cashflow
            print(f"    Annual: {len(cashflow)} periods")

        if not cashflow_q.empty:
            cashflow_q.index.name = 'Date'
            cashflow_q = cashflow_q.reset_index()
            financials['cash_flow_quarterly'] = cashflow_q
            print(f"    Quarterly: {len(cashflow_q)} periods")
    except Exception as e:
        print(f"    WARNING: Could not fetch cash flow: {e}")

    return financials

def fetch_key_statistics(ticker_symbol: str) -> dict:
    """Fetch key statistics and valuation metrics."""
    print(f"\n[3/4] Fetching key statistics for {ticker_symbol}...")

    ticker = yf.Ticker(ticker_symbol)
    info = ticker.info

    # Define the key metrics we want
    key_metrics = [
        'shortRatio', 'shortPercentOfFloat', 'sharesShort', 'shortPercentOfSharesShort',
        'beta', 'sharesOutstanding', 'floatShares',
        'marketCap', 'enterpriseValue',
        'trailingPE', 'forwardPE', 'priceToBook', 'priceToSalesTrailing12Months',
        'enterpriseToRevenue', 'enterpriseToEbitda',
        'profitMargins', 'returnOnAssets', 'returnOnEquity',
        'revenueGrowth', 'earningsGrowth', 'earningsQuarterlyGrowth',
        'totalRevenue', 'totalDebt', 'totalCash', 'freeCashflow',
        'operatingCashflow', 'bookValue', 'debtToEquity',
        'currentRatio', 'quickRatio',
        'fiftyTwoWeekHigh', 'fiftyTwoWeekLow', 'fiftyDayAverage', 'twoHundredDayAverage',
        'averageVolume', 'averageVolume10days',
        'dividendYield', 'payoutRatio',
        'sector', 'industry', 'fullTimeEmployees', 'country'
    ]

    stats = {}
    missing_metrics = []

    for metric in key_metrics:
        if metric in info and info[metric] is not None:
            stats[metric] = info[metric]
        else:
            missing_metrics.append(metric)

    # Add timestamp
    stats['_retrieved_at'] = datetime.now().isoformat()
    stats['_ticker'] = ticker_symbol

    print(f"  Retrieved {len(stats) - 2} metrics")  # -2 for metadata fields
    if missing_metrics:
        print(f"  Missing metrics: {', '.join(missing_metrics[:10])}" +
              (f" (+{len(missing_metrics)-10} more)" if len(missing_metrics) > 10 else ""))

    return stats, missing_metrics

def generate_data_inventory(market_data: pd.DataFrame, financials: dict,
                            stats: dict, missing_metrics: list) -> str:
    """Generate a text summary of the data inventory."""
    print(f"\n[4/4] Generating data inventory summary...")

    lines = [
        "=" * 60,
        "CHEGG (CHGG) DATA INVENTORY",
        f"Generated: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}",
        "=" * 60,
        "",
        "1. MARKET DATA",
        "-" * 40,
    ]

    if not market_data.empty:
        lines.extend([
            f"   File: data/chgg_market_data.csv",
            f"   Trading days: {len(market_data)}",
            f"   Date range: {market_data['Date'].min().strftime('%Y-%m-%d')} to {market_data['Date'].max().strftime('%Y-%m-%d')}",
            f"   Columns: {', '.join(market_data.columns)}",
            f"   Price range: ${market_data['Close'].min():.2f} - ${market_data['Close'].max():.2f}",
            "",
        ])
    else:
        lines.extend(["   WARNING: No market data retrieved", ""])

    lines.extend([
        "2. FINANCIAL STATEMENTS",
        "-" * 40,
    ])

    for key, df in financials.items():
        if not df.empty:
            min_date = df['Date'].min()
            max_date = df['Date'].max()
            if hasattr(min_date, 'strftime'):
                date_range = f"{min_date.strftime('%Y-%m-%d')} to {max_date.strftime('%Y-%m-%d')}"
            else:
                date_range = f"{min_date} to {max_date}"
            lines.append(f"   {key}: {len(df)} periods ({date_range})")

    lines.extend([
        "",
        "3. KEY STATISTICS",
        "-" * 40,
        f"   Total metrics retrieved: {len(stats) - 2}",  # -2 for metadata
    ])

    # Highlight critical metrics
    critical_metrics = ['shortRatio', 'beta', 'marketCap', 'enterpriseValue', 'sharesOutstanding']
    lines.append("   Critical metrics status:")
    for m in critical_metrics:
        status = "AVAILABLE" if m in stats else "MISSING"
        value = f" = {stats[m]}" if m in stats else ""
        lines.append(f"     - {m}: {status}{value}")

    lines.extend([
        "",
        "4. MISSING METRICS",
        "-" * 40,
    ])

    if missing_metrics:
        lines.append(f"   Count: {len(missing_metrics)}")
        lines.append(f"   List: {', '.join(missing_metrics)}")
    else:
        lines.append("   None - all requested metrics retrieved")

    lines.extend([
        "",
        "5. DATA FILES CREATED",
        "-" * 40,
        "   - data/chgg_market_data.csv",
        "   - data/chgg_income_statement.csv",
        "   - data/chgg_balance_sheet.csv",
        "   - data/chgg_cash_flow.csv",
        "   - data/chgg_key_stats.json",
        "   - results/data_inventory.txt",
        "",
        "=" * 60,
        "Data acquisition complete. Ready for downstream analysis.",
        "=" * 60,
    ])

    return "\n".join(lines)

def main():
    """Main execution function."""
    print("=" * 60)
    print("CHEGG (CHGG) DATA ACQUISITION")
    print("=" * 60)

    ticker_symbol = "CHGG"

    # 1. Fetch market data
    market_data = fetch_market_data(ticker_symbol, years=5)

    # 2. Fetch financial statements
    financials = fetch_financial_statements(ticker_symbol)

    # 3. Fetch key statistics
    stats, missing_metrics = fetch_key_statistics(ticker_symbol)

    # 4. Save all data
    print("\n[Saving data files...]")

    # Save market data
    if not market_data.empty:
        market_data_path = DATA_DIR / "chgg_market_data.csv"
        market_data.to_csv(market_data_path, index=False)
        print(f"  Saved: {market_data_path}")

    # Save financial statements
    for key, df in financials.items():
        if not df.empty:
            # Determine file name based on key
            if 'income' in key:
                file_suffix = 'income_statement'
            elif 'balance' in key:
                file_suffix = 'balance_sheet'
            elif 'cash' in key:
                file_suffix = 'cash_flow'
            else:
                file_suffix = key

            # Add annual/quarterly suffix
            if 'annual' in key:
                file_suffix += '_annual'
            elif 'quarterly' in key:
                file_suffix += '_quarterly'

            file_path = DATA_DIR / f"chgg_{file_suffix}.csv"
            df.to_csv(file_path, index=False)
            print(f"  Saved: {file_path}")

    # Save key statistics
    stats_path = DATA_DIR / "chgg_key_stats.json"
    with open(stats_path, 'w') as f:
        json.dump(stats, f, indent=2, default=str)
    print(f"  Saved: {stats_path}")

    # Generate and save data inventory
    inventory = generate_data_inventory(market_data, financials, stats, missing_metrics)
    inventory_path = RESULTS_DIR / "data_inventory.txt"
    with open(inventory_path, 'w') as f:
        f.write(inventory)
    print(f"  Saved: {inventory_path}")

    # Print inventory to console
    print("\n" + inventory)

    print("\n" + "=" * 60)
    print("DATA ACQUISITION COMPLETE")
    print("=" * 60)

if __name__ == "__main__":
    main()