from fastapi import FastAPI, HTTPException, Query
from pydantic import BaseModel
import talib
import numpy as np
from typing import List, Optional, Dict, Any
import yfinance as yf
from datetime import datetime, timedelta
from enum import Enum
import pandas as pd

app = FastAPI(
    title="Elite US Stock Tracker API", 
    version="2.0.0",
    description="Advanced technical analysis for top 5 US companies with actionable trading signals"
)

# Top 5 US Companies by Market Cap (2024)
class EliteStock(str, Enum):
    APPLE = "AAPL"
    MICROSOFT = "MSFT" 
    NVIDIA = "NVDA"
    ALPHABET = "GOOGL"
    AMAZON = "AMZN"

class TimeFrame(str, Enum):
    ONE_WEEK = "7d"
    ONE_MONTH = "1mo"
    THREE_MONTHS = "3mo"
    SIX_MONTHS = "6mo"
    ONE_YEAR = "1y"

class TradingSignal(str, Enum):
    STRONG_BUY = "STRONG_BUY"
    BUY = "BUY"
    HOLD = "HOLD"
    SELL = "SELL"
    STRONG_SELL = "STRONG_SELL"

class PositionRecommendation(BaseModel):
    signal: TradingSignal
    confidence: float  # 0-100%
    entry_price: Optional[float] = None
    stop_loss: Optional[float] = None
    take_profit: Optional[float] = None
    position_size: str  # Small, Medium, Large
    reason: str

class TechnicalIndicators(BaseModel):
    sma_20: float
    sma_50: float
    ema_12: float
    ema_26: float
    rsi: float
    macd: float
    macd_signal: float
    macd_histogram: float
    bollinger_upper: float
    bollinger_middle: float
    bollinger_lower: float
    atr: float
    stoch_k: float
    stoch_d: float
    williams_r: float
    adx: float

class MarketMetrics(BaseModel):
    current_price: float
    price_change_24h: float
    price_change_pct_24h: float
    volume: int
    avg_volume_20d: float
    market_cap: Optional[float] = None
    pe_ratio: Optional[float] = None
    support_level: float
    resistance_level: float

class ComprehensiveAnalysis(BaseModel):
    symbol: str
    company_name: str
    last_updated: datetime
    market_metrics: MarketMetrics
    technical_indicators: TechnicalIndicators
    position_recommendation: PositionRecommendation
    key_levels: Dict[str, float]
    trend_analysis: Dict[str, Any]

@app.get("/")
async def root():
    return {
        "message": "Elite US Stock Tracker API - Top 5 Companies",
        "version": "2.0.0",
        "supported_stocks": [stock.value for stock in EliteStock],
        "features": [
            "Advanced technical analysis",
            "Position recommendations", 
            "Support/Resistance levels",
            "Multi-timeframe analysis",
            "Risk management signals"
        ]
    }

@app.get("/stocks")
async def list_elite_stocks():
    """Get information about all supported elite stocks"""
    stock_info = {
        "AAPL": {"name": "Apple Inc.", "sector": "Technology"},
        "MSFT": {"name": "Microsoft Corporation", "sector": "Technology"},
        "NVDA": {"name": "NVIDIA Corporation", "sector": "Technology"},
        "GOOGL": {"name": "Alphabet Inc.", "sector": "Technology"},
        "AMZN": {"name": "Amazon.com Inc.", "sector": "Consumer Discretionary"}
    }
    
    return {
        "elite_stocks": stock_info,
        "total_count": len(stock_info),
        "last_updated": datetime.now()
    }

def calculate_support_resistance(prices: np.ndarray, window: int = 20) -> tuple:
    """Calculate dynamic support and resistance levels"""
    recent_prices = prices[-window:]
    support = np.min(recent_prices)
    resistance = np.max(recent_prices)
    return support, resistance

