feature_pipeline.py

from __future__ import annotations

import logging
import re
import threading
from datetime import datetime
from dataclasses import dataclass
from typing import Optional, Tuple
from urllib.parse import quote

import numpy as np
import pandas as pd
import requests
from bs4 import BeautifulSoup
from requests.adapters import HTTPAdapter
from urllib3.util.retry import Retry

logger = logging.getLogger(__name__)

def _browser_headers() -> dict:
    return {
        "User-Agent": (
            "Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
            "AppleWebKit/537.36 (KHTML, like Gecko) "
            "Chrome/125.0.0.0 Safari/537.36"
        ),
        "Accept": (
            "text/html,application/xhtml+xml,application/xml;"
            "q=0.9,image/avif,image/webp,image/apng,*/*;q=0.8"
        ),
        "Accept-Language": "en-US,en;q=0.9",
        "Accept-Encoding": "gzip, deflate, br",
        "Connection": "keep-alive",
        "Cache-Control": "no-cache",
        "Pragma": "no-cache",
        "Referer": "https://www.screener.in/",
        "Upgrade-Insecure-Requests": "1",
        "Sec-Fetch-Dest": "document",
        "Sec-Fetch-Mode": "navigate",
        "Sec-Fetch-Site": "none",
        "Sec-Fetch-User": "?1",
    }

def _clean_number(text: str) -> float:
    return float(text.replace(",", "").strip())

@dataclass(frozen=True)
class CapThresholds:
    large: float = 20000.0
    mid: float = 5000.0

class _ThreadLocalSessionFactory:
    def __init__(
        self,
        headers: Optional[dict] = None,
        timeout: Tuple[float, float] = (5.0, 15.0),
        pool_maxsize: int = 32,
        prime_homepage: bool = True,
    ):
        self._local = threading.local()
        self.headers = headers or _browser_headers()
        self.timeout = timeout
        self.pool_maxsize = pool_maxsize
        self.prime_homepage = prime_homepage

    def _build_session(self) -> requests.Session:
        session = requests.Session()
        session.headers.update(self.headers)

        retry = Retry(
            total=3,
            connect=3,
            read=3,
            backoff_factor=0.35,
            status_forcelist=(429, 500, 502, 503, 504),
            allowed_methods=frozenset(["GET"]),
            raise_on_status=False,
        )

        adapter = HTTPAdapter(
            max_retries=retry,
            pool_connections=self.pool_maxsize,
            pool_maxsize=self.pool_maxsize,
        )

        session.mount("https://", adapter)
        session.mount("http://", adapter)
        return session

    def get(self) -> requests.Session:
        if not hasattr(self._local, "session"):
            self._local.session = self._build_session()

            if self.prime_homepage:
                try:
                    self._local.session.get(
                        "https://www.screener.in/",
                        timeout=self.timeout,
                    )
                except requests.RequestException:
                    pass

        return self._local.session

