app.py

import os
import json
from dotenv import load_dotenv
load_dotenv()
from nova_prompt import NOVA_SYSTEM_PROMPT_BASE, NOVA_SYSTEM_PROMPT_MASTER, NOVA_SYSTEM_PROMPT_USER, GENERATIVE_SYSTEM_PROMPT, ORACLE_SYSTEM_PROMPT# Set thread limits to prevent OpenBLAS/MKL deadlock inside FastAPI threads
os.environ["OMP_NUM_THREADS"] = "1"
os.environ["OPENBLAS_NUM_THREADS"] = "1"
os.environ["MKL_NUM_THREADS"] = "1"
os.environ["VECLIB_MAXIMUM_THREADS"] = "1"
os.environ["NUMEXPR_NUM_THREADS"] = "1"

import faulthandler
faulthandler.enable()

from concurrent.futures import ProcessPoolExecutor

from fastapi import FastAPI, HTTPException, Request, Header, Depends
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse, JSONResponse, StreamingResponse
from pydantic import BaseModel
from typing import List, Optional
import secrets
from sqlalchemy.orm import Session
import database
from database import get_pg_engine, SavedPortfolio, BacktestHistory, WebhookConfig, UserMemory
import time
import threading
import uuid
import yfinance as yf
from datetime import datetime
import traceback
import pandas as pd
import logging
import cache
redis_client = cache.get_redis()

try:
    from diagnostics import TraceManager
    tracer = TraceManager()
except ImportError:
    tracer = None

try:
    from huggingface_hub import InferenceClient
    has_hf_hub = True
except ImportError:
    has_hf_hub = False

try:
    import core_engine
except ImportError:
    core_engine = None

from config import OUTPUT_DIR, logger
import access_manager

class RedisTaskStore:
    def __init__(self, prefix="task:"):
        self.prefix = prefix
    
    def get(self, key, default=None):
        try:
            data = cache.redis_client.get(f"{self.prefix}{key}")
            if data:
                return json.loads(data)
        except Exception as e:
            logger.error(f"Redis get error: {e}")
        return default

    def __setitem__(self, key, value):
        try:
            cache.redis_client.setex(f"{self.prefix}{key}", 86400, json.dumps(value))
        except Exception as e:
            logger.error(f"Redis set error: {e}")

    def __contains__(self, key):
        return bool(cache.redis_client.exists(f"{self.prefix}{key}"))

BACKGROUND_TASKS = RedisTaskStore()

def update_task(tid, **kwargs):
    task = BACKGROUND_TASKS.get(tid, {})
    task.update(kwargs)
    BACKGROUND_TASKS[tid] = task

process_pool = ProcessPoolExecutor(max_workers=4)

BASE_DIR = os.path.dirname(os.path.abspath(__file__))
STATIC_DIR = os.path.join(BASE_DIR, "static")

app = FastAPI(title="Portfolio Engine API")

app.mount("/static", StaticFiles(directory=STATIC_DIR), name="static")

from ai_chat import chat_router, memory_router
from options import options_router

app.include_router(chat_router)
app.include_router(memory_router)
app.include_router(options_router)
def get_db():
    from sqlalchemy.orm import sessionmaker
    engine = get_pg_engine()
    SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()

from constants import MASTER_KEY

def get_current_user(x_access_key: Optional[str] = Header(None), x_username: Optional[str] = Header(None)):
    if not x_access_key:
        raise HTTPException(status_code=401, detail="Access Key missing")
    if not x_username and x_access_key != MASTER_KEY:
        raise HTTPException(status_code=401, detail="Username missing")
    if not access_manager.validate_key(x_access_key, username=x_username):
        raise HTTPException(status_code=401, detail="Invalid Key or Username")
    return x_username if x_username else "admin"

@app.get("/api/portfolios")
def get_portfolios(username: str = Depends(get_current_user), db: Session = Depends(get_db)):
    portfolios = db.query(SavedPortfolio).filter(SavedPortfolio.username == username).all()
    return portfolios

class PortfolioSaveRequest(BaseModel):
    name: str
    tickers: list
    weights: dict
    html_report: Optional[str] = None

@app.post("/api/portfolios")
def save_portfolio(req: PortfolioSaveRequest, username: str = Depends(get_current_user), db: Session = Depends(get_db)):
    portfolio = SavedPortfolio(
        username=username,
        name=req.name,
        tickers=req.tickers,
        weights=req.weights,
        html_report=req.html_report
    )
    db.add(portfolio)
    db.commit()
    return {"status": "success", "id": portfolio.id}

@app.delete("/api/portfolios/{portfolio_id}")
def delete_portfolio(portfolio_id: int, username: str = Depends(get_current_user), db: Session = Depends(get_db)):
    portfolio = db.query(SavedPortfolio).filter(SavedPortfolio.id == portfolio_id, SavedPortfolio.username == username).first()
    if not portfolio:
        raise HTTPException(status_code=404, detail="Portfolio not found")
    db.delete(portfolio)
    db.commit()
    return {"status": "success"}



@app.get("/api/backtests")
def get_backtests(username: str = Depends(get_current_user), db: Session = Depends(get_db)):
    backtests = db.query(BacktestHistory).filter(BacktestHistory.username == username).order_by(BacktestHistory.executed_at.desc()).all()
    return backtests

@app.post("/api/webhooks/config")
def update_webhook(payload: dict, username: str = Depends(get_current_user), db: Session = Depends(get_db)):
    config = db.query(WebhookConfig).filter(WebhookConfig.username == username).first()
    if not config:
        config = WebhookConfig(
            username=username, 
            webhook_url=payload.get("url"), 
            api_secret_key=f"whsec_{secrets.token_hex(16)}"
        )
        db.add(config)
    else:
        config.webhook_url = payload.get("url")
    db.commit()
    masked_key = f"whsec_{config.api_secret_key.replace('whsec_', '')[:6]}***" if config.api_secret_key else ""
    return {"status": "success", "api_secret_key": masked_key}


class AuthRequest(BaseModel):
    username: str = ""
    key: str

class KeyGenRequest(BaseModel):
    admin_key: str
    new_key: str

class PortfolioRequest(BaseModel):
    tickers: List[str]
    capital: float = 100000.0
    risk_input: int = 5
    model: int = 1
    allocation_engine: int = 1
    allow_shorting: bool = False
    tax_enabled: bool = False
    garch_enabled: bool = True
    currency: str = "$"
    custom_constraints: Optional[List[dict]] = None
    fixed_weights: Optional[dict] = None
    rebalance_freq_months: int = 3