def generate_trading_signal(indicators: dict, market_data: dict) -> PositionRecommendation:
    """Generate intelligent trading signals based on multiple indicators"""
    
    score = 0
    reasons = []
    
    # RSI Analysis
    rsi = indicators['rsi']
    if rsi < 30:
        score += 2
        reasons.append("RSI oversold (bullish)")
    elif rsi > 70:
        score -= 2
        reasons.append("RSI overbought (bearish)")
    elif 40 <= rsi <= 60:
        score += 1
        reasons.append("RSI neutral zone")
    
    # MACD Analysis
    if indicators['macd'] > indicators['macd_signal']:
        score += 1
        reasons.append("MACD bullish crossover")
    else:
        score -= 1
        reasons.append("MACD bearish signal")
    
    # Moving Average Analysis
    current_price = market_data['current_price']
    if current_price > indicators['sma_50']:
        score += 1
        reasons.append("Price above 50-day SMA")
    else:
        score -= 1
        reasons.append("Price below 50-day SMA")
    
    # Bollinger Bands Analysis
    bb_position = (current_price - indicators['bollinger_lower']) / (indicators['bollinger_upper'] - indicators['bollinger_lower'])
    if bb_position < 0.2:
        score += 1
        reasons.append("Near lower Bollinger Band (potential bounce)")
    elif bb_position > 0.8:
        score -= 1
        reasons.append("Near upper Bollinger Band (potential reversal)")
    
    # ADX Trend Strength
    if indicators['adx'] > 25:
        if score > 0:
            score += 1
            reasons.append("Strong trend confirms bullish bias")
        else:
            score -= 1
            reasons.append("Strong trend confirms bearish bias")
    
    # Determine signal and confidence
    if score >= 4:
        signal = TradingSignal.STRONG_BUY
        confidence = min(95, 70 + score * 5)
        position_size = "Large"
    elif score >= 2:
        signal = TradingSignal.BUY
        confidence = min(85, 60 + score * 5)
        position_size = "Medium"
    elif score <= -4:
        signal = TradingSignal.STRONG_SELL
        confidence = min(95, 70 + abs(score) * 5)
        position_size = "Large"
    elif score <= -2:
        signal = TradingSignal.SELL
        confidence = min(85, 60 + abs(score) * 5)
        position_size = "Medium"
    else:
        signal = TradingSignal.HOLD
        confidence = 50 + abs(score) * 10
        position_size = "Small"
    
    # Calculate risk management levels
    atr = indicators['atr']
    if signal in [TradingSignal.STRONG_BUY, TradingSignal.BUY]:
        entry_price = current_price
        stop_loss = current_price - (2 * atr)
        take_profit = current_price + (3 * atr)
    elif signal in [TradingSignal.STRONG_SELL, TradingSignal.SELL]:
        entry_price = current_price
        stop_loss = current_price + (2 * atr)
        take_profit = current_price - (3 * atr)
    else:
        entry_price = None
        stop_loss = None
        take_profit = None
    
    return PositionRecommendation(
        signal=signal,
        confidence=confidence,
        entry_price=entry_price,
        stop_loss=stop_loss,
        take_profit=take_profit,
        position_size=position_size,
        reason="; ".join(reasons)
    )

