import os
import json
import time
import asyncio
import requests
import pandas as pd
import yfinance as yf
from fastapi import FastAPI, HTTPException, Query
from fastapi.middleware.cors import CORSMiddleware
import google.generativeai as genai
from groq import Groq
from datetime import datetime, timedelta
import pytz

# =================================================================
# CONFIGURATION & KEYS (EXACTLY AS REQUESTED)
# =================================================================
GEMINI_KEY = "AIzaSyD3frbI4K0gsNq6azH-jNGqHoqCzmH8jnE"
GROQ_KEY = "gsk_APdDBgydECWVAdbFAOi7WGdyb3FYUZmiSqYkcS8sjiaqOXNDU21tq"

genai.configure(api_key=GEMINI_KEY)
groq_client = Groq(api_key=GROQ_KEY)

app = FastAPI(title="Vantage Ultra-Prism Pro API", version="4.0.0")

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

# =================================================================
# AGGRESSIVE CACHING UTILITY
# =================================================================
cache_store = {}

def get_cache(key: str, ttl: int):
    """
    ttl in seconds: 
    - Option Chain: 60s
    - Technicals: 30s
    - Prices: 5s
    - AI Strategy: 120s
    """
    if key in cache_store:
        data, timestamp = cache_store[key]
        if time.time() - timestamp < ttl:
            return data
    return None

def set_cache(key: str, data: any):
    cache_store[key] = (data, time.time())

# =================================================================
# MARKET STATUS UTILITY
# =================================================================
def get_market_status():
    ist = pytz.timezone('Asia/Kolkata')
    now = datetime.now(ist)
    
    # Check if Weekend
    if now.weekday() >= 5:
        return False, "Market Closed (Weekend)"
    
    # Check Market Hours (9:15 AM to 3:30 PM)
    start_time = now.replace(hour=9, minute=15, second=0, microsecond=0)
    end_time = now.replace(hour=15, minute=30, second=0, microsecond=0)
    
    if start_time <= now <= end_time:
        return True, "Market Open"
    elif now < start_time:
        return False, f"Market Opens at 09:15 AM"
    else:
        return False, "Market Closed"

# =================================================================
# NSE SESSION MANAGEMENT
# =================================================================
class NSESession:
    def __init__(self):
        self.headers = {
            'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/121.0.0.0 Safari/537.36',
            'Accept': '*/*',
            'Accept-Language': 'en-US,en;q=0.9',
            'Referer': 'https://www.nseindia.com/option-chain',
        }
        self.session = requests.Session()
        self.session.headers.update(self.headers)
        self.last_init = 0

    def init_session(self):
        if time.time() - self.last_init > 300: # Refresh every 5 mins
            try:
                self.session.get("https://www.nseindia.com", timeout=10)
                # Also hit option chain page to get cookies
                self.session.get("https://www.nseindia.com/option-chain", timeout=10)
                self.last_init = time.time()
            except Exception as e:
                print(f"NSE Session Error: {e}")

    def get_json(self, url: str):
        self.init_session()
        try:
            response = self.session.get(url, timeout=8)
            if response.status_code == 200:
                return response.json()
            elif response.status_code == 401:
                self.last_init = 0 # Force re-init
                return self.get_json(url)
        except Exception as e:
            print(f"API Fetch Error: {e}")
        return None

nse_client = NSESession()

# =================================================================
# DATA PROCESSING HELPERS
# =================================================================
def get_historical_data(symbol: str, period="5d", interval="5m"):
    cache_key = f"hist_{symbol}_{period}_{interval}"
    cached = get_cache(cache_key, 30) # Internal buffer
    if cached is not None: return cached

    try:
        df = yf.download(symbol, period=period, interval=interval, progress=False)
        if df.empty:
            df = yf.download(symbol, period="1mo", interval="1d", progress=False)
        
        if not df.empty:
            set_cache(cache_key, df)
            return df
    except:
        pass
    return None