# ─────────────────────────────────────────────────────────────
# SIMPLE UUID SESSION STORE — backed by a JSON file
# /tmp is always writable on Linux (HuggingFace Spaces, Docker).
# Falls back to app-dir/output on Windows.
# ─────────────────────────────────────────────────────────────
def _make_session_token(access_key: str, username: str = "") -> str:
    """Create a new UUID session, persist it in Redis, return the token."""
    token = str(uuid.uuid4()).replace("-", "")  # 32 hex chars, no special chars
    SESSION_MAX_AGE = 43200  # 12 hours
    try:
        session_data = {"access_key": access_key, "username": username}
        cache.redis_client.setex(f"session:{token}", SESSION_MAX_AGE, json.dumps(session_data))
        print(f"[SESSION] created token for '{access_key}' ('{username}') in Redis")
    except Exception as e:
        logger.error(f"[SESSION] Redis set error: {e}")
    return token


def _get_session_data(request: Request) -> Optional[dict]:
    token = request.cookies.get("we_session", "").strip().strip('"').strip("'")
    if not token or len(token) < 8:
        return None
    try:
        data = cache.redis_client.get(f"session:{token}")
        if data:
            return json.loads(data)
    except Exception as e:
        logger.error(f"[SESSION] Redis get error: {e}")
    return None


def _validate_session_cookie(request: Request) -> bool:
    """Check if the session cookie token is valid and not expired."""
    data = _get_session_data(request)
    if not data:
        token = request.cookies.get("we_session", "").strip().strip('"').strip("'")
        print(f"[SESSION] INVALID token={token[:16]}...")
        return False
    return True


def _get_session_access_key(request: Request) -> Optional[str]:
    data = _get_session_data(request)
    return data.get("access_key") if data else None


def _get_session_username(request: Request) -> Optional[str]:
    data = _get_session_data(request)
    return data.get("username") if data else None





@app.get("/")
async def read_login(request: Request):
    return FileResponse(os.path.join(STATIC_DIR, "login.html"), media_type="text/html; charset=utf-8")

@app.get("/main")
async def read_index(request: Request):
    # Serve the app — real security is enforced on all /api/* endpoints per-request.
    # Client-side guard in index.html checks sessionStorage on load and redirects
    # to / if no valid key is found (prevents blank/broken UI without a key).
    return FileResponse(os.path.join(STATIC_DIR, "index.html"), media_type="text/html; charset=utf-8")

@app.get("/options")
@app.get("/options.html")
async def read_options(request: Request):
    return FileResponse(os.path.join(STATIC_DIR, "options.html"), media_type="text/html; charset=utf-8")

@app.get("/api/verify")
async def api_verify(request: Request, authorization: Optional[str] = Header(None)):
    """Quick key-validity check used by the client-side session guard on page load."""
    key = None
    if authorization and authorization.startswith("Bearer "):
        key = authorization[7:]
    if not key:
        raise HTTPException(status_code=401, detail="No key provided")
    if access_manager.validate_key(key, silent=True):
        return JSONResponse({"valid": True, "is_master": access_manager.is_master_key(key)})
    raise HTTPException(status_code=401, detail="Invalid key")

@app.post("/api/logout")
async def api_logout(request: Request):
    response = JSONResponse({"status": "logged_out"})
    response.delete_cookie("we_session", path="/")
    return response


@app.get("/admin")
async def read_admin():
    return FileResponse(os.path.join(STATIC_DIR, "admin.html"))

@app.get("/api/market_ticker")
async def market_ticker():
    cached_data = cache.cache_get_json("market_ticker")
    if cached_data:
        return cached_data
    
    results = []
    tickers = ["SPY", "QQQ", "TLT", "GLD", "BTC-USD", "UUP", "USO"]
    try:
        import yfinance as yf
        # Try batch download first
        df = yf.download(tickers, period="5d", progress=False, timeout=5)
        if not df.empty:
            if isinstance(df.columns, pd.MultiIndex):
                close_df = df['Close'] if 'Close' in df.columns.get_level_values(0) else df
            else:
                close_df = df
            
            for t in tickers:
                try:
                    s = close_df[t].dropna()
                    if len(s) >= 2:
                        p1, p0 = float(s.iloc[-1]), float(s.iloc[-2])
                        chg = (p1 - p0) / p0
                        results.append({
                            "name": t,
                            "price": f"{p1:.2f}",
                            "change": round(chg, 6)
                        })
                except Exception as te:
                    logger.warning(f"Ticker {t} extraction failed: {te}")
    except Exception as e:
        logger.error(f"Batch market ticker fetch failed: {e}")
        # Fallback: try each ticker individually
        try:
            import yfinance as yf
            for t in tickers:
                try:
                    df_single = yf.download(t, period="5d", progress=False, timeout=5)
                    if not df_single.empty:
                        s = df_single['Close'].dropna()
                        if len(s) >= 2:
                            p1, p0 = float(s.iloc[-1]), float(s.iloc[-2])
                            chg = (p1 - p0) / p0
                            results.append({
                                "name": t,
                                "price": f"{p1:.2f}",
                                "change": round(chg, 6)
                            })
                except Exception as te2:
                    logger.warning(f"Individual ticker {t} failed: {te2}")
        except Exception as e2:
            logger.error(f"Individual ticker fallback also failed: {e2}")
    
    if results:
        cache.cache_set_json("market_ticker", results, ttl=300)
    
    return results if results else (cache.cache_get_json("market_ticker") or [])

@app.get("/api/finance_news")
async def finance_news():
    cached_data = cache.cache_get_json("finance_news")
    if cached_data:
        return cached_data
    
    results = []
    try:
        import yfinance as yf
        news = yf.Ticker("SPY").news
        for item in news[:15]:
            content = item.get("content", {})
            title = content.get("title", item.get("title", "No Title"))
            pub_date = content.get("pubDate", item.get("providerPublishTime", ""))
            if isinstance(pub_date, int):
                import datetime
                pub_date = datetime.datetime.fromtimestamp(pub_date).strftime("%Y-%m-%d")
            elif "T" in str(pub_date):
                pub_date = str(pub_date).split("T")[0]
            provider = item.get("provider", {})
            source = provider.get("displayName", item.get("publisher", "Yahoo Finance"))
            link = item.get("clickThroughUrl", {}).get("url", item.get("link", ""))
            results.append({
                "title": title,
                "source": source,
                "time": pub_date,
                "url": link
            })
        if results:
            cache.cache_set_json("finance_news", results, ttl=300)
    except Exception as e:
        logger.error(f"News fetch failed: {e}")
        
    return results or cache.cache_get_json("finance_news") or []


from api import RateLimiter