class ScreenerScraper:
    def __init__(
        self,
        timeout: Tuple[float, float] = (5.0, 15.0),
        pool_maxsize: int = 32,
        thresholds: CapThresholds = CapThresholds(),
    ):
        self._session_factory = _ThreadLocalSessionFactory(
            timeout=timeout,
            pool_maxsize=pool_maxsize,
            prime_homepage=True,
        )
        self.timeout = timeout
        self.thresholds = thresholds

    def _session(self) -> requests.Session:
        return self._session_factory.get()

    @staticmethod
    def _normalize_symbol(symbol: str) -> str:
        return quote(symbol.strip().upper(), safe="-._~")

    def _fetch(self, url: str) -> str:
        session = self._session()
        response = session.get(url, timeout=self.timeout)

        if response.status_code == 200:
            return response.text

        preview = ""
        try:
            preview = response.text[:250].replace("\n", " ").replace("\r", " ")
        except Exception:
            preview = "<no preview>"

        raise RuntimeError(
            f"Request failed\n"
            f"URL: {url}\n"
            f"Status: {response.status_code}\n"
            f"Reason: {response.reason}\n"
            f"Preview: {preview}"
        )

    @staticmethod
    def _parse_market_cap_crore(html: str) -> Optional[float]:
        patterns = [
            r"Mkt Cap:\s*([0-9][0-9,]*(?:\.[0-9]+)?)\s*Crore",
            r"Market Cap\s*₹\s*([0-9][0-9,]*(?:\.[0-9]+)?)\s*Cr\.?",
            r"Market Cap:\s*₹\s*([0-9][0-9,]*(?:\.[0-9]+)?)\s*Cr\.?",
            r"Market Capitalization.*?₹\s*([0-9][0-9,]*(?:\.[0-9]+)?)\s*Cr\.?",
        ]

        for pattern in patterns:
            match = re.search(pattern, html, flags=re.IGNORECASE | re.DOTALL)
            if match:
                try:
                    return _clean_number(match.group(1))
                except Exception:
                    continue

        return None

    def _fetch_screener_html(self, symbol: str, consolidated: bool = True) -> str:
        s = self._normalize_symbol(symbol)
        urls = (
            [
                f"https://www.screener.in/company/{s}/consolidated/",
                f"https://www.screener.in/company/{s}/",
            ]
            if consolidated
            else [
                f"https://www.screener.in/company/{s}/",
                f"https://www.screener.in/company/{s}/consolidated/",
            ]
        )

        last_error: Optional[Exception] = None
        for url in urls:
            try:
                html = self._fetch(url)
                return html
            except Exception as e:
                last_error = e

        raise RuntimeError(f"Failed to fetch Screener page for {symbol}. Last error: {last_error}") from last_error

    def _fetch_nse_quote_html(self, symbol: str) -> str:
        s = self._normalize_symbol(symbol)
        url = f"https://www.nseindia.com/get-quotes/equity?symbol={s}"
        session = self._session()
        try:
            session.get("https://www.nseindia.com/", timeout=self.timeout)
        except requests.RequestException:
            pass
        return self._fetch(url)

    def get_market_cap_crore(self, symbol: str, consolidated: bool = True) -> Optional[float]:
        try:
            html = self._fetch_screener_html(symbol, consolidated=consolidated)
            cap = self._parse_market_cap_crore(html)
            if cap is not None:
                return cap
        except Exception as e:
            logger.warning("Screener lookup failed for %s: %s", symbol, e)

        try:
            html = self._fetch_nse_quote_html(symbol)
            cap = self._parse_market_cap_crore(html)
            if cap is not None:
                return cap
        except Exception as e:
            logger.warning("NSE lookup failed for %s: %s", symbol, e)
        return None

    def classify_market_cap(self, market_cap_crore: Optional[float]) -> str:
        if market_cap_crore is None:
            return "Unknown"
        if market_cap_crore >= self.thresholds.large:
            return "Large Cap"
        if market_cap_crore >= self.thresholds.mid:
            return "Mid Cap"
        return "Small Cap"

    def get_cap_info(self, symbol: str, consolidated: bool = True) -> dict:
        cap = self.get_market_cap_crore(symbol, consolidated=consolidated)
        return {
            "symbol": symbol.upper().strip(),
            "market_cap_crore": cap,
            "market_cap_class": self.classify_market_cap(cap),
        }

    # -- Keep original helper methods for get_stock_info --
    def _fetch_html(self, ticker: str, consolidated: bool = True) -> str:
        return self._fetch_screener_html(ticker, consolidated)

    @staticmethod
    def _make_soup(html: str):
        try:
            return BeautifulSoup(html, "lxml")
        except Exception:
            return BeautifulSoup(html, "html.parser")
    @staticmethod
    def _clean_text(value: str) -> str:
        return " ".join(value.split()).replace("₹", "Rs.")

    def get_stock_info(self, ticker):
        html = self._fetch_html(ticker, consolidated=True)
        soup = self._make_soup(html)

        data = {
            "ticker": ticker.upper(),
            "key_metrics": {},
            "history": {},
            "documents": {},
        }

        top_ratios = soup.find("ul", id="top-ratios")
        if top_ratios:
            for li in top_ratios.find_all("li"):
                n_span = li.find("span", class_="name")
                v_span = li.find("span", class_="value")
                if n_span and v_span:
                    name = n_span.get_text(strip=True).replace("₹", "Rs.")
                    val = self._clean_text(v_span.get_text(" ", strip=True))
                    data["key_metrics"][name] = val

        sections = {
            "quarters": "Quarterly Results",
            "profit-loss": "Profit & Loss",
            "balance-sheet": "Balance Sheet",
            "cash-flow": "Cash Flows",
            "ratios": "Financial Ratios",
        }

        for sec_id, sec_name in sections.items():
            sec = soup.find("section", id=sec_id)
            if not sec:
                continue
            tbl = sec.find("table")
            if not tbl:
                continue

            thead = tbl.find("thead")
            headers_list = [th.get_text(" ", strip=True) for th in thead.find_all("th")] if thead else []
            tbody = tbl.find("tbody")
            if not tbody:
                continue

            rows = []
            for tr in tbody.find_all("tr"):
                tds = tr.find_all("td")
                if not tds:
                    continue
                cols = [td.get_text(" ", strip=True) for td in tds]
                rname = tr.find("td", class_="text")
                if rname and cols:
                    cols[0] = rname.get_text(" ", strip=True).replace("+", "").strip()
                rows.append(cols)

            if not rows:
                continue

            if headers_list:
                if len(headers_list) == len(rows[0]) - 1:
                    headers_list = ["Metric"] + headers_list
                elif len(headers_list) == len(rows[0]):
                    headers_list[0] = "Metric"
                else:
                    headers_list = ["Metric"] + headers_list[1:]

            data["history"][sec_id] = {
                "title": sec_name,
                "headers": headers_list,
                "rows": rows,
            }

        docs = soup.find_all("div", class_="documents")
        for doc in docs:
            h3 = doc.find("h3")
            if not h3:
                continue
            sec_name = h3.get_text(" ", strip=True)
            data["documents"][sec_name] = []
            ul = doc.find("ul", class_="list-links")
            if not ul:
                continue
            for li in ul.find_all("li"):
                text_div = li.find("div")
                a = li.find("a")
                date_str = " ".join(text_div.get_text(" ", strip=True).replace("\n", " ").split()) if text_div else ""
                link_str = a.get_text(" ", strip=True) if a else ""
                if link_str and date_str.endswith(link_str):
                    date_str = date_str[:-len(link_str)].strip()
                data["documents"][sec_name].append({"date": date_str, "title": link_str})

        return data


