main.py

"""
Stock Prediction API — FastAPI application entry point.

Run with:
    uvicorn main:app --reload --host 0.0.0.0 --port 8000
"""

import asyncio
import logging
import os
from contextlib import asynccontextmanager
from pathlib import Path

try:
    from dotenv import load_dotenv
    load_dotenv(Path(__file__).parent / ".env")
except ImportError:
    pass
from datetime import datetime, timezone
from pathlib import Path

from fastapi import FastAPI
from fastapi import HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse, JSONResponse
from fastapi.staticfiles import StaticFiles
from slowapi import _rate_limit_exceeded_handler
from slowapi.errors import RateLimitExceeded

from limiter import limiter
from routers.stock import router as stock_router
from routers.line_webhook import router as line_router
from routers.metal import router as metal_router
from routers.oil import router as oil_router
from routers.paper_trading import router as paper_trading_router

APP_VERSION = "1.1.0"

# ---------------------------------------------------------------------------
# Logging
# ---------------------------------------------------------------------------

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
)

# ---------------------------------------------------------------------------
# Startup warmup — pre-train popular stock models in background to avoid
# Render cold-start timeout (first request needs 45s+, Render limit is 30s).
# ---------------------------------------------------------------------------

_WARMUP_STOCKS = [
    code.strip()
    for code in os.getenv("WARMUP_STOCKS", "2330,0050").split(",")
    if code.strip()
]


async def _warmup_models() -> None:
    """Pre-train models for popular stocks at startup (background task)."""
    from services.predictor_service import get_prediction

    _logger = logging.getLogger("startup.warmup")
    if not _WARMUP_STOCKS:
        _logger.info("Startup warmup disabled")
        return

    _logger.info("Starting warmup for stocks: %s", _WARMUP_STOCKS)
    for code in _WARMUP_STOCKS:
        try:
            await get_prediction(code, raise_http=False)
            _logger.info("Warmup complete for %s", code)
        except Exception as exc:
            _logger.warning("Warmup failed for %s: %s", code, exc)
        # Small delay between stocks to avoid hammering data sources
        await asyncio.sleep(1)
    _logger.info("Warmup finished")


async def _maybe_setup_rich_menu():
    if os.environ.get("RICH_MENU_SETUP_NOW") != "1":
        return
    _log = logging.getLogger("startup.rich_menu")
    try:
        import subprocess
        result = subprocess.run(
            ["python", "setup_rich_menu.py"],
            capture_output=True, text=True, timeout=90,
            cwd=os.path.dirname(os.path.abspath(__file__)),
        )
        _log.warning("rich_menu stdout:\n%s", result.stdout)
        if result.stderr:
            _log.warning("rich_menu stderr:\n%s", result.stderr)
        _log.warning("rich_menu exit=%d — remove RICH_MENU_SETUP_NOW secret now", result.returncode)
    except Exception as e:
        _log.error("rich_menu setup failed: %s", e)


async def _warmup_from_queue():
    """On HF Space: pre-load all Mac 'done' results from queue.json into prediction cache."""
    from services.hf_queue_service import is_on_hf_space, read_queue
    from services.predictor_service import _set_cached_prediction, is_publishable_result

    if not is_on_hf_space():
        return
    try:
        queue = await asyncio.to_thread(read_queue)
        loaded = 0
        for stock_no, item in queue.items():
            if item.get("status") != "done" or not is_publishable_result(item.get("result")):
                continue
            result = dict(item["result"])
            result["optimized"] = True
            result["source"] = "mac_worker"
            _set_cached_prediction(stock_no, result)
            loaded += 1
        logging.getLogger(__name__).info("Queue warmup: loaded %d stocks from HF Dataset", loaded)
    except Exception as e:
        logging.getLogger(__name__).warning("Queue warmup failed: %s", e)


@asynccontextmanager
async def lifespan(app: FastAPI):
    """Lifespan handler: launch warmup background task at startup."""
    # Queue warmup is awaited — loads Mac results into cache BEFORE accepting requests
    await _warmup_from_queue()
    asyncio.create_task(_warmup_models())
    asyncio.create_task(_maybe_setup_rich_menu())
    yield


# ---------------------------------------------------------------------------
# App
# ---------------------------------------------------------------------------

