import argparse
import pandas as pd
import numpy as np
import concurrent.futures
import time
import joblib
import datetime
import os
import warnings

warnings.filterwarnings('ignore')

from feature_pipeline import ScreenerScraper

TIERS = {
    "Large": "https://archives.nseindia.com/content/indices/ind_nifty100list.csv",
    "Mid": "https://archives.nseindia.com/content/indices/ind_niftymidcap150list.csv",
    "Small": "https://archives.nseindia.com/content/indices/ind_niftysmallcap250list.csv",
    "Nifty50": "https://archives.nseindia.com/content/indices/ind_nifty50list.csv"
}

def get_current_historic_val(table_data, row_name):
    if not table_data: return np.nan
    for row in table_data.get('rows', []):
        if row and row[0].lower() == row_name.lower():
            for val_raw in reversed(row[1:]):
                val = str(val_raw).replace(',', '').replace('%', '').strip()
                if val not in ['-', '']:
                    try: return float(val)
                    except: continue
    return np.nan

def fetch_live_features(ticker):
    import requests
    from bs4 import BeautifulSoup
    import re
    import numpy as np
    
    try:
        url = f"https://ticker.finology.in/company/{ticker}"
        html = requests.get(url, headers={'User-Agent': 'Mozilla/5.0'}).text
        soup = BeautifulSoup(html, 'html.parser')
        
        def clean(val):
            if not val: return np.nan
            try: return float(re.sub(r'[^\d.]', '', val))
            except: return np.nan
            
        data = {}
        for div in soup.find_all('div', class_=re.compile(r'compess')):
            text = div.get_text(" ", strip=True)
            if "P/E" in text: data['PE_Ratio'] = clean(text.replace('P/E', ''))
            if "Sales Growth" in text: data['Sales_Growth'] = clean(text.replace('Sales Growth', ''))
            if "ROE" in text: data['ROE'] = clean(text.replace('ROE', ''))
            if "ROCE" in text: data['ROCE'] = clean(text.replace('ROCE', ''))
            if "CASH" in text: data['CASH'] = clean(text.replace('CASH', ''))
            if "DEBT " in text: data['DEBT'] = clean(text.replace('DEBT', ''))
            if "Book Value" in text: data['BV'] = clean(text.replace('Book Value', ''))
            if "No. of Shares" in text: data['Shares'] = clean(text.replace('No. of Shares', ''))

        debt_to_equity = np.nan
        if 'DEBT' in data and 'BV' in data and 'Shares' in data and data['BV'] > 0 and data['Shares'] > 0:
            total_equity = data['BV'] * data['Shares']
            debt_to_equity = data['DEBT'] / total_equity if total_equity > 0 else 0

        return {
            'Ticker': ticker,
            'Sales_Growth': data.get('Sales_Growth', np.nan),
            'OPM': np.nan, # Imputer will naturally handle this gracefully
            'ROCE': data.get('ROCE', np.nan),
            'ROE': data.get('ROE', np.nan),
            'Debt_to_Equity': debt_to_equity,
            'PE_Ratio': data.get('PE_Ratio', np.nan)
        }
    except Exception:
        return None

def get_market_cap_tier(ticker):
    print(f"[{ticker}] Resolving market cap classification via Finology...")
    try:
        import requests
        from bs4 import BeautifulSoup
        import re
        url = f"https://ticker.finology.in/company/{ticker}"
        html = requests.get(url, headers={'User-Agent': 'Mozilla/5.0'}).text
        soup = BeautifulSoup(html, 'html.parser')
        
        cap = 0
        for div in soup.find_all('div', class_=re.compile(r'compess')):
            text = div.get_text(" ", strip=True)
            if "Market Cap" in text:
                try: cap = float(re.sub(r'[^\d.]', '', text.replace('Market Cap', '')))
                except: pass
                
        if cap >= 20000: return "Large"
        if cap >= 5000: return "Mid"
        return "Small"
    except Exception as e:
        print(f"Warning: Could not fetch market cap ({e}). Defaulting to Small.")
        
    return "Small" # Default to small cap model for everything else

