main.py

from state import TraderState
from agents import hunter_agent, skeptic_agent, calibrator_agent, fetch_data
from dotenv import load_dotenv
from langchain_openai import ChatOpenAI
from langchain_groq import ChatGroq
import os
import re
import json
import pytz
from datetime import datetime, timedelta
import yfinance as yf

load_dotenv()
# os.environ["OPENAI_API_KEY"] = os.getenv("OPENROUTER_API_KEY", "")
# llm_parser = ChatOpenAI(model="openai/gpt-3.5-turbo", openai_api_base="https://openrouter.ai/api/v1", temperature=0.0)

llm_parser = ChatGroq(model = "moonshotai/kimi-k2-instruct-0905",api_key=os.getenv("GROQ_API_KEY"), temperature=0.0)

def resolve_symbol(user_input: str) -> str:
    """
    Handles Indian stocks (.NS/.BO).
    """
    prompt = f"""
    Analyze this user query: "{user_input}".
    - Extract the company/brand name.
    - Infer the likely stock market.
    - Suggest the stock ticker symbol with the appropriate exchange suffix 
      (e.g., "GOOGL" for US, "TCS.NS" for NSE India).
    - If it is an Indian stock, it MUST end with ".NS" or ".BO".
    - Respond only with the symbol. If unsure, say "UNKNOWN".
    """
    try:
        suggested_symbol = llm_parser.invoke(prompt).content.strip().upper()
        if suggested_symbol == "UNKNOWN":
            return None
        
        if not (suggested_symbol.endswith(".NS") or suggested_symbol.endswith(".BO")):
            is_india_prompt = f"Is the stock/company '{suggested_symbol}' from the query '{user_input}' traded on Indian exchanges? Answer only YES or NO."
            is_india = llm_parser.invoke(is_india_prompt).content.strip().upper()
            
            if "YES" in is_india:
                suggested_symbol += ".NS"

        # Try NSE, fallback to BSE
        ticker = yf.Ticker(suggested_symbol)

        if not ticker.history(period="1d").empty:
            return suggested_symbol
        
        # If NSE (.NS) failed, it might be a BSE-only stock
        if suggested_symbol.endswith(".NS"):
            bse_symbol = suggested_symbol.replace(".NS", ".BO")
            if not yf.Ticker(bse_symbol).history(period="1d").empty:
                return bse_symbol
        # Fallback
        refine_prompt = f"The symbol '{suggested_symbol}' didn't validate. Suggest the EXACT ticker for: '{user_input}'. Respond ONLY with the ticker."
        refined_symbol = llm_parser.invoke(refine_prompt).content.strip().upper()
        
        # The same suffix logic
        if not (refined_symbol.endswith(".NS") or refined_symbol.endswith(".BO")):
            if "YES" in is_india: 
                refined_symbol += ".NS"
                
        if not yf.Ticker(refined_symbol).history(period="1d").empty:
            return refined_symbol
            
    except Exception as e:
        print(f"Symbol resolution error: {e}")
    return None


