import os, time, math, datetime as dt
from typing import Dict, Any, List

import numpy as np
import httpx
from fastapi import FastAPI, Body
from fastapi.responses import HTMLResponse
from fastapi.staticfiles import StaticFiles

FINNHUB_API_KEY = os.getenv("FINNHUB_API_KEY", "").strip()
TWELVEDATA_API_KEY = os.getenv("TWELVEDATA_API_KEY", "").strip()

HF_TOKEN = os.getenv("HF_TOKEN", "").strip()
HF_TEXT_MODEL = os.getenv("HF_TEXT_MODEL", "google/gemma-2-2b-it").strip()

USER_AGENT = os.getenv("USER_AGENT", "AhsanPortfolioAI/1.0 (+https://huggingface.co/spaces)")

app = FastAPI(title="Ahsan Stock Portfolio AI Center", version="3.0.0")
app.mount("/static", StaticFiles(directory="static"), name="static")


class TTLCache:
    def __init__(self, ttl_seconds: int):
        self.ttl = ttl_seconds
        self._store: Dict[str, Any] = {}

    def get(self, key: str):
        v = self._store.get(key)
        if not v:
            return None
        ts, payload = v
        if time.time() - ts > self.ttl:
            self._store.pop(key, None)
            return None
        return payload

    def set(self, key: str, payload: Any):
        self._store[key] = (time.time(), payload)


quotes_cache = TTLCache(20)
candles_cache = TTLCache(600)
news_cache = TTLCache(300)

profile_cache = TTLCache(3600)
metrics_cache = TTLCache(3600)


async def _get_json(url: str, params: Dict[str, Any], method: str = "GET", json_body: Any = None, headers: Dict[str, str] | None = None) -> Any:
    h = {"User-Agent": USER_AGENT}
    if headers:
        h.update(headers)
    async with httpx.AsyncClient(timeout=30.0, headers=h) as client:
        if method == "GET":
            r = await client.get(url, params=params)
        else:
            r = await client.post(url, params=params, json=json_body)
        r.raise_for_status()
        return r.json()


async def finnhub_quote(symbol: str) -> Dict[str, Any]:
    data = await _get_json("https://finnhub.io/api/v1/quote", {"symbol": symbol, "token": FINNHUB_API_KEY})
    return {"symbol": symbol, "price": data.get("c"), "open": data.get("o"), "high": data.get("h"), "low": data.get("l"),
            "prevClose": data.get("pc"), "timestamp": data.get("t"), "provider": "finnhub"}


async def twelvedata_quote(symbol: str) -> Dict[str, Any]:
    data = await _get_json("https://api.twelvedata.com/quote", {"symbol": symbol, "apikey": TWELVEDATA_API_KEY})
    def f(x):
        try: return float(x)
        except: return None
    return {"symbol": symbol, "price": f(data.get("close")), "open": f(data.get("open")), "high": f(data.get("high")), "low": f(data.get("low")),
            "prevClose": None, "timestamp": None, "provider": "twelvedata", "raw": {"datetime": data.get("datetime"), "exchange": data.get("exchange")}}


async def get_quote(symbol: str) -> Dict[str, Any]:
    key = f"q:{symbol}"
    c = quotes_cache.get(key)
    if c: return c
    last_err = None
    if FINNHUB_API_KEY:
        try:
            q = await finnhub_quote(symbol)
            quotes_cache.set(key, q)
            return q
        except Exception as e:
            last_err = str(e)
    if TWELVEDATA_API_KEY:
        try:
            q = await twelvedata_quote(symbol)
            quotes_cache.set(key, q)
            return q
        except Exception as e:
            last_err = str(e)
    return {"symbol": symbol, "price": None, "provider": "none", "error": "No market data API key configured.", "detail": last_err}


async def finnhub_closes(symbol: str, days: int = 220) -> List[float]:
    key = f"fhc:{symbol}:{days}"
    c = candles_cache.get(key)
    if c: return c
    to_ts = int(time.time())
    from_ts = to_ts - days*24*3600
    data = await _get_json("https://finnhub.io/api/v1/stock/candle", {"symbol": symbol, "resolution":"D", "from":from_ts, "to":to_ts, "token":FINNHUB_API_KEY})
    if data.get("s") != "ok": return []
    closes = [float(x) for x in data.get("c", [])]
    candles_cache.set(key, closes)
    return closes


