app/main.py

import os
import json
import logging
import asyncio
from typing import Optional
from datetime import datetime
from fastapi import FastAPI, Query, HTTPException, WebSocket, WebSocketDisconnect, Depends
from fastapi.middleware.cors import CORSMiddleware
from sqlalchemy.orm import Session
from dotenv import load_dotenv

import sys
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

load_dotenv(os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), ".env"))

from routes.copilot import router as copilot_router
from .load_data import load_all_csv_data
from .database.db import engine, Base, get_db
from .database import models  # noqa — registers tables
from .database.services import MarketDataService, NewsService, seed_postgres
from .database.repositories import MarketNewsRepo

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = FastAPI(title="MiraiTrade AI — Backend API", version="2.0.0")

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

app.include_router(copilot_router)

# ─── PostgreSQL init ──────────────────────────────────────────────────────────
try:
    Base.metadata.create_all(bind=engine)
    logger.info("PostgreSQL tables ready.")
except Exception as e:
    logger.warning(f"PostgreSQL table creation skipped: {e}")

# ─── CSV fallback store (used only by legacy routes NOT yet migrated) ─────────
csv_store = load_all_csv_data()

# ─── Seed PostgreSQL on startup ───────────────────────────────────────────────
@app.on_event("startup")
def startup_seed():
    try:
        from .database.db import SessionLocal
        db = SessionLocal()
        result = seed_postgres(db)
        db.close()
        logger.info(f"Startup seed: {result}")
    except Exception as e:
        logger.warning(f"Startup seed failed (will retry on next start): {e}")

# ─── WebSocket connection manager ────────────────────────────────────────────
class NewsConnectionManager:
    def __init__(self):
        self.active: list[WebSocket] = []

    async def connect(self, ws: WebSocket):
        await ws.accept()
        self.active.append(ws)

    def disconnect(self, ws: WebSocket):
        if ws in self.active:
            self.active.remove(ws)

    async def broadcast(self, message: dict):
        payload = json.dumps(message)
        dead = []
        for ws in self.active:
            try:
                await ws.send_text(payload)
            except Exception:
                dead.append(ws)
        for ws in dead:
            self.disconnect(ws)

news_manager = NewsConnectionManager()

# ─── WebSocket: live news stream ──────────────────────────────────────────────
@app.websocket("/ws/news")
async def ws_news(websocket: WebSocket):
    await news_manager.connect(websocket)
    try:
        while True:
            # Keep connection alive; server pushes via news_manager.broadcast()
            await asyncio.sleep(30)
            await websocket.send_text(json.dumps({"type": "ping"}))
    except WebSocketDisconnect:
        news_manager.disconnect(websocket)

# ─── MongoDB AI analysis cache ────────────────────────────────────────────────
_mongo_client = None

def _get_mongo_collection():
    global _mongo_client
    try:
        import pymongo
        mongo_uri = os.getenv("MONGODB_URI", "")
        if not mongo_uri:
            return None
        if _mongo_client is None:
            _mongo_client = pymongo.MongoClient(mongo_uri, serverSelectionTimeoutMS=3000)
        return _mongo_client["MiraiTrade"]["ai_analysis_cache"]
    except Exception:
        return None

def _gemini_generate(prompt: str) -> str:
    import urllib.request
    api_key = os.getenv("GEMINI_API_KEY", "")
    for model in ["gemini-2.5-flash", "gemini-2.0-flash", "gemini-1.5-flash"]:
        try:
            url = f"https://generativelanguage.googleapis.com/v1beta/models/{model}:generateContent?key={api_key}"
            body = json.dumps({"contents": [{"parts": [{"text": prompt}]}], "generationConfig": {"temperature": 0.15, "maxOutputTokens": 2048}}).encode()
            req = urllib.request.Request(url, data=body, headers={"Content-Type": "application/json"})
            with urllib.request.urlopen(req, timeout=20) as resp:
                data = json.loads(resp.read())
            return data["candidates"][0]["content"]["parts"][0]["text"]
        except Exception as e:
            logger.warning(f"[_gemini_generate] Model {model} failed: {str(e)[:80]}")
            continue
    raise RuntimeError("All Gemini model endpoints failed")

INSTRUMENTS = {
    "AAPL": {"symbol": "AAPL", "name": "Apple Inc.", "sector": "Technology", "industry": "Consumer Electronics", "exchange": "NASDAQ", "currency": "USD", "marketCap": 2980000000000, "beta": 1.21, "pe": 31.4},
    "TSLA": {"symbol": "TSLA", "name": "Tesla Inc.", "sector": "Consumer Discretionary", "industry": "Automotive", "exchange": "NASDAQ", "currency": "USD", "marketCap": 692000000000, "beta": 2.34, "pe": 68.2},
    "MSFT": {"symbol": "MSFT", "name": "Microsoft Corp.", "sector": "Technology", "industry": "Software—Infrastructure", "exchange": "NASDAQ", "currency": "USD", "marketCap": 3140000000000, "beta": 0.89, "pe": 37.8},
    "GOOG": {"symbol": "GOOG", "name": "Alphabet Inc.", "sector": "Communication Services", "industry": "Internet Content & Information", "exchange": "NASDAQ", "currency": "USD", "marketCap": 2210000000000, "beta": 1.06, "pe": 24.9},
    "IBM":  {"symbol": "IBM", "name": "IBM Corp.", "sector": "Technology", "industry": "Information Technology Services", "exchange": "NYSE", "currency": "USD", "marketCap": 218000000000, "beta": 0.71, "pe": 22.6},
    "WMT":  {"symbol": "WMT", "name": "Walmart Inc.", "sector": "Consumer Staples", "industry": "Discount Stores", "exchange": "NYSE", "currency": "USD", "marketCap": 785000000000, "beta": 0.58, "pe": 41.2},
    "UL":   {"symbol": "UL", "name": "Unilever PLC", "sector": "Consumer Staples", "industry": "Household & Personal Products", "exchange": "NYSE", "currency": "USD", "marketCap": 132000000000, "beta": 0.52, "pe": 18.3}
}