def calculate_technicals(df: pd.DataFrame):
    if df is None or df.empty:
        return {}
    
    try:
        # Extract series
        close = df['Close'].iloc[:, 0] if isinstance(df['Close'], pd.DataFrame) else df['Close']
        high = df['High'].iloc[:, 0] if isinstance(df['High'], pd.DataFrame) else df['High']
        low = df['Low'].iloc[:, 0] if isinstance(df['Low'], pd.DataFrame) else df['Low']
        
        current_price = float(close.iloc[-1])
        prev_close = float(close.iloc[-2]) if len(close) > 1 else current_price
        
        # RSI (14)
        delta = close.diff()
        gain = (delta.where(delta > 0, 0)).rolling(window=14).mean()
        loss = (-delta.where(delta < 0, 0)).rolling(window=14).mean()
        rs = gain / (loss + 1e-9)
        rsi = 100 - (100 / (1 + rs))
        current_rsi = round(float(rsi.iloc[-1]), 2) if not rsi.empty else 50
        
        # EMA (20)
        ema20 = close.ewm(span=20, adjust=False).mean().iloc[-1]
        
        # Support/Resistance (Pivot)
        max_h = float(high.max())
        min_l = float(low.min())
        pivot = (max_h + min_l + current_price) / 3
        r1 = (2 * pivot) - min_l
        s1 = (2 * pivot) - max_h
        
        return {
            "price": round(current_price, 2),
            "change": round(current_price - prev_close, 2),
            "changePercent": round(((current_price - prev_close) / prev_close) * 100, 2),
            "rsi": current_rsi,
            "ema20": round(float(ema20), 2),
            "pivot": round(pivot, 2),
            "r1": round(r1, 2),
            "s1": round(s1, 2),
            "trend": "Bullish" if current_rsi > 55 else "Bearish" if current_rsi < 45 else "Neutral",
            "strength": "Strong" if abs(current_rsi - 50) > 15 else "Weak"
        }
    except Exception as e:
        print(f"Technical Error: {e}")
        return {}

# =================================================================
# AI STRATEGY ENGINE
# =================================================================
async def generate_fno_strategy(symbol: str, techs: dict, oc_data: dict = None):
    cache_key = f"ai_{symbol}"
    cached = get_cache(cache_key, 120)
    if cached: return cached

    prompt = f"""
    ROLE: Institutional F&O Quant Researcher
    TASK: Generate a high-conviction intraday/F&O strategy for {symbol}.
    
    TECHNICALS:
    Price: ₹{techs.get('price')}
    Trend: {techs.get('trend')} ({techs.get('strength')})
    RSI: {techs.get('rsi')}
    Pivot: {techs.get('pivot')}
    Levels: R1={techs.get('r1')}, S1={techs.get('s1')}
    
    OPTION DATA (PCR/OI):
    {oc_data if oc_data else 'No data - Use Technicals'}
    
    Output exactly ONE JSON object in a list. DO NOT include markdown formatting or extra text.
    Schema:
    [{{
        "stockName": "{symbol} Pro Analysis",
        "symbol": "{symbol}",
        "tradeDate": "{datetime.now().strftime('%d %b %Y')}",
        "isMarketOpen": true,
        "tradeType": "Intraday F&O",
        "prevClose": {techs.get('price', 0) - techs.get('change', 0)},
        "openPrice": {techs.get('price', 0)},
        "currentMarketPrice": {techs.get('price', 0)},
        "exactEntryPrice": <float>,
        "exactEntryTime": "ASAP / On Breakout",
        "target1": <float>,
        "target2": <float>,
        "stopLoss": <float>,
        "estimatedProfitPercentage": <float>,
        "estimatedRiskPercentage": <float>,
        "riskLevel": "Low/Medium/High",
        "strategyLogic": "Detailed Quant analysis explanation..."
    }}]
    """

    # Try Gemini First
    try:
        model = genai.GenerativeModel('gemini-1.5-flash')
        response = await asyncio.to_thread(model.generate_content, prompt)
        res_text = response.text.strip()
        if "```json" in res_text:
            res_text = res_text.split("```json")[1].split("```")[0].strip()
        elif "```" in res_text:
            res_text = res_text.split("```")[1].strip()
        data = json.loads(res_text)
        set_cache(cache_key, data)
        return data
    except:
        # Fallback to Groq
        try:
            completion = groq_client.chat.completions.create(
                model="llama-3.3-70b-versatile",
                messages=[{"role": "user", "content": prompt}]
            )
            res_text = completion.choices[0].message.content.strip()
            if "```json" in res_text:
                res_text = res_text.split("```json")[1].split("```")[0].strip()
            data = json.loads(res_text)
            set_cache(cache_key, data)
            return data
        except:
            return [{"stockName": symbol, "strategyLogic": "AI Engine busy. Follow Technical Levels."}]

# =================================================================
# API ENDPOINTS
# =================================================================

@app.get("/api/trade")
async def get_trade(symbol: str = "NIFTY"):
    """ High-Speed AI Strategy Endpoint """
    yf_map = {"NIFTY": "^NSEI", "BANKNIFTY": "^NSEBANK", "FINNIFTY": "NIFTY_FIN_SERVICE.NS"}
    yf_sym = yf_map.get(symbol.upper(), f"{symbol.upper()}.NS")
    
    df = get_historical_data(yf_sym)
    techs = calculate_technicals(df)
    
    if not techs:
        raise HTTPException(status_code=404, detail="Symbol Data Error")
    
    # Try to get OC summary for better AI input
    oc_summary = None
    if symbol.upper() in ["NIFTY", "BANKNIFTY", "FINNIFTY"]:
        try:
            oc = await get_option_chain(symbol)
            oc_summary = {"pcr": oc.get("pcr"), "max_pain": "High Volume Zone"}
        except: pass
        
    return await generate_fno_strategy(symbol, techs, oc_summary)