def parse_query(user_input: str) -> tuple:
    """FIXED: Proper horizon detection based on correct trading logic"""
    prompt = f"""
    Analyze this user query: "{user_input}".
    
    TRADING HORIZONS (CRITICAL):
    - "intraday": same day buy and same day sell (hours)
    - "scalping": ultra-short (minutes)
    - "swing": buy this week, sell next week (1-4 weeks)
    - "momentum": trend-following over weeks (2-4 weeks)
    - "long_term": months to years (position trading, e.g., "bought 1 month ago")
    
    Extract:
    - stock symbol (e.g., "HINDUNILVR.NS")
    - action: "buy" or "sell" (default "buy")
    - horizon: match to above definitions
    - date: "today", "tomorrow", or future (e.g., "after 2 weeks")
    - holding_period: If mentioned (e.g., "1 week ago", "a week ago"), provide as "X unit ago". Else, None.
    
    CRITICAL: If user says "bought X time ago" → horizon = "long_term"
    If user says "tomorrow" without history → horizon = "intraday" (same-day strategy for next day)
    
    Respond ONLY in JSON: {{"symbol": "HINDUNILVR.NS", "action": "sell", "horizon": "long_term", "date": "2025-12-27", "holding_period": "1 week ago"}}. Respond sholud be clean.
    """
    try:
        response = llm_parser.invoke(prompt).content.strip()
        parsed = json.loads(response)
        symbol = parsed.get("symbol")
        action = parsed.get("action", "buy")
        horizon = parsed.get("horizon", "intraday")
        date_str = parsed.get("date")
        holding_period = parsed.get("holding_period")
        
        # Override horizon detection with explicit logic
        if "ago" in user_input.lower() or "bought" in user_input.lower() or "purchased" in user_input.lower():
            horizon = "long_term"
        elif "tomorrow" in user_input.lower() and "ago" not in user_input.lower():
            horizon = "intraday"
        elif "week" in user_input.lower() and "ago" not in user_input.lower():
            horizon = "swing"
        elif "minute" in user_input.lower() or "scalp" in user_input.lower():
            horizon = "scalping"
        elif "month" in user_input.lower() and "ago" in user_input.lower():
            horizon = "long_term"
        
        ist = pytz.timezone('Asia/Kolkata')
        now_ist = datetime.now(ist)
        date = None
        if date_str == "today":
            date = now_ist.strftime("%Y-%m-%d")
        elif date_str == "tomorrow":
            date = (now_ist + timedelta(days=1)).strftime("%Y-%m-%d")
        elif date_str and "after" in user_input.lower():
            match = re.search(r'after\s+(\d+)\s*(week|month)', user_input.lower())
            if match:
                num, unit = int(match.group(1)), match.group(2)
                days = num * 7 if unit == "week" else num * 30
                date = (now_ist + timedelta(days=days)).strftime("%Y-%m-%d")
        
        return symbol, date, action, horizon, holding_period
    except Exception as e:
        print(f"Parsing error: {e}. Fallback.")
        symbol = resolve_symbol(user_input)
        action = "sell" if "sell" in user_input.lower() else "buy"
        
        # CRITICAL FIX: Correct horizon fallback
        if "ago" in user_input.lower() or "month" in user_input.lower() or "bought" in user_input.lower():
            horizon = "long_term"
        elif "week" in user_input.lower() and "ago" not in user_input.lower():
            horizon = "swing"
        elif "minute" in user_input.lower():
            horizon = "scalping"
        elif "tomorrow" in user_input.lower():
            horizon = "intraday"
        else:
            horizon = "intraday"
        
        ist = pytz.timezone('Asia/Kolkata')
        now_ist = datetime.now(ist)
        date = now_ist.strftime("%Y-%m-%d") if "today" in user_input.lower() else (now_ist + timedelta(days=1)).strftime("%Y-%m-%d") if "tomorrow" in user_input.lower() else None
        
        match = re.search(r'(\d+)\s*(week|month|day)\s*ago', user_input.lower())
        holding_period = f"{match.group(1)} {match.group(2)} ago" if match else None
        
        return symbol, date, action, horizon, holding_period