PORTFOLIOS = [
    {"id": "PORT-001", "name": "Growth Alpha", "managerId": "USR-001", "strategy": "Long/Short Equity", "benchmark": "NASDAQ", "inception": "2026-01-15", "currency": "USD", "cashBalance": 284750.00},
    {"id": "PORT-002", "name": "Defensive Income", "managerId": "USR-001", "strategy": "Dividend Growth", "benchmark": "S&P 500", "inception": "2026-02-01", "currency": "USD", "cashBalance": 142300.00}
]

RAW_HOLDINGS = [
    {"portfolioId": "PORT-001", "symbol": "AAPL", "quantity": 450, "avgCost": 172.40, "side": "LONG"},
    {"portfolioId": "PORT-001", "symbol": "TSLA", "quantity": 120, "avgCost": 198.60, "side": "LONG"},
    {"portfolioId": "PORT-001", "symbol": "MSFT", "quantity": 85,  "avgCost": 408.20, "side": "LONG"},
    {"portfolioId": "PORT-001", "symbol": "GOOG", "quantity": 200, "avgCost": 165.80, "side": "LONG"},
    {"portfolioId": "PORT-002", "symbol": "IBM",  "quantity": 300, "avgCost": 211.50, "side": "LONG"},
    {"portfolioId": "PORT-002", "symbol": "WMT",  "quantity": 600, "avgCost": 93.20,  "side": "LONG"},
    {"portfolioId": "PORT-002", "symbol": "UL",   "quantity": 800, "avgCost": 49.35,  "side": "LONG"}
]

ORDERS_STORE = [
    {"id": "ORD-5041", "portfolioId": "PORT-001", "symbol": "AAPL", "side": "BUY",  "type": "LIMIT", "quantity": 50,  "filledQty": 50,  "avgFillPrice": 172.40, "limitPrice": 173.00, "status": "FILLED",    "timeInForce": "DAY", "createdAt": "2026-07-02T09:30:00Z"},
    {"id": "ORD-5042", "portfolioId": "PORT-001", "symbol": "TSLA", "side": "BUY",  "type": "MARKET", "quantity": 30, "filledQty": 30,  "avgFillPrice": 208.10, "limitPrice": None, "status": "FILLED",    "timeInForce": "DAY", "createdAt": "2026-07-05T10:15:00Z"},
    {"id": "ORD-5043", "portfolioId": "PORT-001", "symbol": "MSFT", "side": "BUY",  "type": "LIMIT", "quantity": 20,  "filledQty": 20,  "avgFillPrice": 415.60, "limitPrice": 416.00, "status": "FILLED",    "timeInForce": "DAY", "createdAt": "2026-07-08T11:45:00Z"},
    {"id": "ORD-5044", "portfolioId": "PORT-001", "symbol": "GOOG", "side": "BUY",  "type": "MARKET", "quantity": 80, "filledQty": 80,  "avgFillPrice": 168.20, "limitPrice": None, "status": "FILLED",    "timeInForce": "DAY", "createdAt": "2026-07-10T13:20:00Z"},
    {"id": "ORD-5045", "portfolioId": "PORT-001", "symbol": "AAPL", "side": "SELL", "type": "LIMIT", "quantity": 20,  "filledQty": 20,  "avgFillPrice": 181.30, "limitPrice": 181.00, "status": "FILLED",    "timeInForce": "DAY", "createdAt": "2026-07-18T14:05:00Z"},
    {"id": "ORD-5046", "portfolioId": "PORT-001", "symbol": "TSLA", "side": "BUY",  "type": "STOP",  "quantity": 30,  "filledQty": 30,  "avgFillPrice": 221.40, "stopPrice": 220.00, "status": "FILLED",    "timeInForce": "GTC", "createdAt": "2026-07-22T09:50:00Z"},
    {"id": "ORD-5047", "portfolioId": "PORT-001", "symbol": "GOOG", "side": "SELL", "type": "LIMIT", "quantity": 15,  "filledQty": 0,   "avgFillPrice": None,  "limitPrice": 180.00, "status": "NEW",       "timeInForce": "DAY", "createdAt": "2026-08-29T15:00:00Z"},
    {"id": "ORD-5048", "portfolioId": "PORT-001", "symbol": "MSFT", "side": "BUY",  "type": "STOP_LIMIT", "quantity": 10, "filledQty": 0, "avgFillPrice": None, "stopPrice": 425.00, "limitPrice": 426.00, "status": "NEW", "timeInForce": "GTC", "createdAt": "2026-08-29T15:30:00Z"},
    {"id": "ORD-5049", "portfolioId": "PORT-001", "symbol": "TSLA", "side": "SELL", "type": "LIMIT", "quantity": 40,  "filledQty": 25,  "avgFillPrice": 214.80, "limitPrice": 215.00, "status": "PARTIAL",   "timeInForce": "DAY", "createdAt": "2026-08-29T14:30:00Z"},
    {"id": "ORD-5050", "portfolioId": "PORT-001", "symbol": "AAPL", "side": "BUY",  "type": "LIMIT", "quantity": 25,  "filledQty": 0,   "avgFillPrice": None,  "limitPrice": 175.00, "status": "CANCELLED", "timeInForce": "DAY", "createdAt": "2026-08-28T10:00:00Z"}
]