async def twelvedata_closes(symbol: str, days: int = 220) -> List[float]:
    key = f"tdc:{symbol}:{days}"
    c = candles_cache.get(key)
    if c: return c
    data = await _get_json("https://api.twelvedata.com/time_series", {"symbol":symbol, "interval":"1day", "outputsize":days, "apikey":TWELVEDATA_API_KEY})
    vals = data.get("values") or []
    closes = []
    for row in reversed(vals):  # oldest -> newest
        try: closes.append(float(row.get("close")))
        except: pass
    candles_cache.set(key, closes)
    return closes


async def get_closes(symbol: str, days: int = 220) -> List[float]:
    if FINNHUB_API_KEY:
        try: return await finnhub_closes(symbol, days)
        except: return []
    if TWELVEDATA_API_KEY:
        try: return await twelvedata_closes(symbol, days)
        except: return []
    return []


def sma(prices: np.ndarray, n: int) -> float:
    return float(np.mean(prices[-n:])) if prices.size >= n else float("nan")

def rsi(prices: np.ndarray, n: int = 14) -> float:
    if prices.size < n+1: return float("nan")
    d = np.diff(prices)
    gains = np.where(d>0, d, 0.0)
    losses = np.where(d<0, -d, 0.0)
    ag = np.mean(gains[-n:])
    al = np.mean(losses[-n:])
    if al == 0: return 100.0
    rs = ag/al
    return float(100 - (100/(1+rs)))

def ema(series: np.ndarray, span: int) -> np.ndarray:
    if series.size == 0: return series
    a = 2/(span+1)
    out = np.zeros_like(series, dtype=float)
    out[0] = series[0]
    for i in range(1, series.size):
        out[i] = a*series[i] + (1-a)*out[i-1]
    return out

def macd(prices: np.ndarray) -> Dict[str, float]:
    if prices.size < 40:
        return {"macd": float("nan"), "signal": float("nan"), "hist": float("nan")}
    m = ema(prices,12) - ema(prices,26)
    s = ema(m,9)
    h = m-s
    return {"macd": float(m[-1]), "signal": float(s[-1]), "hist": float(h[-1])}

def forecast_linear_log(prices: np.ndarray, horizon: int = 5) -> Dict[str, Any]:
    if prices.size < 30:
        return {"method":"linear_log","horizon":horizon,"forecast":[],"note":"Not enough data"}
    y = np.log(np.maximum(prices, 1e-9))
    x = np.arange(y.size)
    b,a = np.polyfit(x,y,1)
    fx = np.arange(y.size, y.size+horizon)
    yhat = a + b*fx
    return {"method":"linear_log","horizon":horizon,"forecast":np.exp(yhat).tolist(),"slope":float(b)}

def signal(last: float, sma50: float, sma200: float, rsi14: float, macd_hist: float) -> Dict[str, Any]:
    score=0; reasons=[]
    if not math.isnan(sma50) and last > sma50: score+=1; reasons.append("Price > SMA50 (trend +)")
    if not math.isnan(sma200) and not math.isnan(sma50) and sma50 > sma200: score+=1; reasons.append("SMA50 > SMA200 (uptrend)")
    if not math.isnan(rsi14) and rsi14 < 30: score+=1; reasons.append("RSI < 30 (oversold)")
    if not math.isnan(rsi14) and rsi14 > 70: score-=1; reasons.append("RSI > 70 (overbought)")
    if not math.isnan(macd_hist) and macd_hist > 0: score+=1; reasons.append("MACD hist > 0 (momentum +)")
    if not math.isnan(macd_hist) and macd_hist < 0: score-=1; reasons.append("MACD hist < 0 (momentum -)")
    lbl = "HOLD"
    if score >= 2: lbl="BUY"
    if score <= -2: lbl="SELL"
    return {"label":lbl,"score":score,"reasons":reasons}