def generate_reasoning(features, tier, prob):
    reasons = []
    
    if tier == "Large":
        if pd.notnull(features.get('Sales_Growth')):
            if features['Sales_Growth'] < 5: reasons.append(f"Weak Sales Growth ({features['Sales_Growth']:.1f}%) drags down Large Cap momentum.")
            elif features['Sales_Growth'] > 15: reasons.append(f"Strong Sales Growth ({features['Sales_Growth']:.1f}%) is an excellent indicator for Large Caps.")
        if pd.notnull(features.get('ROE')) and features['ROE'] > 20: 
            reasons.append(f"High ROE ({features['ROE']:.1f}%) shows efficient capital use.")
    elif tier == "Small":
        if pd.notnull(features.get('Debt_to_Equity')):
            if features['Debt_to_Equity'] > 1.5: reasons.append(f"Dangerously high Debt/Equity ({features['Debt_to_Equity']:.2f}) signals severe structural risk.")
            elif features['Debt_to_Equity'] < 0.5: reasons.append(f"Low Debt/Equity ({features['Debt_to_Equity']:.2f}) provides strong survival padding.")
        if pd.notnull(features.get('OPM')) and features['OPM'] < 10: 
            reasons.append(f"Low margins ({features['OPM']:.1f}%) leave little room for error.")
    elif tier == "Mid":
        reasons.append("Mid caps exhibit inverted return logic. Metrics are volatile and evaluated in aggregate.")
        
    if not reasons:
        reasons.append("Fundamentals are mixed or average, showing no extreme strengths or weaknesses.")
        
    return " ".join(reasons)

def get_prediction(ticker):
    tier = get_market_cap_tier(ticker)
    features = fetch_live_features(ticker)
    if not features:
        return {"error": "Could not extract live data."}
        
    df = pd.DataFrame([features])
    model_features = ['Sales_Growth', 'OPM', 'ROCE', 'ROE', 'Debt_to_Equity', 'PE_Ratio']
    X = df[model_features].copy()
    
    try:
        model = joblib.load(f'rf_model_{tier.lower()}.pkl')
        imputer = joblib.load(f'imputer_{tier.lower()}.pkl')
        scaler = joblib.load(f'scaler_{tier.lower()}.pkl')
    except Exception as e:
        return {"error": f"Error loading models: {e}"}
        
    X_imputed = imputer.transform(X)
    X_scaled = scaler.transform(X_imputed)
    
    prob = float(model.predict_proba(X_scaled)[0][1])
    decision = "BUY" if prob > 0.65 else "PASS"
    reasoning = generate_reasoning(features, tier, prob)
    
    import math
    clean_features = {}
    for k, v in features.items():
        if isinstance(v, float) and math.isnan(v):
            clean_features[k] = None
        else:
            clean_features[k] = v

    return {
        "Ticker": ticker,
        "Tier": tier,
        "Decision": decision,
        "Confidence": prob,
        "Reasoning": reasoning,
        "Features": clean_features
    }

def run_inference(ticker):
    res = get_prediction(ticker)
    if "error" in res:
        print(f"[{ticker}] {res['error']}")
        return
        
    print("\n" + "="*50)
    print(f" FORECAST FOR {ticker} ({res['Tier']} Cap Model)")
    print("="*50)
    print(f" Decision:   {res['Decision']} (Confidence: {res['Confidence']*100:.1f}%)")
    print(f" Reasoning:  {res['Reasoning']}")
    print("-" * 50)
    for k, v in res['Features'].items():
        print(f" {k}: {v}")
    print("="*50 + "\n")
    return {"Ticker": ticker, "Tier": tier, "Decision": decision, **features}