def generate_alert_message(state: TraderState, symbol: str, action: str, horizon: str) -> str:
    confidence = state['confidence']
    raw_data = state.get('raw_data', {}) or {}
    indicators = raw_data.get("indicators", {})
    news = raw_data.get("news", [])
    prediction = raw_data.get("prediction", "No prediction available.")
    holding = raw_data.get("holding", {})
    claim = state.get('claim', 'No claim.')
    skepticism = state.get('skepticism', 'No skepticism.')
    features = raw_data.get("features", "")
    
    rsi = indicators.get('rsi', 50)
    macd = indicators.get('macd', 0)
    current_price = indicators.get('current_price', 1)
    
    volatility = 0
    if current_price > 0:
        bb_upper = indicators.get('bb_upper', current_price * 1.1)
        bb_lower = indicators.get('bb_lower', current_price * 0.9)
        volatility = (bb_upper - bb_lower) / current_price * 100
    
    risk_adjustment = min(20, volatility / 5) if volatility > 0 else 0
    
    if holding:
        pnl_adjustment = 10 if holding.get('holding_return', 0) > 5 else -10 if holding.get('holding_return', 0) < -5 else 0
        allocation = max(0, min(100, confidence - risk_adjustment + pnl_adjustment))
    else:
        allocation = max(0, min(100, confidence - risk_adjustment))
    
    if indicators.get("error") == "No data available":
        return f" No Data Available\nSorry, we couldn't fetch reliable data for {symbol}. Check official sources like NSE or Yahoo Finance. Try searching for '{symbol}.NS' if it's an Indian stock."
    
    # HORIZON-AWARE best date suggestion
    horizon_desc = {
        "intraday": "same-day trading (hours)",
        "scalping": "ultra-short (minutes)",
        "swing": "1-4 weeks",
        "momentum": "2-4 weeks",
        "long_term": "months to years"
    }
    
    best_date_prompt = f"""
    Based on prediction: '{prediction}', indicators: RSI {rsi:.1f}, MACD {macd:.2f}, trend: {features.split(';')[0] if features else 'N/A'}.
    User query involves '{action}' for '{horizon}' ({horizon_desc.get(horizon, horizon)}). 
    Suggest the SPECIFIC best date/time to {action} {symbol} for maximum benefit. 
    - For intraday/scalping: suggest time of day (e.g., "Sell tomorrow at 10:30 AM IST for 1.5% gain")
    - For swing/momentum: suggest specific date within the horizon (e.g., "Buy on Jan 2, 2026")
    - For long_term: suggest holding duration (e.g., "Hold until Feb 2026 for 8% gain")
    Include % potential gain/loss. Be accurate and align with user intent.
    """
    try:
        best_date = llm_parser.invoke(best_date_prompt).content.strip()
    except:
        best_date = f"Based on trends, {action} within {horizon_desc.get(horizon, horizon)} for potential 1-3% {'gain' if action == 'buy' else 'exit'}."
    
    why_prompt = f"""
    Based on claim: '{claim}', skepticism: '{skepticism}', features: '{features}', prediction: '{prediction}', holding: '{holding}'.
    RSI is {rsi:.1f} (low = oversold/buy signal, high = overbought/sell signal). MACD is {macd:.2f} (positive = bullish, negative = bearish).
    Explain why the signal is Green/Yellow/Red for {action} {symbol} ({horizon} = {horizon_desc.get(horizon, horizon)}). 
    Use layman terms, reference indicators accurately, tie to data. Be genuine—align with RSI/MACD signals.
    """
    try:
        why = llm_parser.invoke(why_prompt).content.strip()
    except:
        why = f"Based on RSI ({rsi:.1f}) and MACD ({macd:.2f}), conditions are {'favorable' if confidence > 60 else 'mixed' if confidence > 40 else 'unfavorable'} for {action}."
    
    if confidence >= 70:
        color = " Green: 'Yes, Go Ahead!'"
        should_action = f"Yes, {action} now. Allocate {allocation:.0f}% as signals are strong."
    elif confidence >= 40:
        color = " Yellow: 'Wait and Watch'"
        should_action = f"Maybe, monitor closely. Allocate {allocation:.0f}% cautiously due to mixed signals and potential risk."
    else:
        color = " Red: 'No, Stop!'"
        should_action = f"No, avoid {action}. Signals are weak; wait for better conditions."
    
    holding_summary = ""
    if holding:
        h_return = holding['holding_return']
        holding_summary = f"\n**Holding Summary**: Purchased {holding['days_held']} days ago at ₹{holding['purchase_price']:.2f}, current: ₹{holding.get('current_price', current_price):.2f}, return: {h_return:.1f}% ({holding['pnl']})."
        
        
        if h_return > 0:
            advice = "lock in gains"
        elif h_return < 0:
            advice = "wait for recovery or cut losses"
        else:
            advice = "exit at breakeven" 
            
        should_action += f" Based on {h_return:.1f}% {holding['pnl']}, {advice}."

    report = f"""
--- AI Trader Analyzer Report ---
Symbol: {symbol} 
**Current Price**: {current_price:,.4f}
Action: {action.capitalize()} ({horizon.replace('_', '-').capitalize()} = {horizon_desc.get(horizon, horizon)}) "/n"
Date: {state.get('query_date', 'N/A')} (IST) 
Signal: {color} 
Confidence: {confidence}% 
Key Metrics: RSI {rsi:.1f}, MACD {macd:.2f}, Trend {features.split(';')[0] if features else 'N/A'} 
{holding_summary} 
Should I {action.capitalize()}?** {should_action}  
Best Date/Time to {action.capitalize()}**: {best_date} 
Why? {why} 
Prediction: {prediction} 
Recommendation: Allocate {allocation:.0f}% based on analysis. Not financial advice. 
"""
    return report.strip()

def main():
    symbol, date, action, horizon, holding_period = None, None, None, None, None
    while not symbol:
        user_input = input("Describe the stock (e.g., 'I want to buy Apple stock for today' or 'I purchased TCS.NS stock 1 month ago, now sell today'): ").strip()
        symbol, date, action, horizon, holding_period = parse_query(user_input)
        if not symbol:
            print("No symbol detected. Try again.")
    
    print(f"Analyzing: {symbol} on {date or 'recent'} (IST) for {action} ({horizon})" + (f", holding: {holding_period}" if holding_period else ""))
    
    # Fetch data ONCE before the loop
    hist = fetch_data(symbol, horizon)
    
    state = TraderState(
        input_symbol=symbol,
        query_date=date,
        action=action,
        horizon=horizon,
        holding_period=holding_period,
        # Pre-load data
        raw_data={'hist': hist},  
        claim=None,
        skepticism=None,
        confidence=50,
        iterations=0,
        stop=False,
        alert_message=None
    )
    
    max_iterations = 5
    for _ in range(max_iterations):
        state = hunter_agent(state)
        state = skeptic_agent(state)
        state = calibrator_agent(state)
        if state['stop']:
            break
    
    state['alert_message'] = generate_alert_message(state, symbol, action, horizon)
    
    result = {
        "symbol": symbol,
        "date": date,
        "action": action,
        "horizon": horizon,
        "holding_period": holding_period,
        "alert_message": state['alert_message'],
        "iterations": state['iterations'],
        "confidence": state['confidence']
    }
    
    print("\n--- AI Trader Analyzer Result ---")
    print(f"Symbol: {result['symbol']}")
    print(f"Date: {result['date']} (IST)")
    print(f"Action: {result['action']}")
    print(f"Horizon: {result['horizon']}")
    print(f"Holding: {result['holding_period'] or 'None'}")
    print(f"Message:\n{result['alert_message']}")
    print(f"Iterations: {result['iterations']}")
    print(f"Confidence: {result['confidence']}%")

if __name__ == "__main__":
    main()