TRANSACTIONS = [
    {"id": "TXN-1001", "portfolioId": "PORT-001", "symbol": "AAPL", "side": "BUY",  "quantity": 100, "price": 172.40, "fee": 4.95, "date": "2026-07-02", "status": "SETTLED"},
    {"id": "TXN-1002", "portfolioId": "PORT-001", "symbol": "TSLA", "side": "BUY",  "quantity": 50,  "price": 208.10, "fee": 4.95, "date": "2026-07-05", "status": "SETTLED"},
    {"id": "TXN-1003", "portfolioId": "PORT-001", "symbol": "MSFT", "side": "BUY",  "quantity": 30,  "price": 415.60, "fee": 4.95, "date": "2026-07-08", "status": "SETTLED"},
    {"id": "TXN-1004", "portfolioId": "PORT-001", "symbol": "GOOG", "side": "BUY",  "quantity": 80,  "price": 168.20, "fee": 4.95, "date": "2026-07-10", "status": "SETTLED"},
    {"id": "TXN-1005", "portfolioId": "PORT-001", "symbol": "AAPL", "side": "SELL", "quantity": 20,  "price": 181.30, "fee": 4.95, "date": "2026-07-18", "status": "SETTLED"},
    {"id": "TXN-1006", "portfolioId": "PORT-001", "symbol": "TSLA", "side": "BUY",  "quantity": 30,  "price": 221.40, "fee": 4.95, "date": "2026-07-22", "status": "SETTLED"}
]

@app.get("/health")
def health_check():
    return {"status": "ok", "backend": "FastAPI", "csv_loaded": True}

@app.get("/api/instruments")
def list_instruments():
    data = list(INSTRUMENTS.values())
    return {"count": len(data), "data": data}

@app.get("/api/instruments/{symbol}")
def get_instrument(symbol: str):
    sym = symbol.upper()
    if sym in INSTRUMENTS:
        return {"data": INSTRUMENTS[sym]}
    raise HTTPException(status_code=404, detail="Instrument not found")

@app.get("/api/market/summary")
def market_summary(db: Session = Depends(get_db)):
    try:
        summary = MarketDataService.get_market_summary(db)
        if summary:
            return {"count": len(summary), "data": summary}
    except Exception as e:
        logger.warning(f"PG market summary failed, using CSV fallback: {e}")

    # CSV fallback
    summary = []
    prices = csv_store["prices"]
    historical = csv_store["historical"]
    for sym in INSTRUMENTS:
        tick_data = prices.get(sym) or historical.get(sym) or []
        if tick_data and len(tick_data) >= 1:
            latest = tick_data[-1]
            prev = tick_data[-2] if len(tick_data) >= 2 else latest
            price = latest["close"]
            change = round(price - prev["close"], 2)
            change_pct = round((change / prev["close"] * 100), 2) if prev["close"] else 0.0
            summary.append({
                "symbol": sym,
                "price": price,
                "change": change,
                "changePercent": change_pct,
                "volume": latest.get("volume", 0),
                "vwap": latest.get("vwap", price),
            })
    return {"count": len(summary), "data": summary}


@app.get("/api/market/symbols/{symbol}")
def market_symbol(symbol: str, period: str = "3M", db: Session = Depends(get_db)):
    sym = symbol.upper()
    inst = INSTRUMENTS.get(sym)

    try:
        historical = MarketDataService.get_historical(db, sym, limit=500)
        live = MarketDataService.get_live_ticks(db, sym, limit=1000)
        series = live or historical
    except Exception as e:
        logger.warning(f"PG market symbol failed, using CSV fallback: {e}")
        series = csv_store["prices"].get(sym) or csv_store["historical"].get(sym) or []
        historical = series

    latest = series[-1] if series else {"close": 100, "open": 100, "high": 100, "low": 100, "volume": 0}
    prev = series[-2] if len(series) >= 2 else latest
    close_key = "close"
    change = round(latest[close_key] - prev[close_key], 2)
    change_pct = round((change / prev[close_key] * 100), 2) if prev[close_key] else 0.0

    try:
        news_items = NewsService.get_feed(db, symbol=sym, limit=10)
        news_items = [
            {
                "id": n["id"],
                "headline": n["headline"],
                "summary": n["summary"],
                "symbols": [ts["ticker"] for ts in n.get("tickers", [])],
                "sentiment": n["sentiment"],
                "confidence": n.get("confidence", ""),
                "timestamp": n.get("published_at"),
                "category": (n.get("topics") or [{}])[0].get("topic", "General"),
                "source": n.get("source", ""),
            }
            for n in news_items
        ]
    except Exception:
        news_items = [n for n in csv_store["news"] if sym in n.get("symbols", [])]

    return {
        "data": {
            "instrument": inst,
            "latest": latest,
            "change": {"change": change, "changePercent": change_pct},
            "historical": historical,
            "news": news_items,
        }
    }

@app.get("/api/portfolio")
def list_portfolios():
    return {"count": len(PORTFOLIOS), "data": PORTFOLIOS}

@app.get("/api/portfolio/{id}/summary")
def portfolio_summary(id: str):
    pid = "PORT-001" if id in ["active", "PORT-001"] else id
    port = next((p for p in PORTFOLIOS if p["id"] == pid), PORTFOLIOS[0])
    return {
        "data": {
            **port,
            "totalValue": 450000.0,
            "holdingsValue": 165250.0,
            "totalCost": 150000.0,
            "unrealizedPL": 15250.0,
            "unrealizedPLPct": 10.17,
            "todayPL": 1250.0,
            "todayPLPct": 0.83,
            "cashPct": 63.3,
            "positionCount": 4
        }
    }

