utils.py

"""Utility helper functions used across the project.

This module centralizes small helpers such as writing generated code to file,
converting values to floats with defaults, and saving dataframe objects to CSV.
"""
from __future__ import annotations

import os
import tempfile
from datetime import datetime, timedelta
from typing import Any, Optional, Tuple

import gradio as gr

from fmp_python.fmp import FMP

__all__ = [
    "write_output",
    "_to_float_or_default",
    "save_df_to_csv",
    "resolve_market_to_ticker",
    "validate_date_range",
    "validate_ticker_symbol",
    "validate_fmp_key",
    "save_strategy_to_file",
    "yf_interval_info",
    "YF_INTERVALS",
    "YF_TO_FMP_MAP",
    "REPLAY_MAP",
    "REPLAY_INTERVALS",
    "update_replay_intervals",
]

def save_strategy_to_file(code_text: str):
    """Persist generated strategy code into a temp file Gradio can expose."""
    if not code_text:
        return None
    tmp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".py", prefix="strategy_")
    tmp_file.write(code_text.encode("utf-8"))
    tmp_file.flush()
    tmp_file.close()
    return tmp_file.name

def write_output(code: str, filename: str = "bt_strategy.py") -> None:
    """Write a string of code to a file.

    Args:
        code: The string content to write to disk.
        filename: The filename to use (defaults to `code.py`).
    """
    with open(filename, "w") as f:
        f.write(code)


def _to_float_or_default(value: Any, default: Any) -> float:
    """Convert a possibly-empty value to float or return a default.

    This mirrors behavior in the application code: if the input is None or an
    empty string it returns the provided default cast to float.
    """
    try:
        if value is None or value == "":
            raise ValueError
        return float(value)
    except (ValueError, TypeError):
        return float(default)


def resolve_market_to_ticker(market_name: str, market_to_ticker: dict) -> str:
    """Return a ticker symbol for a market name or ticker string.

    If market_name is a friendly name (in market_to_ticker), returns the mapped
    ticker symbol; otherwise assumes the input is a ticker symbol and returns
    the uppercased stripped version.
    """
    if market_name in list(market_to_ticker.keys()):
        return market_to_ticker[market_name]
    return market_name.upper().strip()


def validate_date_range(selected_start: str, selected_end: str, current_date_str: str) -> Tuple[str, str]:
    """Validate date strings in YYYY-MM-DD format and ensure start <= end <= today.

    Returns a tuple (start, end) as strings. Raises ValueError with a user-facing
    message on invalid input.
    """
    try:
        start_dt = datetime.strptime(selected_start, "%Y-%m-%d")
    except ValueError:
        raise ValueError(f"❌ Error: Start Date must be in format YYYY-MM-DD. Got '{selected_start}'.")
    try:
        end_dt = datetime.strptime(selected_end, "%Y-%m-%d")
    except ValueError:
        raise ValueError(f"❌Error: End Date must be in format YYYY-MM-DD. Got '{selected_end}'.")
    today_dt = datetime.strptime(current_date_str, "%Y-%m-%d")
    if start_dt > end_dt:
        raise ValueError("❌ Error: Start Date must be earlier than or equal to End Date.")
    if end_dt > today_dt:
        raise ValueError(f"❌ Error: End Date cannot be in the future. Today is {current_date_str}.")
    return selected_start, selected_end


def validate_ticker_symbol(tckr_symbl: str, data_source: str = "yahoofinance") -> str:
    """Validate a ticker symbol using either FMP or yfinance based on the data source.

    Raises ValueError on invalid tickers or if the provider fails to respond.
    """
    source = (data_source or "").lower()
    if source == "fmp":
        try:
            fmp = FMP()
            result = fmp.get_quote(tckr_symbl)
            if not result:
               raise ValueError(f"❌ Error: Ticker {tckr_symbl} not found.")
            else:
                # If there's a valid quote, the ticker exists
                print(f"{tckr_symbl} exists! Data:", result)
        except Exception as e:
            raise ValueError(f"❌ Error: Invalid ticker symbol '{tckr_symbl}' for FMP. {e}")
    try:
        import yfinance as yf

        ticker_data = yf.Ticker(tckr_symbl)
        info = ticker_data.info
        if not info or "symbol" not in info:
            raise ValueError(f"❌ Error: Invalid ticker symbol '{tckr_symbl}'.")
        return info.get("longName") or info.get("shortName") or tckr_symbl
    except ValueError:
        raise
    except Exception:
        raise ValueError(f"❌ Error: Unable to validate ticker symbol '{tckr_symbl}'.")