@app.get("/analysis/{symbol}", response_model=ComprehensiveAnalysis)
async def get_comprehensive_analysis(
    symbol: EliteStock,
    timeframe: TimeFrame = TimeFrame.THREE_MONTHS
):
    """Get comprehensive technical analysis for an elite stock"""
    try:
        # Fetch stock data
        ticker = yf.Ticker(symbol.value)
        hist = ticker.history(period=timeframe.value)
        info = ticker.info
        
        if hist.empty:
            raise HTTPException(status_code=404, detail=f"No data found for {symbol.value}")
        
        # Extract price data
        closes = hist['Close'].values
        highs = hist['High'].values
        lows = hist['Low'].values
        volumes = hist['Volume'].values
        
        # Calculate technical indicators
        sma_20 = talib.SMA(closes, timeperiod=20)
        sma_50 = talib.SMA(closes, timeperiod=50)
        ema_12 = talib.EMA(closes, timeperiod=12)
        ema_26 = talib.EMA(closes, timeperiod=26)
        rsi = talib.RSI(closes, timeperiod=14)
        macd, macd_signal, macd_hist = talib.MACD(closes)
        bb_upper, bb_middle, bb_lower = talib.BBANDS(closes)
        atr = talib.ATR(highs, lows, closes, timeperiod=14)
        stoch_k, stoch_d = talib.STOCH(highs, lows, closes)
        williams_r = talib.WILLR(highs, lows, closes)
        adx = talib.ADX(highs, lows, closes)
        
        # Get latest values
        current_price = float(closes[-1])
        latest_indicators = {
            'sma_20': float(sma_20[-1]) if not np.isnan(sma_20[-1]) else current_price,
            'sma_50': float(sma_50[-1]) if not np.isnan(sma_50[-1]) else current_price,
            'ema_12': float(ema_12[-1]) if not np.isnan(ema_12[-1]) else current_price,
            'ema_26': float(ema_26[-1]) if not np.isnan(ema_26[-1]) else current_price,
            'rsi': float(rsi[-1]) if not np.isnan(rsi[-1]) else 50.0,
            'macd': float(macd[-1]) if not np.isnan(macd[-1]) else 0.0,
            'macd_signal': float(macd_signal[-1]) if not np.isnan(macd_signal[-1]) else 0.0,
            'macd_histogram': float(macd_hist[-1]) if not np.isnan(macd_hist[-1]) else 0.0,
            'bollinger_upper': float(bb_upper[-1]) if not np.isnan(bb_upper[-1]) else current_price * 1.02,
            'bollinger_middle': float(bb_middle[-1]) if not np.isnan(bb_middle[-1]) else current_price,
            'bollinger_lower': float(bb_lower[-1]) if not np.isnan(bb_lower[-1]) else current_price * 0.98,
            'atr': float(atr[-1]) if not np.isnan(atr[-1]) else current_price * 0.02,
            'stoch_k': float(stoch_k[-1]) if not np.isnan(stoch_k[-1]) else 50.0,
            'stoch_d': float(stoch_d[-1]) if not np.isnan(stoch_d[-1]) else 50.0,
            'williams_r': float(williams_r[-1]) if not np.isnan(williams_r[-1]) else -50.0,
            'adx': float(adx[-1]) if not np.isnan(adx[-1]) else 25.0
        }
        
        # Calculate support and resistance
        support, resistance = calculate_support_resistance(closes)
        
        # Market metrics
        price_change_24h = float(closes[-1] - closes[-2]) if len(closes) > 1 else 0.0
        price_change_pct_24h = (price_change_24h / closes[-2] * 100) if len(closes) > 1 else 0.0
        avg_volume_20d = float(np.mean(volumes[-20:])) if len(volumes) >= 20 else float(volumes[-1])
        
        market_metrics = MarketMetrics(
            current_price=current_price,
            price_change_24h=price_change_24h,
            price_change_pct_24h=price_change_pct_24h,
            volume=int(volumes[-1]),
            avg_volume_20d=avg_volume_20d,
            market_cap=info.get('marketCap'),
            pe_ratio=info.get('trailingPE'),
            support_level=float(support),
            resistance_level=float(resistance)
        )
        
        # Generate trading recommendation
        position_rec = generate_trading_signal(latest_indicators, {'current_price': current_price})
        
        # Key levels analysis
        key_levels = {
            "pivot_point": float((highs[-1] + lows[-1] + closes[-1]) / 3),
            "fibonacci_618": float(support + (resistance - support) * 0.618),
            "fibonacci_382": float(support + (resistance - support) * 0.382),
            "vwap": float(np.average(closes[-20:], weights=volumes[-20:])) if len(closes) >= 20 else current_price
        }
        
        # Trend analysis
        trend_analysis = {
            "short_term_trend": "bullish" if latest_indicators['ema_12'] > latest_indicators['ema_26'] else "bearish",
            "medium_term_trend": "bullish" if current_price > latest_indicators['sma_50'] else "bearish",
            "trend_strength": "strong" if latest_indicators['adx'] > 25 else "weak",
            "volatility": "high" if latest_indicators['atr'] / current_price > 0.03 else "normal"
        }
        
        return ComprehensiveAnalysis(
            symbol=symbol.value,
            company_name=info.get('longName', symbol.value),
            last_updated=datetime.now(),
            market_metrics=market_metrics,
            technical_indicators=TechnicalIndicators(**latest_indicators),
            position_recommendation=position_rec,
            key_levels=key_levels,
            trend_analysis=trend_analysis
        )
        
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Analysis failed: {str(e)}")