@app.get("/api/portfolio/{id}/holdings")
def portfolio_holdings(id: str):
    pid = "PORT-001" if id in ["active", "PORT-001"] else id
    h_list = [h for h in RAW_HOLDINGS if h["portfolioId"] == pid]
    res = []
    prices = csv_store["prices"]
    historical = csv_store["historical"]

    for h in h_list:
        sym = h["symbol"]
        ticks = prices.get(sym) or historical.get(sym) or []
        price = ticks[-1]["close"] if ticks else h["avgCost"]
        mkt_val = price * h["quantity"]
        cost = h["avgCost"] * h["quantity"]
        pl = mkt_val - cost
        pl_pct = round((pl / cost * 100), 2)
        res.append({
            **h,
            "currentPrice": price,
            "bid": round(price * 0.999, 2),
            "ask": round(price * 1.001, 2),
            "mid": price,
            "marketValue": mkt_val,
            "costBasis": cost,
            "unrealizedPL": pl,
            "unrealizedPLPct": pl_pct,
            "todayPL": round(price * 0.005 * h["quantity"], 2),
            "todayPLPct": 0.5,
            "dayChange": 0.5,
            "volume": ticks[-1].get("volume", 0) if ticks else 0
        })
    return {"count": len(res), "data": res}

@app.get("/api/portfolio/{id}/transactions")
def portfolio_transactions(id: str):
    pid = "PORT-001" if id in ["active", "PORT-001"] else id
    txns = [t for t in TRANSACTIONS if t["portfolioId"] == pid]
    return {"count": len(txns), "data": txns}

@app.get("/api/orders")
def list_orders(portfolio_id: Optional[str] = Query(None), portfolioId: Optional[str] = Query(None), status: Optional[str] = None):
    pid = portfolio_id or portfolioId or "PORT-001"
    if pid == "active":
        pid = "PORT-001"
    orders = [o for o in ORDERS_STORE if o.get("portfolioId") == pid or pid is None]
    if status:
        orders = [o for o in orders if o.get("status") == status]
    return {"count": len(orders), "data": orders}

@app.get("/api/orders/{id}")
def get_order(id: str):
    order = next((o for o in ORDERS_STORE if o["id"] == id), None)
    if not order:
        raise HTTPException(status_code=404, detail="Order not found")
    events = [
        {"orderId": id, "event": "CREATED", "timestamp": order.get("createdAt"), "note": "Submitted by USR-001"},
        {"orderId": id, "event": order["status"], "timestamp": order.get("createdAt"), "note": f"State: {order['status']}"}
    ]
    return {"data": {**order, "events": events}}

@app.post("/api/orders")
def create_order(payload: dict):
    new_id = f"ORD-{5050 + len(ORDERS_STORE) + 1}"
    sym = payload.get("symbol", "AAPL")
    prices = csv_store["prices"].get(sym) or csv_store["historical"].get(sym) or []
    px = prices[-1]["close"] if prices else 180.0
    new_order = {
        "id": new_id,
        "portfolioId": payload.get("portfolioId", "PORT-001"),
        "symbol": sym,
        "side": payload.get("side", "BUY"),
        "type": payload.get("type", "LIMIT"),
        "quantity": payload.get("quantity", 10),
        "filledQty": payload.get("quantity", 10) if payload.get("type") == "MARKET" else 0,
        "avgFillPrice": px if payload.get("type") == "MARKET" else None,
        "limitPrice": payload.get("limitPrice"),
        "stopPrice": payload.get("stopPrice"),
        "status": "FILLED" if payload.get("type") == "MARKET" else "NEW",
        "timeInForce": payload.get("timeInForce", "DAY"),
        "createdAt": "2026-08-30T10:00:00Z"
    }
    ORDERS_STORE.insert(0, new_order)

    # Automatically queue into Compliance Review workflow for pre-execution surveillance
    comp_item = {
        "id": f"COMP-{9000 + len(COMPLIANCE_REVIEWS) + 1}",
        "tradeId": new_id,
        "symbol": sym,
        "side": new_order["side"],
        "quantity": new_order["quantity"],
        "price": payload.get("limitPrice") or px,
        "trader": "Active Trader",
        "status": "PENDING",
        "automatedChecks": {
            "restrictedSecurity": "PASS",
            "tradingWindow": "PASS",
            "amlKyc": "PASS",
            "sanctions": "PASS",
            "duplicateOrder": "PASS",
            "employeeTradingPolicy": "PASS"
        },
        "submittedAt": "2026-08-03T00:15:00Z"
    }
    COMPLIANCE_REVIEWS.insert(0, comp_item)

    return {"data": new_order, "events": [{"orderId": new_id, "event": "CREATED", "timestamp": "2026-08-30T10:00:00Z", "note": "Submitted & Queued for Compliance"}]}

@app.delete("/api/orders/{id}")
def cancel_order(id: str):
    order = next((o for o in ORDERS_STORE if o["id"] == id), None)
    if order:
        order["status"] = "CANCELLED"
        return {"data": order}
    raise HTTPException(status_code=404, detail="Order not found")

@app.get("/api/risk/{id}")
def get_risk(id: str, period: str = "3M"):
    return {
        "data": {
            "risk": {"var95": 14250.0, "var99": 21300.0, "sharpe": 1.84, "beta": 1.12, "maxDrawdown": -8.45, "volatility": 14.2},
            "attribution": [{"sector": "Technology", "allocation": 58.5, "return": 14.2}, {"sector": "Consumer Discretionary", "allocation": 22.1, "return": 9.8}],
            "benchmark": {"name": "NASDAQ", "ytdReturn": 16.4, "beta": 1.0},
            "timeSeries": []
        }
    }

@app.get("/api/limits/{id}")
def get_limits(id: str):
    return {
        "data": {
            "portfolio_limits": {"yearly_budget_per_symbol": 500000, "max_order_size_pct": 5, "var_limit_pct": 2, "approval_threshold_pct": 70},
            "symbol_limits": []
        }
    }