async def company_news(symbol: str, days: int = 7) -> List[Dict[str, Any]]:
    key=f"n:{symbol}:{days}"
    c=news_cache.get(key)
    if c: return c
    if not FINNHUB_API_KEY: return []
    today = dt.date.today()
    start = today - dt.timedelta(days=days)
    data = await _get_json("https://finnhub.io/api/v1/company-news", {"symbol":symbol,"from":start.isoformat(),"to":today.isoformat(),"token":FINNHUB_API_KEY})
    out=[]
    for a in (data or [])[:20]:
        out.append({"headline":a.get("headline"),"source":a.get("source"),"url":a.get("url"),"datetime":a.get("datetime"),"summary":a.get("summary")})
    news_cache.set(key,out)
    return out

async def profile(symbol: str) -> Dict[str, Any]:
    if not FINNHUB_API_KEY: return {}
    key=f"p:{symbol}"
    c=profile_cache.get(key)
    if c: return c
    data = await _get_json("https://finnhub.io/api/v1/stock/profile2", {"symbol":symbol,"token":FINNHUB_API_KEY})
    profile_cache.set(key,data)
    return data

async def metrics(symbol: str) -> Dict[str, Any]:
    if not FINNHUB_API_KEY: return {}
    key=f"m:{symbol}"
    c=metrics_cache.get(key)
    if c: return c
    data = await _get_json("https://finnhub.io/api/v1/stock/metric", {"symbol":symbol,"metric":"all","token":FINNHUB_API_KEY})
    metrics_cache.set(key,data)
    return data


async def hf_generate(prompt: str) -> str:
    if not HF_TOKEN: return ""
    url = f"https://api-inference.huggingface.co/models/{HF_TEXT_MODEL}"
    headers = {"Authorization": f"Bearer {HF_TOKEN}", "User-Agent": USER_AGENT}
    payload = {"inputs": prompt, "parameters": {"max_new_tokens": 220, "return_full_text": False}}
    data = await _get_json(url, {}, method="POST", json_body=payload, headers=headers)
    if isinstance(data, list) and data and isinstance(data[0], dict) and "generated_text" in data[0]:
        return data[0]["generated_text"].strip()
    if isinstance(data, dict) and "generated_text" in data:
        return str(data["generated_text"]).strip()
    return str(data)[:1500]


@app.get("/", response_class=HTMLResponse)
async def index():
    with open("static/index.html", "r", encoding="utf-8") as f:
        return HTMLResponse(f.read())

@app.get("/api/health")
async def health():
    return {
        "ok": True,
        "time": dt.datetime.utcnow().isoformat()+"Z",
        "market_data_provider": "finnhub" if FINNHUB_API_KEY else ("twelvedata" if TWELVEDATA_API_KEY else "none"),
        "ai_enabled": bool(HF_TOKEN),
        "hf_model": HF_TEXT_MODEL
    }

@app.post("/api/quotes")
async def quotes(payload: Dict[str, Any] = Body(...)):
    symbols = payload.get("symbols") or []
    out=[]
    for s in symbols:
        s=str(s).strip()
        if s: out.append(await get_quote(s))
    return {"quotes": out}

@app.post("/api/analyze")
async def analyze(payload: Dict[str, Any] = Body(...)):
    symbol = str(payload.get("symbol","")).strip()
    include_ai = bool(payload.get("include_ai", False))

    q = await get_quote(symbol)
    closes = await get_closes(symbol, 220)
    arr = np.array(closes, dtype=float) if closes else np.array([], dtype=float)

    last = q.get("price")
    if last is None and arr.size: last = float(arr[-1])
    last = float(last) if last is not None else float("nan")

    s50 = sma(arr, 50) if arr.size else float("nan")
    s200 = sma(arr, 200) if arr.size else float("nan")
    r14 = rsi(arr, 14) if arr.size else float("nan")
    m = macd(arr) if arr.size else {"hist": float("nan"), "macd": float("nan"), "signal": float("nan")}
    fc = forecast_linear_log(arr, 5) if arr.size else {"method":"linear_log","horizon":5,"forecast":[]}
    sig = signal(last, s50, s200, r14, m.get("hist", float("nan")))

    exp_return = None
    if fc.get("forecast"):
        try: exp_return = (fc["forecast"][-1]/last - 1)*100
        except: exp_return = None

    ai = ""
    if include_ai and HF_TOKEN:
        headlines = []
        if FINNHUB_API_KEY:
            try:
                n = await company_news(symbol, 7)
                headlines = [x.get("headline") for x in n[:6] if x.get("headline")]
            except: headlines = []
        prompt = (
            "You are an investment research assistant. Provide informational output (not advice).\n"
            f"Ticker: {symbol}\nLast price: {last}\nSMA50: {s50}\nSMA200: {s200}\nRSI14: {r14}\nMACD hist: {m.get('hist')}\n"
            f"Signal: {sig.get('label')} (score {sig.get('score')})\n"
            f"5-day expected return (%): {exp_return}\n"
            f"Recent headlines: {headlines}\n\n"
            "Return JSON with keys: thesis, risk, catalyst, action, timeframe, confidence."
        )
        try: ai = await hf_generate(prompt)
        except Exception as e: ai = f"AI error: {e}"

    return {
        "symbol": symbol,
        "quote": q,
        "closes": closes[-60:],
        "signal": sig,
        "indicators": {"sma50": s50, "sma200": s200, "rsi14": r14, "macd": m},
        "forecast": fc,
        "expected_return_pct_5d": exp_return,
        "ai": ai,
        "notes": [
            "Signals are based on simple technical indicators (trend/momentum) and are NOT financial advice.",
            "Forecast is a lightweight statistical projection (linear regression on log prices)."
        ]
    }