def validate_fmp_key() -> tuple[bool, str]:
    """Validate FMP API key. Returns (is_valid, error_message)."""
    # Check if key exists
    api_key = os.getenv('FMP_API_KEY')
    if not api_key:
        return False, "❌ FMP_API_KEY environment variable not set. Please add it to your .env file."

    # Test with lightweight API call
    try:
        fmp = FMP()
        result = fmp.get_quote('SPY')
        if result:
            return True, ""
        else:
            return False, "❌ FMP API key is invalid. Please check your FMP_API_KEY in .env file."
    except Exception as e:
        return False, f"❌ FMP API validation failed: {str(e)}. Please check your FMP_API_KEY."

YF_INTERVALS = ["1m", "5m", "15m", "30m", "1h", "1d"]

# Map YF intervals to FMP intervals (None = not supported)
YF_TO_FMP_MAP = {
    "1m": "1min",
    "5m": "5min",
    "15m": "15min",
    "30m": None,
    "1h": "1hour",
    "1d": None,
}

# Replay mapping: key -> (fmp_interval, yf_interval, compression)
REPLAY_MAP = {
    "1min -> 5min": ("1min", "1m", 5),
    "1min -> 15min": ("1min", "1m", 15),
    "1min -> 30min": ("1min", "1m", 30),
    "5min -> 15min": ("5min", "5m", 15),
    "5min -> 30min": ("5min", "5m", 30),
}
REPLAY_INTERVALS = list(REPLAY_MAP.keys())


def update_replay_intervals(replay_enabled: bool):
    """Update interval dropdown based on replay checkbox."""
    if replay_enabled:
        return gr.update(choices=REPLAY_INTERVALS, value=REPLAY_INTERVALS[0])
    else:
        return gr.update(choices=YF_INTERVALS, value="1d")


def yf_interval_info(date_range: dict, interval: str, auto_period: bool) -> str:
    """Describe the effective yfinance download window for intraday intervals."""
    intraday = ["1m", "2m", "5m", "15m", "30m", "60m", "1h"]
    if not auto_period or interval not in intraday:
        return ""
    end_dt = datetime.strptime(date_range["end"], "%Y-%m-%d")
    start_dt = datetime.strptime(date_range["start"], "%Y-%m-%d")
    if interval == "1m":
        max_days = 7
    elif interval in ["2m", "5m", "15m", "30m"]:
        max_days = 60
    else:  # 60m/1h
        max_days = 730
    desired_days = max(1, (end_dt - start_dt).days or 1)
    used_days = min(desired_days, max_days)
    effective_start = (end_dt - timedelta(days=used_days - 1)).strftime("%Y-%m-%d")
    return f"YFinance download window: {effective_start} → {date_range['end']} ({used_days}d) @ {interval}"


def save_df_to_csv(df: Optional[Any], filename: str) -> Optional[str]:
    """Save a pandas DataFrame (or df-like object) to CSV in the system temp dir.

    Returns the full path to the saved file or None when the input could not be
    converted into a non-empty DataFrame.
    """
    if df is None:
        return None

    try:
        import pandas as pd
    except Exception:
        # pandas is required for this helper; if it's not installed return None to
        # fail gracefully so callers can handle the absence (e.g. in limited
        # environments where CSV exports aren't available).
        return None

    if isinstance(df, pd.DataFrame):
        df_to_save = df
    else:
        try:
            df_to_save = pd.DataFrame(df)
        except Exception:
            return None

    if df_to_save.empty:
        return None

    out_dir = tempfile.gettempdir()
    os.makedirs(out_dir, exist_ok=True)
    filepath = os.path.join(out_dir, f"{filename}.csv")
    df_to_save.to_csv(filepath, index=False)
    return filepath