class FeatureEngineer:
    @staticmethod
    def clean_numeric(val):
        if pd.isna(val):
            return np.nan
        if isinstance(val, (int, float, np.number)):
            return float(val)
        if not isinstance(val, str):
            return val
        val = val.replace(",", "").replace("%", "").replace("₹", "").replace("Rs.", "").strip()
        if val in ["-", ""]:
            return np.nan
        try:
            return float(val)
        except ValueError:
            return val

    @staticmethod
    def parse_date(date_str):
        try:
            return pd.to_datetime(date_str, format="%b %Y")
        except Exception:
            return pd.to_datetime(date_str, errors="coerce")

    @staticmethod
    def _section_to_frame(sec_id, sec_data):
        headers = sec_data.get("headers") or []
        rows = sec_data.get("rows") or []
        if len(headers) < 2 or not rows:
            return None

        df = pd.DataFrame(rows, columns=headers[: len(rows[0])])
        if "Metric" not in df.columns:
            df.columns = ["Metric"] + list(df.columns[1:])

        df = df.set_index("Metric").transpose().reset_index().rename(columns={"index": "Date_Str"})
        metric_cols = [c for c in df.columns if c != "Date_Str"]

        if metric_cols:
            cleaned = df[metric_cols].replace(
                {",": "", "%": "", "₹": "", "Rs.": ""},
                regex=True,
            )
            df[metric_cols] = cleaned.apply(pd.to_numeric, errors="coerce")

        df["Date"] = df["Date_Str"].map(FeatureEngineer.parse_date)
        df = df.dropna(subset=["Date"]).sort_values("Date").set_index("Date")
        df = df.drop(columns=["Date_Str"])

        prefix = sec_id.split("-")[0].upper() + "_"
        df.columns = [f"{prefix}{col}" for col in df.columns]
        return df

    def build_features(self, data):
        ticker = data["ticker"]

        df_dict = {}
        for sec_id, sec_data in data["history"].items():
            df = self._section_to_frame(sec_id, sec_data)
            if df is not None and not df.empty:
                df_dict[sec_id] = df

        if "quarters" not in df_dict or df_dict["quarters"].empty:
            print("No quarterly data found.")
            return None

        main_df = df_dict["quarters"].copy()

        for sec_id in ["profit-loss", "balance-sheet", "cash-flow", "ratios"]:
            if sec_id in df_dict and not df_dict[sec_id].empty:
                annual_df = df_dict[sec_id].sort_index()
                main_df = pd.merge_asof(
                    main_df.sort_index(),
                    annual_df,
                    left_index=True,
                    right_index=True,
                    direction="backward",
                )

        announcements = data.get("documents", {}).get("Announcements", [])
        if announcements:
            event_dates = []
            for ann in announcements:
                ds = ann.get("date", "")
                if "20" in ds:
                    parsed = pd.to_datetime(ds, errors="coerce")
                    if pd.notnull(parsed):
                        event_dates.append(parsed)
            if event_dates:
                events_series = pd.Series(1, index=pd.DatetimeIndex(event_dates))
                quarterly_counts = events_series.resample("Q").sum()
                main_df["ANN_COUNT"] = 0
                for q_date, count in quarterly_counts.items():
                    valid_idx = main_df.index[main_df.index <= q_date]
                    if len(valid_idx) > 0:
                        main_df.loc[valid_idx[-1], "ANN_COUNT"] += int(count)

        for k, v in data.get("key_metrics", {}).items():
            main_df[f"STATIC_{k.replace(' ', '_')}"] = self.clean_numeric(v)

        main_df["Ticker"] = ticker

        if "QUARTERS_Sales" in main_df.columns:
            main_df["QUARTERS_Sales_YoY_Growth"] = main_df["QUARTERS_Sales"].pct_change(periods=4) * 100

        return main_df


def run_pipeline(ticker):
    print(f"[{ticker}] Scraping data...")
    data = ScreenerScraper().get_stock_info(ticker)

    print(f"[{ticker}] Engineering features...")
    df = FeatureEngineer().build_features(data)

    if df is not None:
        filename = f"{ticker}_features.parquet"
        print(f"[{ticker}] Exporting to {filename}...")
        df.to_parquet(filename, engine="pyarrow", index=True)
        print("Success!")
        return filename

    print("Failed to build features.")
    return None


if __name__ == "__main__":
    import sys

    ticker = sys.argv[1] if len(sys.argv) > 1 else "TCS"
    run_pipeline(ticker)