@app.get("/api/limits/{id}/{symbol}")
def get_symbol_limit(id: str, symbol: str):
    return {
        "data": {
            "symbol": symbol.upper(),
            "portfolio_id": id,
            "yearly_budget": 500000,
            "utilized": 150000,
            "remaining": 350000,
            "utilization_pct": 30.0,
            "approval_required": False,
            "approval_threshold_pct": 70
        }
    }

@app.get("/api/news")
def get_news(
    symbol: Optional[str] = None,
    topic: Optional[str] = None,
    sentiment: Optional[str] = None,
    limit: int = 50,
    offset: int = 0,
    db: Session = Depends(get_db),
):
    items = NewsService.get_feed(db, symbol=symbol, topic=topic, sentiment=sentiment, limit=limit, offset=offset)
    # Normalise to legacy shape so existing frontend code keeps working
    normalized = []
    for n in items:
        tickers = [ts["ticker"] for ts in n.get("tickers", []) if ts.get("ticker")]
        first_topic = (n.get("topics") or [{}])[0].get("topic", "General")
        normalized.append({
            "id": n["id"],
            "db_id": n.get("db_id"),
            "headline": n["headline"],
            "summary": n["summary"],
            "source": n["source"],
            "url": n.get("url", ""),
            "timestamp": n.get("published_at"),
            "symbols": tickers if tickers else ["AAPL"],
            "sentiment": n["sentiment"],
            "confidence": f"{int((n.get('confidence_score') or 0.8) * 100)}%",
            "confidence_score": n.get("confidence_score"),
            "importance_score": n.get("importance_score"),
            "overall_sentiment_score": n.get("overall_sentiment_score"),
            "category": first_topic,
            "topics": n.get("topics", []),
            "tickers": n.get("tickers", []),
            "is_breaking": n.get("is_breaking", False),
        })
    return {"count": len(normalized), "data": normalized}


@app.get("/api/news/intelligence")
def news_intelligence(db: Session = Depends(get_db)):
    """Aggregate Market Intelligence panel data."""
    return {
        "trending_topics": NewsService.get_trending_topics(db, hours=24),
        "ticker_sentiment": NewsService.get_ticker_sentiment_summary(db, hours=24),
        "news_velocity": NewsService.get_news_velocity(db, hours=6),
        "sector_sentiment": NewsService.get_sector_sentiment(db, hours=24),
        "breaking_news": NewsService.get_breaking_news(db),
    }


@app.post("/api/news/{news_id}/ai-analysis")
async def get_ai_analysis(news_id: str, db: Session = Depends(get_db)):
    """
    On-demand AI analysis via Gemini.
    Returns cached result from MongoDB if available; otherwise generates and caches.
    """
    col = _get_mongo_collection()

    # Check MongoDB cache first
    if col is not None:
        cached = col.find_one({"news_id": news_id}, {"_id": 0})
        if cached:
            return {"cached": True, "data": cached}

    # Fetch news item from PostgreSQL
    item = NewsService.get_by_id(db, news_id)
    if not item:
        raise HTTPException(status_code=404, detail="News item not found")

    tickers = ", ".join([t["ticker"] for t in item.get("tickers", [])]) or "market"
    prompt = f"""You are an institutional equity research analyst. Analyze this market news article and provide a structured JSON response ONLY (no markdown, no prose outside the JSON).

ARTICLE:
Headline: {item['headline']}
Summary: {item.get('summary', '')}
Source: {item.get('source', '')}
Published: {item.get('published_at', '')}
Tickers: {tickers}
Current Sentiment: {item.get('sentiment', 'neutral')} ({item.get('overall_sentiment_score', 0):.2f})

Return exactly this JSON structure:
{{
  "executive_summary": "2-3 sentence summary",
  "market_impact": "immediate market impact explanation",
  "sentiment": "Bullish" | "Neutral" | "Bearish",
  "sentiment_reasoning": "why this sentiment",
  "intraday_outlook": "price direction next few hours",
  "one_day_outlook": "1-day price impact",
  "one_week_outlook": "1-week trend",
  "confidence_score": 0.0-1.0,
  "key_drivers": ["driver1", "driver2", "driver3"],
  "risks": ["risk1", "risk2"],
  "opportunities": ["opp1", "opp2"],
  "sector_impact": [{{"sector": "...", "impact": "positive|negative|neutral", "explanation": "..."}}],
  "ticker_impact": [{{"ticker": "...", "price_direction": "up|down|sideways", "magnitude": "low|medium|high", "reason": "..."}}]
}}"""

    try:
        raw = _gemini_generate(prompt)
        # Extract JSON from response
        import re
        json_match = re.search(r'\{[\s\S]+\}', raw)
        if not json_match:
            raise ValueError("No JSON in Gemini response")
        analysis = json.loads(json_match.group())
    except Exception as e:
        logger.warning(f"Gemini analysis failed for {news_id}: {e}")
        analysis = {
            "executive_summary": item.get("summary", item["headline"]),
            "market_impact": "Analysis unavailable — AI service error.",
            "sentiment": item.get("sentiment", "neutral").capitalize(),
            "sentiment_reasoning": "Based on source data sentiment score.",
            "intraday_outlook": "Insufficient data for intraday projection.",
            "one_day_outlook": "Monitor price action closely.",
            "one_week_outlook": "Dependent on broader market conditions.",
            "confidence_score": item.get("confidence_score", 0.7),
            "key_drivers": [],
            "risks": [],
            "opportunities": [],
            "sector_impact": [],
            "ticker_impact": [],
        }

    result = {
        "news_id": news_id,
        "headline": item["headline"],
        "generated_at": datetime.utcnow().isoformat(),
        **analysis,
    }

    # Cache in MongoDB
    if col is not None:
        try:
            col.replace_one({"news_id": news_id}, result, upsert=True)
        except Exception as ce:
            logger.warning(f"MongoDB cache write failed: {ce}")

    return {"cached": False, "data": result}

