data/ptt_sentiment.py

"""
PTT Stock board sentiment scraper.
Produces daily sentiment score per stock from [標的] posts.
Score = sum of (推 - 噓) for posts mentioning a stock on that date.
"""
import json
import re
import time
from datetime import datetime, timedelta
from pathlib import Path

import pandas as pd
import requests

CACHE_DIR = Path.home() / ".cache" / "stock_predictor"
CACHE_FILE = CACHE_DIR / "ptt_sentiment.json"
CACHE_TTL_HOURS = 6
PTT_BASE = "https://www.ptt.cc"
SESSION_COOKIES = {"over18": "1"}
HEADERS = {"User-Agent": "Mozilla/5.0 (compatible; stock-predictor/1.0)"}

# Regex for TWSE 4-digit stock codes (avoid matching years like 2026)
STOCK_RE = re.compile(r'\b([2-9]\d{3}|[0-1]\d{3})\b')
KNOWN_STOCKS = {  # common codes to disambiguate
    "2330", "2317", "2454", "2881", "2882", "2886", "2891",
    "0050", "0056", "2303", "2308", "2357", "2382", "2412",
    "3008", "3034", "3045", "4904", "4938", "6505", "6669",
}


def _parse_nrec(nrec_text: str) -> int:
    """Parse PTT nrec field: '爆'→100, 'X{n}'→negative, digits→int."""
    t = nrec_text.strip()
    if t == "爆":
        return 100
    if t.startswith("X"):
        try:
            return -int(t[1:])
        except ValueError:
            return -10
    try:
        return int(t)
    except ValueError:
        return 0


def _extract_stock_codes(title: str) -> list[str]:
    """Extract TWSE-style 4-digit codes from post title."""
    candidates = STOCK_RE.findall(title)
    # Prefer known codes; if none found, return all 4-digit candidates
    known = [c for c in candidates if c in KNOWN_STOCKS]
    return known if known else candidates[:3]  # cap at 3 to avoid noise


def _fetch_page(url: str, session: requests.Session) -> tuple[list[dict], str | None]:
    """
    Fetch one PTT index page. Returns (posts_list, prev_page_url).
    Each post dict: {title, author, date_str, href, nrec, stocks}.
    """
    try:
        resp = session.get(url, headers=HEADERS, timeout=10)
        resp.raise_for_status()
    except Exception:
        return [], None

    from html.parser import HTMLParser

    class _Parser(HTMLParser):
        def __init__(self):
            super().__init__()
            self.posts = []
            self.prev_url = None
            self._in_title = self._in_author = self._in_date = self._in_nrec = False
            self._cur = {}

        def handle_starttag(self, tag, attrs):
            a = dict(attrs)
            cls = a.get("class", "")
            if tag == "div" and "r-ent" in cls:
                self._cur = {}
            if tag == "div" and cls == "title":
                self._in_title = True
            if tag == "div" and cls == "author":
                self._in_author = True
            if tag == "div" and cls == "date":
                self._in_date = True
            if tag == "div" and cls == "nrec":
                self._in_nrec = True
            if tag == "a" and self._in_title and "href" in a:
                self._cur["href"] = a["href"]
            # Prev page button
            if tag == "a" and a.get("class") == "btn wide" and "上頁" in (a.get("title") or ""):
                self.prev_url = a.get("href")
            # Find btn wide links with text containing 上頁 via data
            if tag == "a" and "btn wide" in cls:
                self._last_btn_href = a.get("href")

        def handle_data(self, data):
            if self._in_title:
                self._cur["title"] = self._cur.get("title", "") + data
            if self._in_author:
                self._cur["author"] = data.strip()
                self._in_author = False
            if self._in_date:
                self._cur["date_str"] = data.strip()
                self._in_date = False
            if self._in_nrec:
                self._cur["nrec_text"] = data.strip()
                self._in_nrec = False
            if data.strip() == "上頁" and hasattr(self, "_last_btn_href"):
                self.prev_url = self._last_btn_href

        def handle_endtag(self, tag):
            if tag == "div" and self._in_title:
                self._in_title = False
                if self._cur.get("href") and self._cur.get("title"):
                    self.posts.append(dict(self._cur))

    parser = _Parser()
    parser.feed(resp.text)
    # Fallback: find prev page via regex
    if not parser.prev_url:
        m = re.search(r'href="(/bbs/Stock/index\d+\.html)"[^>]*>.*?上頁', resp.text, re.S)
        if m:
            parser.prev_url = m.group(1)

    posts = []
    for p in parser.posts:
        title = p.get("title", "").strip()
        if "[標的]" not in title:
            continue
        nrec = _parse_nrec(p.get("nrec_text", "0"))
        stocks = _extract_stock_codes(title)
        date_str = p.get("date_str", "").strip()
        posts.append({
            "title":    title,
            "href":     p.get("href", ""),
            "nrec":     nrec,
            "stocks":   stocks,
            "date_str": date_str,
        })
    return posts, parser.prev_url