@app.post("/api/auth", dependencies=[Depends(RateLimiter(limit=5, window=60))])
async def api_auth(req: AuthRequest, request: Request):
    if not req.username.strip() and not access_manager.is_master_key(req.key):
        raise HTTPException(status_code=400, detail="Username is required")
        
    forwarded = request.headers.get("X-Forwarded-For")
    ip = forwarded.split(",")[0].strip() if forwarded else (request.client.host if request.client else "Unknown")
    try:
        if access_manager.validate_key(req.key, username=req.username, ip=ip):
            is_master = access_manager.is_master_key(req.key)
            # Issue a secure HTTP-only session cookie so /main cannot be bypassed via URL
            token = _make_session_token(req.key, req.username)
            response = JSONResponse({"status": "success", "message": "Access Granted", "is_master": is_master, "username": req.username})
            response.set_cookie(
                key="we_session",
                value=token,
                httponly=True,      # JS cannot read this cookie
                samesite="lax",     # lax: allows top-level nav (strict blocks post-login redirect)
                max_age=43200,      # 12 hours
                path="/"
            )
            return response
        raise HTTPException(status_code=401, detail="Invalid or Expired Access Key")
    except ValueError as e:
        raise HTTPException(status_code=429, detail=str(e))

class AdminKeyGenRequest(BaseModel):
    admin_key: str
    username: str = ""
    hours: int = 1

class AdminRevokeRequest(BaseModel):
    admin_key: str
    target_key: str = ""
    confirm_token: str = ""
    target_key: str = ""
    confirm_token: str = ""
    target_key: str = ""
    confirm_token: str = ""
    target_key: str = ""
    confirm_token: str = ""
    target_key: str

@app.post("/api/admin/generate")
async def admin_generate(req: AdminKeyGenRequest):
    new_key = access_manager.generate_otk(req.admin_key, username=req.username, hours=req.hours)
    if new_key:
        return {"status": "success", "key": new_key, "message": f"OTK '{new_key}' created for {req.username}."}
    raise HTTPException(status_code=401, detail="Invalid Admin Key")

@app.post("/api/admin/revoke")
async def admin_revoke(req: AdminRevokeRequest):
    success = access_manager.revoke_otk(req.admin_key, req.target_key)
    if success:
        return {"status": "success", "message": f"Key '{req.target_key}' revoked."}
    raise HTTPException(status_code=401, detail="Invalid Admin Key or Key Not Found")

@app.get("/api/admin/keys")
async def admin_list_keys(admin_key: str = Header(...)):
    keys = access_manager.get_all_keys(admin_key)
    if admin_key != access_manager.MASTER_KEY:
        raise HTTPException(status_code=401, detail="Invalid Admin Key")
    return {"keys": keys}


import hashlib
def get_action_token(action: str, admin_key: str) -> str:
    window = int(time.time() / 120)
    return hashlib.sha256(f"{admin_key}:{action}:{window}".encode()).hexdigest()

@app.get("/api/admin/action_token")
async def get_admin_action_token(action: str, admin_key: str = Header(...)):
    if admin_key != access_manager.MASTER_KEY:
        raise HTTPException(status_code=401, detail="Invalid Admin Key")
    return {"token": get_action_token(action, admin_key)}

@app.post("/api/admin/revoke_all")
async def admin_revoke_all(req: AdminRevokeRequest):
    if req.admin_key != access_manager.MASTER_KEY:
        raise HTTPException(status_code=401, detail="Invalid Admin Key")
    if getattr(req, "confirm_token", "") != get_action_token("revoke_all", req.admin_key):
        raise HTTPException(status_code=403, detail="Invalid or expired confirmation token")
    if getattr(req, "confirm_token", "") != get_action_token("revoke_all", req.admin_key):
        raise HTTPException(status_code=403, detail="Invalid or expired confirmation token")
    if getattr(req, "confirm_token", "") != get_action_token("revoke_all", req.admin_key):
        raise HTTPException(status_code=403, detail="Invalid or expired confirmation token")
    if getattr(req, "confirm_token", "") != get_action_token("revoke_all", req.admin_key):
        raise HTTPException(status_code=403, detail="Invalid or expired confirmation token")
    keys = access_manager.get_all_keys(req.admin_key)
    count = 0
    for k, v in keys.items():
        if not v.get("revoked", False):
            access_manager.revoke_otk(req.admin_key, k)
            count += 1
    return {"status": "success", "revoked_count": count}

@app.get("/api/admin/logs")
async def admin_get_logs(admin_key: str = Header(...)):
    if admin_key != access_manager.MASTER_KEY:
        raise HTTPException(status_code=401, detail="Invalid Admin Key")
    
    from constants import OUTPUT_DIR
    import os
    from datetime import datetime, timedelta
    
    log_file = os.path.join(OUTPUT_DIR, "access.log")
    logs = []
    if os.path.exists(log_file):
        cutoff = datetime.now() - timedelta(hours=24)
        with open(log_file, "r", encoding="utf-8") as f:
            for line in f:
                if not line.strip(): continue
                # Parse the datetime from the log line (e.g. "2026-06-12 22:33:45,491 - ...")
                try:
                    dt_str = line.split(" - ")[0].split(",")[0]
                    log_dt = datetime.strptime(dt_str, "%Y-%m-%d %H:%M:%S")
                    if log_dt >= cutoff:
                        logs.append(line.strip())
                except Exception:
                    # If parsing fails, just append it
                    logs.append(line.strip())
    
    # Return last 500 lines max to prevent huge payloads
    return {"logs": logs[-500:]}