def run_daemon():
    print("Starting NIFTY 50 Forecasting Daemon...")
    
    try:
        df_nifty = pd.read_csv(TIERS["Nifty50"])
        tickers = df_nifty['Symbol'].tolist()
    except Exception as e:
        print(f"Could not load NIFTY 50 tickers: {e}")
        return
        
    print(f"[{datetime.datetime.now()}] Waking up to process NIFTY 50...")
    
    # We don't fetch market cap tiers because NIFTY 50 is all Large Cap
    
    # Verify rate limit status first with a test call
    print("Testing connection...")
    test_feat = fetch_live_features("RELIANCE")
    if not test_feat:
        print(f"[{datetime.datetime.now()}] Error: IP Rate Limit Block detected on startup. Will not proceed.")
        return
        
    results = []
    total = len(tickers)
    processed = 0
    with concurrent.futures.ThreadPoolExecutor(max_workers=5) as executor:
        futures = {executor.submit(fetch_live_features, ticker): ticker for ticker in tickers}
        for future in concurrent.futures.as_completed(futures):
            res = future.result()
            processed += 1
            if res:
                results.append(res)
                
            pct = int((processed / total) * 100)
            bar = '#' * (pct // 5) + '-' * (20 - (pct // 5))
            print(f"\r[NIFTY 50] {bar} {pct}% ({processed}/{total})", end='', flush=True)
    print("\n", flush=True)
                
    if results:
        df = pd.DataFrame(results)
        features = ['Sales_Growth', 'OPM', 'ROCE', 'ROE', 'Debt_to_Equity', 'PE_Ratio']
        
        try:
            # NIFTY 50 is strictly Large Cap
            model = joblib.load('rf_model_large.pkl')
            imputer = joblib.load('imputer_large.pkl')
            scaler = joblib.load('scaler_large.pkl')
            
            X = df[features].copy()
            X_imputed = imputer.transform(X)
            X_scaled = scaler.transform(X_imputed)
            
            probs = model.predict_proba(X_scaled)[:, 1]
            df['Confidence'] = probs
            df['Decision'] = df['Confidence'].apply(lambda x: "BUY" if x > 0.65 else "PASS")
            
            # Generate reasoning for each
            reasonings = []
            for _, row in df.iterrows():
                feat_dict = {
                    'Sales_Growth': row['Sales_Growth'],
                    'ROE': row['ROE'],
                    'ROCE': row['ROCE'],
                    'Debt_to_Equity': row['Debt_to_Equity'],
                    'OPM': row['OPM']
                }
                r = generate_reasoning(feat_dict, "Large", row['Confidence'])
                reasonings.append(r)
            df['Reasoning'] = reasonings
            
            # Format report and save to JSON
            report = df[['Ticker', 'Decision', 'Confidence', 'Reasoning', 'Sales_Growth', 'ROE', 'Debt_to_Equity', 'OPM']]
            report_dict = report.to_dict(orient='records')
            
            import math
            for item in report_dict:
                for k, v in item.items():
                    if isinstance(v, float) and math.isnan(v):
                        item[k] = None

            import json
            with open("nifty50_predictions.json", "w") as f:
                json.dump({"last_updated": datetime.datetime.now().isoformat(), "predictions": report_dict}, f, indent=4)
                
            print(f"[{datetime.datetime.now()}] Successfully generated nifty50_predictions.json with {len(df)} stocks.", flush=True)
        except Exception as e:
            print(f"Error during NIFTY 50 prediction: {e}", flush=True)
    else:
        print(f"[{datetime.datetime.now()}] Error: Failed to extract live data for all 50 stocks (Likely IP Rate Limit Block). Skipping report generation.", flush=True)
    print("Daemon run completed.", flush=True)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Multi-Cap Stock Forecasting Engine")
    parser.add_argument("--ticker", type=str, help="Run an on-demand forecast for a specific ticker")
    parser.add_argument("--daemon", action="store_true", help="Run the background 14-day loop for NIFTY 50")
    
    args = parser.parse_args()
    
    if args.daemon:
        run_daemon()
    elif args.ticker:
        run_inference(args.ticker.upper())
    else:
        print("Please provide a --ticker or run with --daemon. Use -h for help.")