@app.get("/portfolio/overview")
async def portfolio_overview():
    """Get overview of all elite stocks with quick signals"""
    try:
        results = {}
        
        for stock in EliteStock:
            ticker = yf.Ticker(stock.value)
            hist = ticker.history(period="1mo")
            
            if not hist.empty:
                closes = hist['Close'].values
                current_price = float(closes[-1])
                rsi = talib.RSI(closes, timeperiod=14)
                macd, macd_signal, _ = talib.MACD(closes)
                
                # Quick signal
                latest_rsi = float(rsi[-1]) if not np.isnan(rsi[-1]) else 50.0
                latest_macd = float(macd[-1]) if not np.isnan(macd[-1]) else 0.0
                latest_macd_signal = float(macd_signal[-1]) if not np.isnan(macd_signal[-1]) else 0.0
                
                if latest_rsi < 30 and latest_macd > latest_macd_signal:
                    quick_signal = "BUY"
                elif latest_rsi > 70 and latest_macd < latest_macd_signal:
                    quick_signal = "SELL"
                else:
                    quick_signal = "HOLD"
                
                results[stock.value] = {
                    "current_price": current_price,
                    "rsi": latest_rsi,
                    "quick_signal": quick_signal,
                    "price_change_24h": float(closes[-1] - closes[-2]) if len(closes) > 1 else 0.0
                }
        
        return {
            "portfolio_overview": results,
            "market_sentiment": "mixed",  # Could be enhanced with market-wide analysis
            "last_updated": datetime.now()
        }
        
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Portfolio overview failed: {str(e)}")

@app.get("/alerts/{symbol}")
async def get_trading_alerts(symbol: EliteStock):
    """Get real-time trading alerts for a specific stock"""
    try:
        ticker = yf.Ticker(symbol.value)
        hist = ticker.history(period="5d")  # Last 5 days for recent alerts
        
        if hist.empty:
            raise HTTPException(status_code=404, detail=f"No data found for {symbol.value}")
        
        closes = hist['Close'].values
        highs = hist['High'].values
        lows = hist['Low'].values
        
        # Calculate indicators for alerts
        rsi = talib.RSI(closes, timeperiod=14)
        bb_upper, bb_middle, bb_lower = talib.BBANDS(closes)
        
        alerts = []
        current_price = float(closes[-1])
        latest_rsi = float(rsi[-1]) if not np.isnan(rsi[-1]) else 50.0
        
        # RSI alerts
        if latest_rsi <= 30:
            alerts.append({
                "type": "RSI_OVERSOLD",
                "message": f"RSI at {latest_rsi:.1f} - Potential buying opportunity",
                "urgency": "HIGH"
            })
        elif latest_rsi >= 70:
            alerts.append({
                "type": "RSI_OVERBOUGHT", 
                "message": f"RSI at {latest_rsi:.1f} - Consider taking profits",
                "urgency": "HIGH"
            })
        
        # Bollinger Band alerts
        if current_price <= bb_lower[-1]:
            alerts.append({
                "type": "BOLLINGER_LOWER",
                "message": f"Price touching lower Bollinger Band - Potential reversal",
                "urgency": "MEDIUM"
            })
        elif current_price >= bb_upper[-1]:
            alerts.append({
                "type": "BOLLINGER_UPPER",
                "message": f"Price touching upper Bollinger Band - Overbought condition",
                "urgency": "MEDIUM"
            })
        
        # Volume alerts
        volumes = hist['Volume'].values
        avg_volume = np.mean(volumes[:-1])  # Average excluding today
        volume_ratio = volumes[-1] / avg_volume
        
        if volume_ratio > 2:
            alerts.append({
                "type": "HIGH_VOLUME",
                "message": f"Volume spike: {volume_ratio:.1f}x average - Significant interest",
                "urgency": "HIGH"
            })
        
        return {
            "symbol": symbol.value,
            "alerts": alerts,
            "alert_count": len(alerts),
            "last_updated": datetime.now()
        }
        
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Alerts failed: {str(e)}")

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)