def _date_from_str(date_str: str, year: int) -> datetime | None:
    """Parse PTT date like ' 5/11' → datetime(year, 5, 11)."""
    try:
        clean = date_str.strip()
        # PTT format: ' 5/11' or '5/11'
        parts = clean.split("/")
        if len(parts) == 2:
            month, day = int(parts[0].strip()), int(parts[1].strip())
            return datetime(year, month, day)
    except Exception:
        pass
    return None


def scrape_ptt_sentiment(days_back: int = 180) -> pd.DataFrame:
    """
    Scrape PTT [標的] posts going back `days_back` days.
    Returns DataFrame with columns: date (date), stock_no (str), sentiment_score (int).
    """
    session = requests.Session()
    session.cookies.update(SESSION_COOKIES)

    cutoff = datetime.now() - timedelta(days=days_back)
    current_year = datetime.now().year

    all_posts = []
    url = f"{PTT_BASE}/bbs/Stock/index.html"
    pages_fetched = 0
    max_pages = 300  # safety cap

    while url and pages_fetched < max_pages:
        posts, prev_url = _fetch_page(f"{PTT_BASE}{url}" if url.startswith("/") else url, session)
        pages_fetched += 1

        stop = False
        for p in posts:
            dt = _date_from_str(p["date_str"], current_year)
            if dt is None:
                # Try previous year (December posts when we're in January)
                dt = _date_from_str(p["date_str"], current_year - 1)
            if dt is None:
                continue
            if dt < cutoff:
                stop = True
                break
            for stock in p["stocks"]:
                all_posts.append({
                    "date":            dt.date(),
                    "stock_no":        stock,
                    "sentiment_score": p["nrec"],
                })

        time.sleep(0.3)
        if stop or not prev_url:
            break
        url = prev_url

    if not all_posts:
        return pd.DataFrame(columns=["date", "stock_no", "sentiment_score"])

    df = pd.DataFrame(all_posts)
    df["date"] = pd.to_datetime(df["date"])
    # Aggregate: sum scores per (date, stock_no)
    agg = df.groupby(["date", "stock_no"])["sentiment_score"].sum().reset_index()
    return agg


def load_ptt_sentiment(days_back: int = 180, force_refresh: bool = False) -> pd.DataFrame:
    """Load from cache or scrape fresh."""
    CACHE_DIR.mkdir(parents=True, exist_ok=True)
    if not force_refresh and CACHE_FILE.exists():
        age_hours = (time.time() - CACHE_FILE.stat().st_mtime) / 3600
        if age_hours < CACHE_TTL_HOURS:
            try:
                data = json.loads(CACHE_FILE.read_text())
                df = pd.DataFrame(data)
                df["date"] = pd.to_datetime(df["date"])
                return df
            except Exception:
                pass
    df = scrape_ptt_sentiment(days_back=days_back)
    if not df.empty:
        CACHE_FILE.write_text(json.dumps(df.assign(date=df["date"].dt.strftime("%Y-%m-%d")).to_dict("records")))
    return df


def add_ptt_sentiment(df: pd.DataFrame, stock_no: str) -> pd.DataFrame:
    """
    Add ptt_sentiment_1d and ptt_sentiment_5d_ma to df.
    Uses merge_asof(direction="backward") to avoid lookahead.
    Missing dates → 0.0.
    """
    ptt = load_ptt_sentiment()
    if ptt.empty:
        df["ptt_sentiment_1d"]  = 0.0
        df["ptt_sentiment_5d_ma"] = 0.0
        return df

    stock_ptt = ptt[ptt["stock_no"] == str(stock_no)].copy()
    if stock_ptt.empty:
        df["ptt_sentiment_1d"]  = 0.0
        df["ptt_sentiment_5d_ma"] = 0.0
        return df

    stock_ptt = stock_ptt.sort_values("date").rename(
        columns={"sentiment_score": "ptt_sentiment_1d"}
    )
    # 5-day rolling mean on PTT side
    stock_ptt["ptt_sentiment_5d_ma"] = (
        stock_ptt["ptt_sentiment_1d"].rolling(5, min_periods=1).mean()
    )

    # Align to df by date — use merge_asof so we never look ahead
    if "date" in df.columns:
        df_dates = pd.to_datetime(df["date"])
    else:
        df_dates = pd.to_datetime(df.index)

    df_tmp = pd.DataFrame({"date": df_dates})
    merged = pd.merge_asof(
        df_tmp.sort_values("date"),
        stock_ptt[["date", "ptt_sentiment_1d", "ptt_sentiment_5d_ma"]].sort_values("date"),
        on="date", direction="backward",
    )
    # Reindex back to original df order
    df = df.copy()
    df["ptt_sentiment_1d"]   = merged["ptt_sentiment_1d"].fillna(0.0).values
    df["ptt_sentiment_5d_ma"] = merged["ptt_sentiment_5d_ma"].fillna(0.0).values
    return df