@app.get("/api/option-chain")
async def get_option_chain(symbol: str = "NIFTY"):
    """ Professional F&O Option Chain Endpoint """
    cache_key = f"oc_{symbol.upper()}"
    cached = get_cache(cache_key, 60)
    if cached: return cached

    base_url = "https://www.nseindia.com/api/option-chain-"
    url = f"{base_url}indices?symbol={symbol.upper()}"
    if symbol.upper() not in ["NIFTY", "BANKNIFTY", "FINNIFTY", "MIDCPNIFTY"]:
        url = f"{base_url}equities?symbol={symbol.upper()}"
    
    data = nse_client.get_json(url)
    if not data:
        # Fallback Mock/Simulated for Stability
        return {"symbol": symbol, "error": "NSE Latency. Use Technicals.", "status": "Degraded"}
    
    try:
        records = data.get('records', {})
        filtered = data.get('filtered', {})
        ce_oi = filtered.get('CE', {}).get('totOI', 0)
        pe_oi = filtered.get('PE', {}).get('totOI', 0)
        pcr = round(pe_oi / (ce_oi or 1), 2)
        
        result = {
            "symbol": symbol.upper(),
            "underlying": records.get('underlyingValue'),
            "pcr": pcr,
            "sentiment": "Bullish" if pcr > 1.1 else "Bearish" if pcr < 0.8 else "Neutral",
            "timestamp": records.get('timestamp'),
            "expiries": records.get('expiryDates', [])[:3],
            "data": filtered.get('data', [])[:15] # Truncated for speed
        }
        set_cache(cache_key, result)
        return result
    except:
        return {"error": "Processing Error"}

@app.get("/api/indices")
async def get_indices():
    """ Live Market Sentiment Overview """
    cache_key = "indices_v4"
    cached = get_cache(cache_key, 5)
    if cached: return cached

    idx_map = {"^NSEI": "NIFTY 50", "^NSEBANK": "BANK NIFTY", "^BSESN": "SENSEX", "^CNXIT": "NIFTY IT"}
    results = []
    
    try:
        data = yf.download(list(idx_map.keys()), period="1d", interval="5m", progress=False)
        for sym, name in idx_map.items():
            try:
                curr = float(data['Close'][sym].iloc[-1])
                prev = float(data['Open'][sym].iloc[0])
                change = curr - prev
                results.append({
                    "name": name,
                    "symbol": sym,
                    "price": round(curr, 2),
                    "change": round(change, 2),
                    "percent": f"{'+' if change >= 0 else ''}{((change/prev)*100):.2f}%"
                })
            except: continue
    except: pass
    
    set_cache(cache_key, results)
    return results

@app.get("/api/stocks")
async def get_stocks():
    """ Top Trending Stocks """
    cache_key = "stocks_v4"
    cached = get_cache(cache_key, 10)
    if cached: return cached

    stock_list = ["RELIANCE.NS", "TCS.NS", "HDFCBANK.NS", "INFY.NS", "ICICIBANK.NS", "SBIN.NS", "AXISBANK.NS"]
    results = []
    
    try:
        data = yf.download(stock_list, period="1d", interval="5m", progress=False)
        for sym in stock_list:
            try:
                curr = float(data['Close'][sym].iloc[-1])
                prev = float(data['Open'][sym].iloc[0])
                change = ((curr-prev)/prev)*100
                results.append({
                    "symbol": sym.replace(".NS", ""),
                    "price": round(curr, 2),
                    "changePercent": f"{'+' if change >= 0 else ''}{change:.2f}%",
                    "trend": "Up" if change > 0 else "Down"
                })
            except: continue
    except: pass
    
    set_cache(cache_key, results)
    return results

@app.get("/api/search")
async def search(query: str):
    """ Ultra-Fast Search """
    try:
        s = yf.Search(query, max_results=5)
        return [{"symbol": q['symbol'].replace(".NS", ""), "name": q['shortname']} for q in s.quotes if ".NS" in q['symbol']]
    except: return []

@app.get("/api/technical")
async def technical(symbol: str):
    yf_sym = f"{symbol.upper()}.NS" if "^" not in symbol else symbol
    df = get_historical_data(yf_sym)
    return calculate_technicals(df)

@app.get("/")
def health():
    isOpen, msg = get_market_status()
    return {
        "engine": "Vantage-Ultra-Pro-V4",
        "market_status": msg,
        "is_open": isOpen,
        "cache_objects": len(cache_store),
        "server_time": datetime.now().strftime("%H:%M:%S")
    }

if __name__ == "__main__":
    import uvicorn
    # Use 7860 for Hugging Face standard
    uvicorn.run(app, host="0.0.0.0", port=7860)