USERS_STORE = [
    {"id": "USR-001", "name": "Sam Trader", "email": "trader@mirai.ai", "role": "trader", "status": "active", "pnl": 8500.0, "portfolioValue": 450000.0},
    {"id": "USR-002", "name": "Aarav Mehta", "email": "aarav@mirai.ai", "role": "trader", "status": "active", "pnl": 24500.0, "portfolioValue": 1250000.0},
    {"id": "USR-003", "name": "Priya Sharma", "email": "priya@mirai.ai", "role": "trader", "status": "active", "pnl": 18200.0, "portfolioValue": 850000.0},
    {"id": "USR-004", "name": "Vikram Singh", "email": "vikram@mirai.ai", "role": "trader", "status": "active", "pnl": 12400.0, "portfolioValue": 650000.0},
    {"id": "USR-005", "name": "Elena Rostova", "email": "elena@mirai.ai", "role": "manager", "status": "active"},
    {"id": "USR-006", "name": "David Vance", "email": "david@mirai.ai", "role": "risk", "status": "active"},
    {"id": "USR-007", "name": "Rachel Zane", "email": "rachel@mirai.ai", "role": "compliance", "status": "active"},
    {"id": "USR-008", "name": "Alexander Pierce", "email": "admin@mirai.ai", "role": "admin", "status": "active"}
]

@app.get("/api/users")
def get_all_users():
    return {"users": USERS_STORE}

@app.get("/api/users/traders")
def get_traders():
    traders = [u for u in USERS_STORE if u["role"] == "trader"]
    return {"count": len(traders), "data": traders}

@app.get("/api/stp/metrics")
def get_stp_metrics():
    return {
        "success_rate": 98.75,
        "avg_latency_ms": 38,
        "total_trades_today": 1347,
        "manual_interventions": 3,
        "settlement_success": 99.2
    }

# ─── WORKFLOW QUEUES & STORES FOR COMPLIANCE, RISK, LEGAL ───────────────────

COMPLIANCE_REVIEWS = [
    {
        "id": "COMP-9001",
        "tradeId": "TRD-8841",
        "symbol": "AAPL",
        "side": "BUY",
        "quantity": 1500,
        "price": 194.50,
        "trader": "Aarav Mehta",
        "status": "PENDING",
        "automatedChecks": {
            "restrictedSecurity": "PASS",
            "tradingWindow": "PASS",
            "amlKyc": "PASS",
            "sanctions": "PASS",
            "duplicateOrder": "PASS",
            "employeeTradingPolicy": "PASS"
        },
        "submittedAt": "2026-08-02T10:15:00Z"
    },
    {
        "id": "COMP-9002",
        "tradeId": "TRD-8842",
        "symbol": "MSFT",
        "side": "BUY",
        "quantity": 1000,
        "price": 412.30,
        "trader": "Priya Sharma",
        "status": "PENDING",
        "automatedChecks": {
            "restrictedSecurity": "PASS",
            "tradingWindow": "PASS",
            "amlKyc": "PASS",
            "sanctions": "PASS",
            "duplicateOrder": "PASS",
            "employeeTradingPolicy": "WARN"
        },
        "submittedAt": "2026-08-02T10:20:00Z"
    }
]

RISK_REVIEWS = [
    {
        "id": "RSK-7001",
        "tradeId": "TRD-8839",
        "symbol": "TSLA",
        "side": "BUY",
        "quantity": 2500,
        "price": 240.00,
        "trader": "Vikram Singh",
        "status": "PENDING",
        "checks": {
            "positionLimit": "PASS",
            "exposureLimit": "HIGH_RISK",
            "sectorExposure": "PASS",
            "stockConcentration": "BREACH",
            "marketVolatility": "PASS",
            "liquidityCheck": "PASS",
            "marginRequirement": "PASS",
            "counterpartyRisk": "PASS",
            "varCheck": "HIGH_RISK",
            "stressTest": "FAIL",
            "stopLossRequirement": "PASS",
            "leverageLimit": "PASS"
        },
        "requiresLegal": False,
        "submittedAt": "2026-08-02T09:45:00Z"
    }
]

LEGAL_REVIEWS = [
    {
        "id": "LGL-5001",
        "tradeId": "TRD-8835",
        "symbol": "NVDA",
        "side": "BUY",
        "quantity": 3000,
        "price": 125.00,
        "trader": "Sam Trader",
        "status": "PENDING",
        "checks": {
            "insiderTrading": "PASS",
            "conflictOfInterest": "REVIEW_NEEDED",
            "restrictedProjectAccess": "PASS",
            "chineseWallViolation": "PASS",
            "regulatoryCompliance": "PASS"
        },
        "submittedAt": "2026-08-02T08:30:00Z"
    }
]

RESTRICTED_SECURITIES_STORE = [
    {"symbol": "GME", "reason": "Meme stock volatility lockout", "added": "2026-07-15"},
    {"symbol": "AMC", "reason": "Excessive margin volatility", "added": "2026-07-20"},
    {"symbol": "INSIDER_RESTRICTED", "reason": "M&A Blackout Window", "added": "2026-08-01"}
]

EXPOSURE_MONITORING_STORE = [
    {"sector": "Banking", "exposure": 40.0, "limit": 50.0, "utilized": 80.0, "status": "NORMAL"},
    {"sector": "IT", "exposure": 58.0, "limit": 60.0, "utilized": 97.0, "status": "NEAR_LIMIT"},
    {"sector": "Pharma", "exposure": 15.0, "limit": 40.0, "utilized": 37.5, "status": "NORMAL"},
    {"sector": "Energy", "exposure": 28.0, "limit": 30.0, "utilized": 93.3, "status": "NEAR_LIMIT"}
]