@app.post("/api/preview")
async def preview_portfolio(req: PortfolioRequest, x_access_key: Optional[str] = Header(None), x_username: Optional[str] = Header(None)):
    if not access_manager.validate_key(x_access_key, username=x_username):
        raise HTTPException(status_code=401, detail="Unauthorized")
        
    try:
        overrides = {
            'tickers': req.tickers,
            'capital': req.capital,
            'risk_input': req.risk_input,
            'risk_factor': {1:0.1, 2:0.5, 3:1.0, 4:2.0, 5:3.0, 6:5.0, 7:7.5, 8:10.0, 9:15.0, 10:25.0}.get(req.risk_input, 3.0),
            'model': req.model,
            'allocation_engine': req.allocation_engine,
            'single_asset_min': -1.0 if req.allow_shorting else 0.0,
            'tax_enabled': req.tax_enabled,
            'garch_enabled': req.garch_enabled,
            'custom_constraints': req.custom_constraints,
            'fixed_weights': req.fixed_weights
        }
        
        result = core_engine.run_engine(overrides=overrides, serve=False, preview_only=True)
        
        return {
            "status": "success", 
            "target_weights": result.get("target_weights", {}),
            "efficient_frontier": result.get("efficient_frontier", {"vols": [], "rets": []})
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

def _run_optimization(tid, request_dict, x_access_key, x_username):
    try:
        def _format_currency(t):
            t = t.upper().strip()
            if len(t) == 3 and t in ['EUR', 'GBP', 'JPY', 'CHF', 'AUD', 'CAD', 'SEK', 'NOK', 'DKK', 'NZD']:
                return f"{t}USD=X"
            if len(t) == 3 and t in ['BTC', 'ETH', 'SOL', 'XRP', 'ADA', 'DOT', 'LTC', 'BNB', 'LINK', 'UNI']:
                return f"{t}-USD"
            if len(t) == 6 and t.endswith('USD'):
                return f"{t}=X"
            return t
        
        clean_tickers = [_format_currency(t) for t in request_dict["tickers"]]
        
        overrides = {
            'universe': clean_tickers,
            'tickers': clean_tickers,
            'capital': request_dict["capital"],
            'risk_input': request_dict["risk_input"],
            'risk_factor': {1:0.1, 2:0.5, 3:1.0, 4:2.0, 5:3.0, 6:5.0, 7:7.5, 8:10.0, 9:15.0, 10:25.0}.get(request_dict["risk_input"], 3.0),
            'model': request_dict["model"],
            'allocation_engine': request_dict["allocation_engine"],
            'single_asset_min': -1.0 if request_dict["allow_shorting"] else 0.0,
            'tax_enabled': request_dict["tax_enabled"],
            'garch_enabled': request_dict["garch_enabled"],
            'currency_symbol': request_dict["currency"],
            'custom_constraints': request_dict["custom_constraints"],
            'fixed_weights': request_dict["fixed_weights"],
            'rebalance_freq_months': request_dict["rebalance_freq_months"]
        }
        
        tracer.start_trace(tid)
        tracer.add_flag(tid, "TASK_INIT", f"Target Universe: {overrides.get('universe')}")
        
        # Determine if we are acting as a Proxy (e.g., on Render) or if we should run the math locally.
        is_proxy = os.getenv("RENDER") == "true" or os.getenv("IS_PROXY") == "true"
        is_backend = not is_proxy
        
        if is_backend:
            tracer.add_flag(tid, "BACKEND_START", "Running compute engine locally")
            # We are the backend (either HF or localhost). Run the heavy math.
            result = core_engine.run_engine(overrides=overrides, serve=False, task_id=tid)
            
            try:
                engine = database.get_pg_engine()
                from database import get_db_session
                with get_db_session() as db:
                    stats = result.get('stats', {})
                    bt_stats = result.get('bt_stats', {})
                    html_report_content = result.get('html_report', "")
                    if html_report_content:
                        report_path = os.path.join(OUTPUT_DIR, "portfolio_report.html")
                        os.makedirs(OUTPUT_DIR, exist_ok=True)
                        with open(report_path, "w", encoding="utf-8") as f:
                            f.write(html_report_content)
                            
                    history = database.BacktestHistory(
                        username=x_username or "admin",
                        model_used=str(request_dict["model"]),
                        return_pct=bt_stats.get('Annualized Return', 0),
                        sharpe_ratio=bt_stats.get('Sharpe Ratio', 0),
                        max_drawdown=bt_stats.get('Max Drawdown', 0),
                        tickers=list(result.get('target_weights', {}).keys()),
                        weights=result.get('target_weights', {}),
                        html_report=html_report_content
                    )
                    db.add(history)
                    db.commit()
                    
                    # Keep only the last 15 backtests per user (FIFO)
                    user_history = db.query(database.BacktestHistory).filter(
                        database.BacktestHistory.username == (x_username or "admin")
                    ).order_by(database.BacktestHistory.executed_at.desc()).all()
                    
                    if len(user_history) > 15:
                        for old_run in user_history[15:]:
                            db.delete(old_run)
                        db.commit()

                    # --- WEBHOOK EXECUTION ---
                    config = db.query(database.WebhookConfig).filter(database.WebhookConfig.username == (x_username or "anonymous")).first()
                    if config and config.webhook_url:
                        import requests
                        payload = {
                            "model": request_dict["model"],
                            "weights": result.get("target_weights", {}),
                            "stats": stats,
                            "status": "completed"
                        }
                        try:
                            requests.post(config.webhook_url, json=payload, headers={"Authorization": f"Bearer {config.api_secret_key}"}, timeout=5)
                        except Exception as wh_err:
                            logger.error(f"Webhook failed to fire: {wh_err}")

            except Exception as e:
                logger.error(f"Failed to save backtest history or fire webhook: {e}")
            
            tracer.add_flag(tid, "HF_BACKEND_COMPLETE", "Math engine finished")
            
            # Fetch, update, and re-assign the entire dict to trigger FileBackedDict.__setitem__ and save()
            task_data = BACKGROUND_TASKS.get(tid, {})
            task_data["status"] = "completed"
            task_data["message"] = "Report generated."
            task_data["target_weights"] = result.get("target_weights", {})
            task_data["stats"] = result.get("stats", {})
            BACKGROUND_TASKS[tid] = task_data

        else:
                import requests
                hf_url = os.getenv("HF_BACKEND_URL", "").rstrip('/')
                if not hf_url:
                    tracer.add_flag(tid, "ERROR", "HF_BACKEND_URL not set in environment")
                    BACKGROUND_TASKS[tid]["status"] = "error"
                    BACKGROUND_TASKS[tid]["error"] = "HF_BACKEND_URL not configured for proxying."
                    return
                # Use the master key to bypass HF API restrictions and authenticate internally
                hf_key = os.getenv("HF_MASTER_KEY", "")
                
                tracer.add_flag(tid, "PROXY_START", f"Forwarding to HF backend: {hf_url}")
                try:
                    proxy_res = requests.post(
                        f"{hf_url}/api/generate", 
                        json=request_dict,
                        headers={"X-Access-Key": hf_key},
                        timeout=120
                    )
                except requests.exceptions.Timeout:
                    raise Exception("Hugging Face Backend timed out while queuing the optimization. This is likely due to the free tier spinning up from sleep.")
                except requests.exceptions.RequestException as req_e:
                    raise Exception(f"Hugging Face Backend connection failed: {req_e}")
                
                if not proxy_res.ok:
                    raise Exception(f"Hugging Face Backend Error ({proxy_res.status_code}) at {hf_url}/api/generate. Check HF_BACKEND_URL.")
                
                proxy_data = proxy_res.json()
                remote_task_id = proxy_data.get("task_id")
                tracer.add_flag(tid, "PROXY_HANDOFF_SUCCESS", f"HF Task ID: {remote_task_id}")
                
                if not remote_task_id:
                    raise Exception("Failed to get remote task ID from Hugging Face.")
                
                import time
                retries = 0
                while True:
                    time.sleep(2)
                    try:
                        status_res = requests.get(
                            f"{hf_url}/api/status/{remote_task_id}",
                            headers={"X-Access-Key": hf_key},
                            timeout=15
                        )
                        retries = 0 # Reset on success
                    except requests.exceptions.RequestException as e:
                        retries += 1
                        if retries > 10:
                            tracer.add_flag(tid, "PROXY_POLL_TIMEOUT", f"Poll failed 10 times: {e}")
                            BACKGROUND_TASKS[tid]["status"] = "error"
                            BACKGROUND_TASKS[tid]["message"] = "Hugging Face Backend is unreachable (Timeout). It might have crashed or restarted."
                            break
                        continue
                    
                    if status_res.ok:
                        s_data = status_res.json()
                        BACKGROUND_TASKS[tid]["status"] = s_data["status"]
                        BACKGROUND_TASKS[tid]["message"] = s_data["message"]
                        
                        if s_data["status"] == "completed":
                            BACKGROUND_TASKS[tid]["target_weights"] = s_data.get("target_weights", {})
                            BACKGROUND_TASKS[tid]["stats"] = s_data.get("stats", {})
                            BACKGROUND_TASKS[tid]["bt_stats"] = s_data.get("bt_stats", {})
                            
                            # Download the completed HTML report from HF to Render
                            report_res = requests.get(f"{hf_url}/report")
                            if report_res.ok:
                                report_path = os.path.join(OUTPUT_DIR, "portfolio_report.html")
                                os.makedirs(OUTPUT_DIR, exist_ok=True)
                                with open(report_path, "wb") as f:
                                    f.write(report_res.content)
                            break
                        elif s_data["status"] == "error":
                            error_msg = s_data.get("message", "Unknown error from HF")
                            tracer.add_flag(tid, "PROXY_POLL_ERROR", error_msg)
                            BACKGROUND_TASKS[tid]["status"] = "error"
                            BACKGROUND_TASKS[tid]["message"] = error_msg
                            break
                    else:
                        tracer.add_flag(tid, "PROXY_CONNECTION_LOST", f"Status Code: {status_res.status_code}")
                        retries += 1
                        if retries > 10:
                            BACKGROUND_TASKS[tid]["status"] = "error"
                            BACKGROUND_TASKS[tid]["message"] = f"Lost connection to Hugging Face backend (HTTP {status_res.status_code})."
                            raise Exception(f"Lost connection to Hugging Face backend. (HTTP {status_res.status_code})")
                        continue
                        
    except (Exception, SystemExit) as e:
        error_trace = traceback.format_exc()
        tracer.add_flag(tid, "FATAL_ERROR", f"{str(e)}\n\nTraceback:\n{error_trace}")
        logger.error(f"Optimization failed: {error_trace}")
        task_data = BACKGROUND_TASKS.get(tid, {})
        task_data["status"] = "error"
        task_data["message"] = f"Error: {str(e)} (Check console logs for details)"
        BACKGROUND_TASKS[tid] = task_data
            
@app.post("/api/generate")
async def generate_portfolio(req: PortfolioRequest, x_access_key: Optional[str] = Header(None), x_username: Optional[str] = Header(None)):
    if not access_manager.validate_key(x_access_key, username=x_username):
        raise HTTPException(status_code=401, detail="Unauthorized")
        
    task_id = str(uuid.uuid4())
    BACKGROUND_TASKS[task_id] = {"status": "running", "message": "Initializing...", "target_weights": {}}
    
    request_dict = req.model_dump() if hasattr(req, "model_dump") else req.dict()
    process_pool.submit(_run_optimization, task_id, request_dict, x_access_key, x_username)

    return {
        "status": "queued", 
        "task_id": task_id,
        "message": "Optimization started in background."
    }





# ─────────────────────────────────────────────
# HFT SIMULATOR API
# ─────────────────────────────────────────────
class HFTRequest(BaseModel):
    symbols: list[str]
    duration_ms: int = 1000
    latency_ms: int = 5
    tick_ms: int = 10
    strategy: Optional[str] = None
    target_qty: float = 100.0

@app.post("/api/hft/simulate")
def hft_simulate(req: HFTRequest, x_access_key: Optional[str] = Header(None)):
    if not access_manager.validate_key(x_access_key, silent=True):
        raise HTTPException(status_code=401, detail="Unauthorized")
        
    try:
        from hft_simulator import HFTSimulator
        from hft_strategies import MarketMakingStrategy, MomentumStrategy, MeanReversionStrategy, ExecutionBridgeStrategy
        import yfinance as yf
        
        # Initialize simulation
        start_time = datetime.now()
        duration_sec = req.duration_ms / 1000.0
        sim = HFTSimulator(req.symbols, start_time, duration_sec, req.tick_ms, req.latency_ms)
        
        # Get live prices as starting point, fallback to 100
        initial_prices = {}
        from concurrent.futures import ThreadPoolExecutor
        
        def fetch_price(sym):
            try:
                hist = yf.download(sym, period="1d", timeout=5)
                return sym, float(hist['Close'].iloc[-1]) if not hist.empty else 100.0
            except Exception:
                return sym, 100.0
                
        executor = ThreadPoolExecutor(max_workers=10)
        futures = {executor.submit(fetch_price, sym): sym for sym in req.symbols}
        from concurrent.futures import as_completed, TimeoutError
        try:
            for future in as_completed(futures.keys(), timeout=2.0):
                sym, price = future.result()
                initial_prices[sym] = price
        except TimeoutError:
            logging.warning("yfinance fetch timed out in HFT, falling back to synthetic prices.")
            for sym in req.symbols:
                if sym not in initial_prices:
                    initial_prices[sym] = 100.0
        finally:
            executor.shutdown(wait=False)
                
        sim.initialize_books(initial_prices)
        
        # Attach requested strategy
        if req.strategy == 'market_making':
            for sym in req.symbols:
                sim.add_strategy(MarketMakingStrategy(sym))
        elif req.strategy == 'momentum':
            for sym in req.symbols:
                sim.add_strategy(MomentumStrategy(sym))
        elif req.strategy == 'mean_reversion':
            for sym in req.symbols:
                sim.add_strategy(MeanReversionStrategy(sym))
        elif req.strategy == 'execution':
            for sym in req.symbols:
                sim.add_strategy(ExecutionBridgeStrategy(sym, req.target_qty, 'buy', chunks=10))
                
        results = sim.run()
        res_dict = results.to_dict()
        res_dict['initial_prices'] = initial_prices
        return {"status": "success", "results": res_dict}
    except Exception:
        import traceback
        logging.error(f"HFT Simulation failed: {traceback.format_exc()}")
        
        # Fallback to prevent UI crash
        from datetime import timedelta
        now = datetime.now()
        fallback_results = {
            "metrics": {
                "total_trades": 0,
                "volume": 0.0,
                "avg_spread": 0.01
            },
            "times": [now.isoformat(), (now + timedelta(milliseconds=req.duration_ms)).isoformat()],
            "mid_prices": [100.0, 100.0],
            "spreads": [0.01, 0.01],
            "trade_prices": [],
            "trade_times": [],
            "initial_prices": {sym: 100.0 for sym in req.symbols},
            "final_depth": None
        }
        return {"status": "success", "results": fallback_results}

def _cpp_worker(returns_array, weights_array, expected_returns_array, queue):
    import time
    try:
        import quant_engine_cpp
        cpp_results = {}
        
        # 1. Ledoit-Wolf
        start = time.time()
        cpp_cov = quant_engine_cpp.compute_ledoit_wolf_covariance(returns_array)
        cpp_results["ledoit_wolf_ms"] = (time.time() - start) * 1000
        
        # 2. Monte Carlo (scaled down slightly for HF environments)
        start = time.time()
        quant_engine_cpp.run_monte_carlo(weights_array, expected_returns_array, cpp_cov, 5000, 252, 10000.0)
        cpp_results["monte_carlo_ms"] = (time.time() - start) * 1000
        
        # 3. GARCH
        start = time.time()
        quant_engine_cpp.batch_fit_garch(returns_array)
        cpp_results["garch_ms"] = (time.time() - start) * 1000
        
        queue.put({"status": "success", "results": cpp_results})
    except Exception as e:
        queue.put({"status": "error", "error": str(e)})

@app.get("/api/benchmark")
def get_benchmark(x_access_key: Optional[str] = Header(None)):
    if not access_manager.validate_key(x_access_key, silent=True):
        raise HTTPException(status_code=401, detail="Unauthorized")
        
    import time
    import numpy as np
    
    # Generate random returns for benchmarking
    # 250 days, 100 assets
    rng = np.random.default_rng(42)
    returns = rng.normal(0, 0.02, (250, 100))
    weights = np.ones(100) / 100.0
    expected_returns = rng.normal(0.05, 0.02, 100)
    
    results = {
        "python": {},
        "cpp": {}
    }
    
    # --- PYTHON BENCHMARKS ---
    
    # 1. Ledoit-Wolf (simplified for benchmark to prevent OpenBLAS threadpool deadlocks)
    start = time.time()
    try:
        py_cov = np.cov(returns, rowvar=False)
    except Exception:
        py_cov = np.eye(100)
    results["python"]["ledoit_wolf_ms"] = (time.time() - start) * 1000
    
    # 2. Monte Carlo (Pure Python)
    def py_monte_carlo(weights, expected_returns, cov_matrix, num_simulations, days, initial_portfolio_value, rng):
        L = np.linalg.cholesky(cov_matrix + np.eye(len(weights))*1e-8)
        mu_daily = expected_returns / 252.0
        final_vals = np.zeros(num_simulations)
        for s in range(num_simulations):
            port_val = initial_portfolio_value
            for d in range(days):
                z = rng.normal(0, 1, len(weights))
                shocks = L @ z
                daily_returns = mu_daily + shocks
                port_ret = weights @ daily_returns
                port_val *= (1.0 + port_ret)
            final_vals[s] = port_val
        return final_vals
        
    start = time.time()
    py_monte_carlo(weights, expected_returns, py_cov, 5000, 252, 10000.0, rng)
    results["python"]["monte_carlo_ms"] = (time.time() - start) * 1000
    
    # 3. GARCH (Pure Python fallback)
    def py_garch(returns_1d):
        t = len(returns_1d)
        cond_vol = np.zeros(t)
        var0 = np.var(returns_1d)
        h = var0
        cond_vol[0] = np.sqrt(var0)
        ll = 0.0
        alpha, beta, omega = 0.1, 0.8, var0 * 0.1
        for i in range(1, t):
            h = omega + alpha * returns_1d[i-1]**2 + beta * h
            if h < 1e-8: h = 1e-8
            cond_vol[i] = np.sqrt(h)
            ll += -0.5 * (np.log(2 * np.pi) + np.log(h) + (returns_1d[i]**2) / h)
        return ll
        
    start = time.time()
    import concurrent.futures
    with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor:
        try:
            list(executor.map(py_garch, [returns[:, i] for i in range(100)], timeout=5.0))
        except concurrent.futures.TimeoutError:
            pass
    results["python"]["garch_ms"] = (time.time() - start) * 1000
    
    # --- C++ BENCHMARKS ---
    import multiprocessing as mp
    cpp_error_holder = {}
    
    # _cpp_worker is now defined at module level

    # Use 'spawn' context if possible to ensure clean memory state, or default Process
    try:
        ctx = mp.get_context('spawn')
    except Exception:
        ctx = mp
        
    q = ctx.Queue()
    p = ctx.Process(target=_cpp_worker, args=(returns, weights, expected_returns, q))
    p.daemon = True
    p.start()
    p.join(timeout=8.0)
    
    if p.is_alive():
        p.terminate()
        p.join()
        cpp_error_holder['error'] = "C++ module timed out or hung"
    else:
        if not q.empty():
            worker_res = q.get()
            if worker_res["status"] == "success":
                results["cpp"] = worker_res["results"]
                results["cpp_available"] = True
            else:
                cpp_error_holder['error'] = worker_res["error"]
        else:
            cpp_error_holder['error'] = "C++ process crashed unexpectedly (Segfault/OOM)"
    
    if 'error' in cpp_error_holder:
        import logging
        logging.getLogger(__name__).warning(f"C++ engine benchmark failed: {cpp_error_holder['error']}")
        results["cpp_available"] = False
        results["cpp"]["ledoit_wolf_ms"] = 0
        results["cpp"]["monte_carlo_ms"] = 0
        results["cpp"]["garch_ms"] = 0
        
    return {"status": "success", "results": results}

@app.get("/api/traces")
async def get_all_traces():
    return tracer.traces

@app.get("/api/trace/{task_id}")
async def get_task_trace(task_id: str):
    return {"task_id": task_id, "trace": tracer.get_trace(task_id)}

def is_running_on_hf() -> bool:
    return os.environ.get("SPACE_ID") is not None

@app.get("/api/status/{task_id}")
async def get_task_status(task_id: str, x_access_key: Optional[str] = Header(None)):
    if not access_manager.validate_key(x_access_key, silent=True):
        raise HTTPException(status_code=401, detail="Unauthorized")
        
    is_proxy = os.getenv("RENDER") == "true" or os.getenv("IS_PROXY") == "true"
    is_backend = not is_proxy
    
    if is_backend:
        task = BACKGROUND_TASKS.get(task_id)
        if not task:
            raise HTTPException(status_code=404, detail="Task not found")
        return task
    else:
        # Check local task state first in case proxying the generate request failed
        local_task = BACKGROUND_TASKS.get(task_id)
        if local_task and local_task.get("status") == "error":
            return local_task

        hf_url = os.getenv("HF_BACKEND_URL", "").rstrip('/')
        if not hf_url:
            return {"status": "error", "error": "HF_BACKEND_URL not configured for proxying status check."}
        try:
            import requests
            hf_res = requests.get(
                f"{hf_url}/api/status/{task_id}",
                headers={"X-Access-Key": x_access_key},
                timeout=5
            )
            if not hf_res.ok:
                return {"status": "error", "message": f"Backend returned {hf_res.status_code}", "task_id": task_id}
            return hf_res.json()
        except requests.exceptions.Timeout:
            return {"status": "running", "message": "Backend polling timeout...", "task_id": task_id}
        except Exception as e:
            return {"status": "error", "message": f"Proxy error: {str(e)}", "task_id": task_id}

@app.get("/report")
async def get_report():
    is_proxy = os.getenv("RENDER") == "true" or os.getenv("IS_PROXY") == "true"
    is_backend = not is_proxy
    
    if is_backend:
        report_path = os.path.join(OUTPUT_DIR, "portfolio_report.html")
        if os.path.exists(report_path):
            return FileResponse(report_path)
        raise HTTPException(status_code=404, detail="Report not generated yet.")
    else:
        hf_url = os.getenv("HF_BACKEND_URL", "").rstrip('/')
        if not hf_url:
            raise HTTPException(status_code=500, detail="HF_BACKEND_URL not configured.")
        import requests
        try:
            res = requests.get(f"{hf_url}/report", timeout=10)
            if res.ok:
                from fastapi.responses import HTMLResponse
                return HTMLResponse(content=res.text)
            else:
                raise HTTPException(status_code=res.status_code, detail="Report not ready on backend.")
        except Exception as e:
            raise HTTPException(status_code=500, detail=str(e))

def _alert_daemon():
    """Background daemon to check for market drops and perform monthly maintenance."""
    import time
    last_cleanup = time.time()
    
    while True:
        try:
            # Wake up every 1 hour (3600 seconds)
            time.sleep(3600)
            
            # 1. Check for SPY drops
            try:
                ticker = yf.Ticker("SPY")
                hist = ticker.history(period="2d")
                if len(hist) >= 2:
                    current = float(hist['Close'].iloc[-1])
                    prev = float(hist['Close'].iloc[-2])
                    pct_change = ((current - prev) / prev) * 100
                    
                    if pct_change <= -5.0:
                        access_manager.send_telegram_alert(f"🚨 **MARKET ALERT**\nSPY has dropped by {pct_change:.2f}%!\nCheck the portfolio engine.")
            except Exception:
                pass
                
            # 2. Monthly DB/Log Cleanup
            # 30 days = 2592000 seconds
            if time.time() - last_cleanup > 2592000:
                try:
                    logger.info("Performing monthly database and log cleanup...")
                    # access_manager should clean up its expired keys and sync to Redis
                    if hasattr(access_manager, 'cleanup_expired_keys'):
                        access_manager.cleanup_expired_keys()
                    last_cleanup = time.time()
                except Exception as e:
                    logger.error(f"Monthly cleanup failed: {e}")
                    
        except Exception:
            pass # Suppress daemon errors

import warnings
warnings.filterwarnings('ignore', category=DeprecationWarning, module='fastapi')

@app.on_event("startup")
def startup_event():
    import threading
    import sqlite3
    import os
    
    # Auto-migrate SQLite schema to include backtest_history JSON columns
    db_path = os.path.join(OUTPUT_DIR, 'portfolio_db.sqlite3')
    if os.path.exists(db_path):
        try:
            conn = sqlite3.connect(db_path)
            c = conn.cursor()
            try:
                c.execute("ALTER TABLE backtest_history ADD COLUMN tickers JSON")
            except Exception:
                pass
            try:
                c.execute("ALTER TABLE backtest_history ADD COLUMN weights JSON")
            except Exception:
                pass
            try:
                c.execute("ALTER TABLE backtest_history ADD COLUMN html_report TEXT")
            except Exception:
                pass
            try:
                c.execute("ALTER TABLE saved_portfolios ADD COLUMN html_report TEXT")
            except Exception:
                pass
            try:
                c.execute("""
                    CREATE TABLE IF NOT EXISTS user_memory (
                        id INTEGER PRIMARY KEY AUTOINCREMENT,
                        access_key VARCHAR UNIQUE,
                        memory_text VARCHAR NOT NULL,
                        updated_at DATETIME DEFAULT CURRENT_TIMESTAMP
                    )
                """)
            except Exception:
                pass
            conn.commit()
            conn.close()
        except Exception as e:
            logger.error(f"Failed to auto-migrate SQLite database: {e}")

    # Start the background alert daemon
    alert_thread = threading.Thread(target=_alert_daemon, daemon=True)
    alert_thread.start()


from sqlalchemy import text

@app.get('/health', tags=['System'])
async def health_check(db: Session = Depends(get_db)):
    health_status = {'status': 'ok', 'redis': 'ok', 'postgres': 'ok'}
    try:
        if redis_client:
            redis_client.ping()
        else:
            health_status['redis'] = 'disabled'
    except Exception as e:
        health_status['redis'] = f'error: {e}'
        health_status['status'] = 'degraded'
    
    try:
        db.execute(text('SELECT 1'))
    except Exception as e:
        health_status['postgres'] = f'error: {e}'
        health_status['status'] = 'degraded'
        
    if health_status['status'] == 'degraded':
        raise HTTPException(status_code=503, detail=health_status)
        
    return health_status




from sqlalchemy import text

@app.get('/health', tags=['System'])
async def health_check(db: Session = Depends(get_db)):
    health_status = {'status': 'ok', 'redis': 'ok', 'postgres': 'ok'}
    try:
        if redis_client:
            redis_client.ping()
        else:
            health_status['redis'] = 'disabled'
    except Exception as e:
        health_status['redis'] = f'error: {e}'
        health_status['status'] = 'degraded'
    
    try:
        db.execute(text('SELECT 1'))
    except Exception as e:
        health_status['postgres'] = f'error: {e}'
        health_status['status'] = 'degraded'
        
    if health_status['status'] == 'degraded':
        raise HTTPException(status_code=503, detail=health_status)
        
    return health_status



@app.on_event("shutdown")
def shutdown_event():
    logger.info("Shutting down application...")
    logger.info("Shutting down process pool executor...")
    process_pool.shutdown(wait=True)
    
    try:
        if redis_client:
            redis_client.close()
            logger.info("Redis connection closed.")
    except Exception as e:
        logger.error(f"Failed closing Redis: {e}")
        
    try:
        from database import engine
        engine.dispose()
        logger.info("Database engine disposed.")
    except Exception as e:
        logger.error(f"Failed disposing DB engine: {e}")
        
    logger.info("Application shutdown complete.")



if __name__ == "__main__":
    import uvicorn
    # Use reload=False by default to prevent thread killing on file writes (like tasks.json or sqlite)
    # If reload is needed, exclude the output directory and database files.
    should_reload = os.getenv("DEBUG_RELOAD") == "true"
    uvicorn.run("app:app", host="0.0.0.0", port=8000, reload=should_reload, reload_excludes=["output/*", "*.db", "*.sqlite3"])


class AdminClearRequest(BaseModel):
    admin_key: str
    confirm_token: str = ""
    confirm_token: str = ""
    confirm_token: str = ""
    confirm_token: str = ""

@app.post("/api/admin/clear_backtests")
def admin_clear_backtests(req: AdminClearRequest, db: Session = Depends(get_db)):
    if req.admin_key != access_manager.MASTER_KEY:
        raise HTTPException(status_code=401, detail="Invalid Master Key")
    count = db.query(BacktestHistory).delete()
    db.commit()
    return {"status": "success", "deleted_count": count}

# ─────────────────────────────────────────────
# ADVANCED QUANTITATIVE FEATURES (OPTIONS)
# ─────────────────────────────────────────────



# ─────────────────────────────────────────────
# ADVANCED QUANTITATIVE FEATURES (STAT ARB)
# ─────────────────────────────────────────────

class StatArbRequest(BaseModel):
    tickers: List[str]
    p_value_threshold: float = 0.05
    run_backtest: bool = True

@app.post("/api/statarb/scan")
def scan_stat_arb(req: StatArbRequest, x_access_key: Optional[str] = Header(None)):
    if not access_manager.validate_key(x_access_key, silent=True):
        raise HTTPException(status_code=401, detail="Unauthorized")
    try:
        from stat_arb import find_cointegrated_pairs
        from backtest import run_stat_arb_backtest
        import yfinance as yf
        
        # 1. Fetch Data
        hist_data = yf.download(req.tickers, period="2y", timeout=5)['Close']
        if hist_data.empty:
            raise HTTPException(status_code=400, detail="Failed to fetch historical data")
            
        # Clean data (drop columns with all NaNs and forward fill)
        hist_data = hist_data.dropna(axis=1, how='all').ffill()
        
        # 2. Find Pairs
        pairs = find_cointegrated_pairs(hist_data, p_value_threshold=req.p_value_threshold)
        
        if not pairs:
            return {"pairs": [], "message": "No cointegrated pairs found."}
            
        # 3. Optional Backtest on top pair
        top_pair_result = None
        if req.run_backtest and pairs:
            best_pair = pairs[0]
            t1, t2 = best_pair["pair"]
            top_pair_result = run_stat_arb_backtest(hist_data, t1, t2, best_pair["hedge_ratio"])
            
        return {
            "pairs": pairs,
            "top_pair_backtest": top_pair_result
        }
    except Exception as e:
        logger.error(f"Error running stat arb scan: {e}")
        raise HTTPException(status_code=500, detail=str(e))

# ─────────────────────────────────────────────
# ADVANCED QUANTITATIVE FEATURES (CRYPTO ARB)
# ─────────────────────────────────────────────

class CryptoArbRequest(BaseModel):
    symbol: str = "BTC/USDT"
    capital: float = 10000.0

@app.post("/api/cryptoarb/scan")
def scan_crypto_arb(req: CryptoArbRequest, x_access_key: Optional[str] = Header(None)):
    if not access_manager.validate_key(x_access_key, silent=True):
        raise HTTPException(status_code=401, detail="Unauthorized")
    try:
        from crypto_arb import fetch_exchange_prices, find_arbitrage_opportunities
        from execution import execute_crypto_arbitrage
        
        # 1. Fetch live prices
        prices = fetch_exchange_prices(req.symbol)
        
        # 2. Find opportunities
        opps = find_arbitrage_opportunities(prices)
        
        # 3. Simulate execution for the best one
        execution_result = None
        if opps:
            best_opp = opps[0]
            execution_result = execute_crypto_arbitrage(best_opp, req.capital)
            
        if execution_result and 'profit_usd' not in execution_result:
            execution_result['profit_usd'] = execution_result.get('net_profit_usd', 0)
            
        return {
            "symbol": req.symbol,
            "prices": prices,
            "opportunities": opps,
            "execution": execution_result
        }
    except Exception as e:
        logger.error(f"Error running crypto arb scan: {e}")
        raise HTTPException(status_code=500, detail=str(e))




from pydantic import BaseModel
class EggDiscoverRequest(BaseModel):
    egg_id: str

@app.get("/api/eggs")
def get_eggs(username: str = Depends(get_current_user), db: Session = Depends(get_db)):
    import database
    progress = db.query(database.UserEggProgress).filter(database.UserEggProgress.username == username).first()
    if progress:
        return {"discovered_eggs": progress.discovered_eggs, "vault_unlocked": bool(progress.vault_unlocked)}
    return {"discovered_eggs": [], "vault_unlocked": False}

@app.post("/api/eggs/discover")
def discover_egg(req: EggDiscoverRequest, username: str = Depends(get_current_user), db: Session = Depends(get_db)):
    import database
    progress = db.query(database.UserEggProgress).filter(database.UserEggProgress.username == username).first()
    if not progress:
        progress = database.UserEggProgress(username=username, discovered_eggs=[req.egg_id], vault_unlocked=0)
        db.add(progress)
    else:
        eggs = list(progress.discovered_eggs)
        if req.egg_id not in eggs:
            eggs.append(req.egg_id)
            progress.discovered_eggs = eggs
    db.commit()
    return {"status": "success"}

@app.post("/api/eggs/unlock_vault")
def unlock_vault(username: str = Depends(get_current_user), db: Session = Depends(get_db)):
    import database
    progress = db.query(database.UserEggProgress).filter(database.UserEggProgress.username == username).first()
    if not progress:
        progress = database.UserEggProgress(username=username, discovered_eggs=[], vault_unlocked=1)
        db.add(progress)
    else:
        progress.vault_unlocked = 1
    db.commit()
    return {"status": "success"}