data_input.py

# data_input.py
""" Helpers to keep transactions.csv and holdings.csv in sync"""

from __future__ import annotations

from pathlib import Path
from typing import Literal

import pandas as pd
from pandas.errors import EmptyDataError

# Base paths (match Static_Portfolio)
BASE_DIR = Path(__file__).resolve().parent
DATA_DIR = BASE_DIR / "data"
TX_PATH = DATA_DIR / "transactions.csv"
HOLDINGS_PATH = DATA_DIR / "holdings.csv"

TX_COLUMNS = ["Ticker", "Side", "Quantity", "Price", "Date"]
HOLDINGS_COLUMNS = [
    "Ticker",
    "Quantity",
    "Buy_Price",
    "Buy_Date",
    "Cost_Basis",
    "Realized_PL",
]


def _ensure_data_dir() -> None:
    """Ensure the data directory exists."""
    DATA_DIR.mkdir(parents=True, exist_ok=True)


def _read_csv(path: Path, columns: list[str]) -> pd.DataFrame:
    if path.exists():
        try:
            df = pd.read_csv(path)
        except EmptyDataError:
            # File exists but is empty – treat as brand-new file
            df = pd.DataFrame(columns=columns)

        # Make sure all expected cols exist
        for col in columns:
            if col not in df.columns:
                df[col] = 0 if col in ("Quantity", "Price", "Cost_Basis", "Realized_PL") else ""
        return df[columns]
    else:
        # No file yet – start with an empty DataFrame with correct columns
        return pd.DataFrame(columns=columns)


def append_transaction_row(
    ticker: str,
    side: Literal["BUY", "SELL"],
    quantity: float,
    price: float,
    date: str,
    tx_path: Path = TX_PATH,
) -> None:
    """
    Append a single BUY/SELL row into transactions.csv.
    """
    _ensure_data_dir()
    df = _read_csv(tx_path, TX_COLUMNS)

    new_row = {
        "Ticker": ticker.upper(),
        "Side": side.upper(),
        "Quantity": float(quantity),
        "Price": float(price),
        "Date": date,
    }
    df = pd.concat([df, pd.DataFrame([new_row])], ignore_index=True)
    df.to_csv(tx_path, index=False)


def update_holdings_from_transaction(
    ticker: str,
    side: Literal["BUY", "SELL"],
    quantity: float,
    price: float,
    date: str,
    holdings_path: Path = HOLDINGS_PATH,
) -> None:
    """Update holdings.csv to reflect a new BUY/SELL"""
    _ensure_data_dir()
    df = _read_csv(holdings_path, HOLDINGS_COLUMNS)

    ticker = ticker.upper()
    quantity = float(quantity)
    price = float(price)

    mask = df["Ticker"] == ticker
    if not mask.any():
        # No existing row for this ticker
        if side.upper() == "BUY":
            new_row = {
                "Ticker": ticker,
                "Quantity": quantity,
                "Buy_Price": price,
                "Buy_Date": date,
                "Cost_Basis": quantity * price,
                "Realized_PL": 0.0,
            }
            df = pd.concat([df, pd.DataFrame([new_row])], ignore_index=True)
        else:
            # Selling something you don't own — ignore to avoid crashing the web app
            df.to_csv(holdings_path, index=False)
            return
    else:
        # There is an existing position
        idx = df.index[mask][0]
        old_qty = float(df.at[idx, "Quantity"])
        old_price = float(df.at[idx, "Buy_Price"])

        if side.upper() == "BUY":
            new_qty = old_qty + quantity
            if new_qty <= 0:
                # All sold (or weird input) – drop row
                df = df.drop(index=idx)
            else:
                # Weighted average price
                new_buy_price = (old_qty * old_price + quantity * price) / new_qty
                df.at[idx, "Quantity"] = new_qty
                df.at[idx, "Buy_Price"] = new_buy_price
                df.at[idx, "Buy_Date"] = date
                df.at[idx, "Cost_Basis"] = new_qty * new_buy_price
                # Realized_PL untouched here
        else:  # SELL
            new_qty = old_qty - quantity
            if new_qty <= 0:
                df = df.drop(index=idx)
            else:
                df.at[idx, "Quantity"] = new_qty
                # Keep Buy_Price and Buy_Date as original
                df.at[idx, "Cost_Basis"] = new_qty * df.at[idx, "Buy_Price"]

    df.to_csv(holdings_path, index=False)


def record_new_position_from_form(form) -> None:
    """ Record a new BUY/SELL position from a web form.
    Expects a form with:
      - ticker
      - quantity
      - price
      - date
      - side (optional; defaults to BUY)

    updating BOTH transactions.csv and holdings.csv """
    ticker = (form.get("ticker") or "").strip().upper()
    quantity_raw = form.get("quantity", 0) or 0
    price_raw = form.get("price", 0) or 0
    date = (form.get("date") or "").strip()
    side = (form.get("side") or "BUY").upper()

    try:
        quantity = float(quantity_raw)
        price = float(price_raw)
    except (TypeError, ValueError):
        # Bad numeric input – ignore
        return

    if not ticker or quantity <= 0 or price <= 0 or not date:
        # Invalid input — ignore 
        return

    append_transaction_row(ticker, side, quantity, price, date)
    update_holdings_from_transaction(ticker, side, quantity, price, date)