@app.get("/api/workflow/compliance")
def get_compliance_queue():
    return {"count": len(COMPLIANCE_REVIEWS), "data": COMPLIANCE_REVIEWS, "restricted": RESTRICTED_SECURITIES_STORE}

@app.post("/api/workflow/compliance/{id}/decision")
def compliance_decision(id: str, payload: dict):
    decision = payload.get("decision", "APPROVED")
    review = next((r for r in COMPLIANCE_REVIEWS if r["id"] == id), None)
    if not review:
        raise HTTPException(status_code=404, detail="Compliance review item not found")
    review["status"] = decision
    if decision == "APPROVED":
        # Pass forward to Risk Queue as per workflow
        risk_item = {
            "id": f"RSK-{8000 + len(RISK_REVIEWS) + 1}",
            "tradeId": review["tradeId"],
            "symbol": review["symbol"],
            "side": review["side"],
            "quantity": review["quantity"],
            "price": review["price"],
            "trader": review["trader"],
            "status": "PENDING",
            "checks": {
                "positionLimit": "PASS",
                "exposureLimit": "PASS",
                "sectorExposure": "PASS",
                "stockConcentration": "PASS",
                "marketVolatility": "PASS",
                "liquidityCheck": "PASS",
                "marginRequirement": "PASS",
                "counterpartyRisk": "PASS",
                "varCheck": "PASS",
                "stressTest": "PASS",
                "stopLossRequirement": "PASS",
                "leverageLimit": "PASS"
            },
            "requiresLegal": False,
            "submittedAt": "2026-08-02T11:00:00Z"
        }
        RISK_REVIEWS.insert(0, risk_item)
    return {"message": f"Compliance decision {decision} saved", "data": review}

@app.get("/api/workflow/risk")
def get_risk_queue():
    return {"count": len(RISK_REVIEWS), "data": RISK_REVIEWS, "exposure": EXPOSURE_MONITORING_STORE}

@app.post("/api/workflow/risk/{id}/decision")
def risk_decision(id: str, payload: dict):
    decision = payload.get("decision", "APPROVED")
    review = next((r for r in RISK_REVIEWS if r["id"] == id), None)
    if not review:
        raise HTTPException(status_code=404, detail="Risk review item not found")
    review["status"] = decision
    review["requiresLegal"] = False
    if decision == "APPROVED":
        settlement_item = {
            "settlementId": f"STL-{Date.now() if 'Date' in globals() else '7001'}",
            "tradeId": review["tradeId"],
            "symbol": review["symbol"],
            "side": review["side"],
            "quantity": review["quantity"],
            "price": review["price"],
            "trader": review["trader"],
            "status": "SETTLED"
        }
        # Trade moves directly from risk to settlement
    return {"message": f"Risk decision {decision} logged and sent to settlement", "data": review}

@app.get("/api/workflow/legal")
def get_legal_queue():
    return {"count": len(LEGAL_REVIEWS), "data": LEGAL_REVIEWS}

@app.post("/api/workflow/legal/{id}/decision")
def legal_decision(id: str, payload: dict):
    decision = payload.get("decision", "APPROVED")
    review = next((r for r in LEGAL_REVIEWS if r["id"] == id), None)
    if not review:
        raise HTTPException(status_code=404, detail="Legal review item not found")
    review["status"] = decision
    return {"message": f"Legal decision {decision} processed", "data": review}

# ─── EXTRA ENDPOINTS FOR COMPLIANCE, RISK & LEGAL FEATURES ───────────────────

@app.get("/api/compliance/restricted")
def get_restricted_securities():
    return {"count": len(RESTRICTED_SECURITIES_STORE), "data": RESTRICTED_SECURITIES_STORE}

@app.post("/api/compliance/restricted")
def add_restricted_security(payload: dict):
    sym = payload.get("symbol", "").upper()
    reason = payload.get("reason", "Manual compliance lock")
    if not sym:
        raise HTTPException(status_code=400, detail="Symbol required")
    item = {"symbol": sym, "reason": reason, "added": "2026-08-03"}
    RESTRICTED_SECURITIES_STORE.insert(0, item)
    return {"message": f"Security {sym} added to restricted list", "data": item}

@app.get("/api/risk/exposure")
def get_sector_exposure():
    return {"count": len(EXPOSURE_MONITORING_STORE), "data": EXPOSURE_MONITORING_STORE}

@app.get("/api/risk/var-metrics")
def get_var_metrics():
    return {
        "dailyVaR": 14250.0,
        "weeklyVaR": 31800.0,
        "monthlyVaR": 68500.0,
        "breachedLimit": False,
        "historicalLoss": "-4.2% Max (2020-03)"
    }

@app.get("/api/legal/cases")
def get_legal_cases():
    return {
        "cases": LEGAL_REVIEWS,
        "insiderSurveillance": [
          {"ticker": "AAPL", "status": "CLEARED", "alert": "Pre-earnings trade screened"},
          {"ticker": "INSIDER_RESTRICTED", "status": "BLACKOUT", "alert": "M&A blackout window enforced"}
        ],
        "chineseWall": {"status": "ACTIVE", "breaches": 0, "activeProjects": ["Project Titan", "M&A Alpha"]}
    }

# ─── INSTITUTION ONBOARDING & STP EXTENSIONS ───────────────────────────────