@app.post("/api/news")
async def news(payload: Dict[str, Any] = Body(...)):
    symbol = str(payload.get("symbol","")).strip()
    days = int(payload.get("days", 7))
    return {"symbol": symbol, "articles": await company_news(symbol, days), "provider": "finnhub" if FINNHUB_API_KEY else "none"}

@app.post("/api/fundamentals")
async def fundamentals(payload: Dict[str, Any] = Body(...)):
    symbol = str(payload.get("symbol","")).strip()
    return {"symbol": symbol, "profile": await profile(symbol), "metrics": await metrics(symbol), "provider":"finnhub" if FINNHUB_API_KEY else "none"}

@app.post("/api/portfolio/analyze")
async def portfolio_analyze(payload: Dict[str, Any] = Body(...)):
    holdings = payload.get("holdings") or []
    results=[]
    for h in holdings:
        sym = str(h.get("providerSymbol") or h.get("symbol") or "").strip()
        if not sym: continue
        q = await get_quote(sym)
        price = q.get("price")
        qty = float(h.get("quantity") or 0.0)
        cb = h.get("costBasis")
        if cb is None and isinstance(h.get("lastValue"), (int,float)) and isinstance(h.get("lastReturn"), (int,float)):
            cb = float(h["lastValue"]) - float(h["lastReturn"])
        try: cb = float(cb) if cb is not None else None
        except: cb = None

        mv = float(price)*qty if (price is not None) else None
        pnl = (mv - cb) if (mv is not None and cb is not None) else None
        pnl_pct = (pnl/cb*100) if (pnl is not None and cb not in (None,0.0)) else None

        closes = await get_closes(sym, 220)
        arr = np.array(closes, dtype=float) if closes else np.array([], dtype=float)
        last = float(price) if price is not None else (float(arr[-1]) if arr.size else float("nan"))
        s50 = sma(arr,50) if arr.size else float("nan")
        s200 = sma(arr,200) if arr.size else float("nan")
        r14 = rsi(arr,14) if arr.size else float("nan")
        m = macd(arr) if arr.size else {"hist": float("nan")}
        sig = signal(last, s50, s200, r14, m.get("hist", float("nan")))

        results.append({"symbol": h.get("symbol"), "providerSymbol": sym, "quote": q, "quantity": qty, "costBasis": cb,
                        "marketValue": mv, "pnl": pnl, "pnlPct": pnl_pct, "signal": sig})

    winners = sorted([r for r in results if r.get("pnlPct") is not None], key=lambda x: x["pnlPct"], reverse=True)[:5]
    losers = sorted([r for r in results if r.get("pnlPct") is not None], key=lambda x: x["pnlPct"])[:5]
    counts={"BUY":0,"HOLD":0,"SELL":0}
    for r in results:
        lbl=(r.get("signal") or {}).get("label")
        if lbl in counts: counts[lbl]+=1

    return {"results": results, "topWinners": winners, "topLosers": losers, "signalCounts": counts,
            "provider":"finnhub" if FINNHUB_API_KEY else ("twelvedata" if TWELVEDATA_API_KEY else "none"),
            "aiEnabled": bool(HF_TOKEN)}