app = FastAPI(
    title="Taiwan Stock Predictor API",
    description=(
        "Fetches OHLCV data from TWSE/TPEX official APIs, computes technical "
        "indicators, and uses a RandomForest ML model to generate directional "
        "signals. NOT financial advice."
    ),
    version=APP_VERSION,
    lifespan=lifespan,
)

# ---------------------------------------------------------------------------
# CORS — allow Vite dev server and any localhost origin
# ---------------------------------------------------------------------------

_extra = os.getenv("FRONTEND_URL", "")
_render_url = os.getenv("RENDER_EXTERNAL_URL", "")
ALLOWED_ORIGINS = [
    "http://localhost:5173",
    "http://localhost:5174",
    "http://127.0.0.1:5173",
    "http://127.0.0.1:5174",
    # Render production domain
    "https://stock-predictor-api.onrender.com",
]
if _extra:
    ALLOWED_ORIGINS.append(_extra)
if _render_url and _render_url not in ALLOWED_ORIGINS:
    ALLOWED_ORIGINS.append(_render_url)

app.state.limiter = limiter
app.add_exception_handler(RateLimitExceeded, _rate_limit_exceeded_handler)

app.add_middleware(
    CORSMiddleware,
    allow_origins=ALLOWED_ORIGINS,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# ---------------------------------------------------------------------------
# Routers (must be registered BEFORE static file catch-all)
# ---------------------------------------------------------------------------

app.include_router(stock_router)
app.include_router(line_router)
app.include_router(metal_router)
app.include_router(oil_router)
app.include_router(paper_trading_router)

# ---------------------------------------------------------------------------
# System endpoints
# ---------------------------------------------------------------------------

@app.get("/health", tags=["system"])
async def health():
    return {"status": "ok", "version": APP_VERSION, "timestamp": datetime.now(timezone.utc).isoformat()}


REPORTS_DIR = Path(__file__).parent / "docs" / "validation_runs"


@app.get("/reports/validation/{file_path:path}", tags=["system"])
async def validation_report(file_path: str):
    """Serve generated validation reports through the app server."""
    target = (REPORTS_DIR / file_path).resolve()
    try:
        target.relative_to(REPORTS_DIR.resolve())
    except ValueError:
        raise HTTPException(status_code=404, detail="Report not found")
    if not target.is_file():
        raise HTTPException(status_code=404, detail="Report not found")
    return FileResponse(str(target), headers={"Cache-Control": "no-cache"})

# ---------------------------------------------------------------------------
# API 404 handler — ensures /api/* typos return JSON 404 instead of being
# swallowed by the SPA catch-all route below.
# ---------------------------------------------------------------------------


@app.get("/api/{rest_of_path:path}", include_in_schema=False)
async def api_not_found(rest_of_path: str):
    """Return a proper 404 for any unmatched /api/* route."""
    return JSONResponse(
        status_code=404,
        content={"detail": f"API endpoint not found: /api/{rest_of_path}"},
    )


# ---------------------------------------------------------------------------
# Serve React frontend static files (only when dist/ folder exists)
# This is used in production (Docker / HuggingFace Spaces)
# ---------------------------------------------------------------------------

STATIC_DIR = Path(__file__).parent / "static"

if STATIC_DIR.exists():
    # Serve JS/CSS/image assets
    app.mount("/assets", StaticFiles(directory=str(STATIC_DIR / "assets")), name="assets")

    @app.get("/{full_path:path}", include_in_schema=False)
    async def serve_spa(full_path: str):
        """Catch-all: serve React index.html for all non-API routes (SPA routing)."""
        return FileResponse(
            str(STATIC_DIR / "index.html"),
            headers={"Cache-Control": "no-cache, no-store, must-revalidate"},
        )
else:
    @app.get("/", tags=["system"])
    async def root():
        return {
            "name": "Taiwan Stock Predictor API",
            "version": APP_VERSION,
            "docs": "/docs",
            "endpoints": {
                "history": "/api/stock/{stock_no}/history?months=6",
                "predict": "/api/stock/{stock_no}/predict",
                "info": "/api/stock/{stock_no}/info",
                "search": "/api/stock/search?q={query}",
            },
        }