INSTITUTIONS_STORE = [
    {
        "id": "INST-1001",
        "name": "Global Macro Fund LP",
        "lei": "5493006MHB84DD0ZWV18",
        "country": "US",
        "entityType": "Hedge Fund",
        "aum": 2500000000,
        "status": "ACTIVE",
        "mandate": "Global Equities",
        "permissions": ["Long", "Short", "Margin"],
        "assetClasses": ["Equities", "FX"],
        "tradingLimits": {"dailyVaR": 5000000, "maxLeverage": 2.5},
        "submittedAt": "2026-07-15T08:00:00Z"
    }
]

BASKET_ORDERS_STORE = []
CHILD_ORDERS_STORE = []

@app.get("/api/institutions")
def get_institutions():
    return {"count": len(INSTITUTIONS_STORE), "data": INSTITUTIONS_STORE}

@app.post("/api/institutions")
def create_institution(payload: dict):
    new_id = f"INST-{1000 + len(INSTITUTIONS_STORE) + 1}"
    new_inst = {
        "id": new_id,
        "name": payload.get("name", "Unknown"),
        "lei": payload.get("lei", ""),
        "country": payload.get("country", ""),
        "entityType": payload.get("entityType", ""),
        "aum": payload.get("aum", 0),
        "status": "DRAFT",
        "mandate": payload.get("mandate", ""),
        "permissions": payload.get("permissions", []),
        "assetClasses": payload.get("assetClasses", []),
        "tradingLimits": payload.get("tradingLimits", {}),
        "submittedAt": "2026-08-04T10:00:00Z"
    }
    INSTITUTIONS_STORE.insert(0, new_inst)
    return {"message": "Draft created", "data": new_inst}

@app.post("/api/institutions/{id}/approve")
def approve_institution(id: str, payload: dict):
    step = payload.get("step") # "LEGAL", "COMPLIANCE", "RISK", "ADMIN_ACTIVATE"
    decision = payload.get("decision", "APPROVED")
    
    inst = next((i for i in INSTITUTIONS_STORE if i["id"] == id), None)
    if not inst:
        raise HTTPException(status_code=404, detail="Institution not found")
        
    if decision == "REJECTED":
        inst["status"] = "REJECTED"
        return {"message": "Institution rejected", "data": inst}
        
    if step == "LEGAL":
        inst["status"] = "PENDING_COMPLIANCE"
    elif step == "COMPLIANCE":
        inst["status"] = "PENDING_RISK"
    elif step == "RISK":
        inst["status"] = "APPROVED"
    elif step == "ADMIN_ACTIVATE":
        inst["status"] = "ACTIVE"
        
    return {"message": f"Step {step} marked as {decision}", "data": inst}

@app.post("/api/orders/basket")
def create_basket_order(payload: dict):
    new_id = f"BASKET-{6000 + len(BASKET_ORDERS_STORE) + 1}"
    basket = {
        "id": new_id,
        "institutionId": payload.get("institutionId"),
        "basketName": payload.get("basketName", "Untitled Basket"),
        "executionDeadline": payload.get("executionDeadline"),
        "priority": payload.get("priority", "NORMAL"),
        "style": payload.get("style", "VWAP"),
        "status": "PENDING_ALLOCATION",
        "createdAt": "2026-08-04T10:00:00Z",
        "orders": payload.get("orders", [])
    }
    BASKET_ORDERS_STORE.insert(0, basket)
    return {"message": "Basket created", "data": basket}

@app.get("/api/orders/basket")
def get_basket_orders():
    return {"count": len(BASKET_ORDERS_STORE), "data": BASKET_ORDERS_STORE}

@app.post("/api/orders/basket/{id}/allocate")
def allocate_basket(id: str, payload: dict):
    basket = next((b for b in BASKET_ORDERS_STORE if b["id"] == id), None)
    if not basket:
        raise HTTPException(status_code=404, detail="Basket not found")
        
    trader_id = payload.get("traderId")
    
    child_orders = []
    for order in basket["orders"]:
        child_id = f"CHILD-{7000 + len(CHILD_ORDERS_STORE) + 1}"
        child = {
            "id": child_id,
            "basketId": basket["id"],
            "traderId": trader_id,
            "symbol": order.get("symbol"),
            "side": order.get("side"),
            "quantity": order.get("quantity"),
            "status": "ASSIGNED",
            "filledQty": 0,
            "avgFillPrice": None,
            "createdAt": "2026-08-04T10:05:00Z"
        }
        child_orders.append(child)
        CHILD_ORDERS_STORE.insert(0, child)
        
    basket["status"] = "ALLOCATED"
    return {"message": "Allocated to trader", "data": child_orders}

@app.get("/api/orders/child")
def get_child_orders(traderId: Optional[str] = None):
    orders = CHILD_ORDERS_STORE
    if traderId:
        orders = [o for o in orders if o["traderId"] == traderId]
    return {"count": len(orders), "data": orders}

@app.post("/api/orders/child/{id}/status")
def update_child_order_status(id: str, payload: dict):
    status = payload.get("status")
    order = next((o for o in CHILD_ORDERS_STORE if o["id"] == id), None)
    if not order:
        raise HTTPException(status_code=404, detail="Child order not found")
        
    order["status"] = status
    if status == "COMPLETED":
        order["filledQty"] = order["quantity"]
        order["avgFillPrice"] = 150.0 
        
    return {"message": f"Status updated to {status}", "data": order}

@app.get("/api/manager/ai-recommendation")
def ai_recommendation(basketId: str):
    return {
        "data": {
            "recommendedTrader": "USR-002",
            "name": "Aarav Mehta",
            "reason": "High historical execution quality (98%) for Tech sector and current low workload.",
            "metrics": {
                "historicalSlippage": "1.2 bps",
                "workload": "Low",
                "sectorExpertise": ["Technology", "Consumer Discretionary"]
            }
        }
    }