#!/usr/bin/env python3
"""
Shadow LLM Improver
Runs after EOD training completes (~6pm ET).
Reviews shadow performance + research findings, proposes specific code changes.
Saves proposals to /opt/services/shadow_llm/proposals/ for human review.
"""

import json, os, glob, requests, subprocess
from datetime import datetime
import difflib

OLLAMA   = "http://localhost:11434/api/chat"
MODEL    = "deepseek-r1:14b"
BASE     = "/opt/services/shadow_llm"
SHADOW   = "/opt/services/bots/shadow_bot/Stock_Bot"
REAL_BOT = "/opt/services/bots/trading_bot/Stock_Bot"

SYSTEM = """You are a careful LightGBM trading bot developer.
You propose minimal, targeted improvements to a shadow (test) bot.
You always output structured JSON proposals, never raw code.
You understand that small targeted changes are safer than large refactors.
You are conservative — you prefer changes with high confidence and low risk.
"""

# Approved modification zones with safe bounds
SAFE_ZONES = {
    "stop_loss_pct": {"min": 0.03, "max": 0.15, "current_line_pattern": "DEFAULT_STOP_LOSS_PCT"},
    "take_profit_pct": {"min": 0.08, "max": 0.30, "current_line_pattern": "DEFAULT_TAKE_PROFIT_PCT"},
    "n_estimators": {"min": 50, "max": 200, "current_line_pattern": "N_ESTIMATORS"},
    "learning_rate": {"min": 0.01, "max": 0.1, "current_line_pattern": "LEARNING_RATE"},
    "num_leaves": {"min": 15, "max": 63, "current_line_pattern": "NUM_LEAVES"},
    "min_confidence": {"min": 0.50, "max": 0.70, "current_line_pattern": "IC_MIN_VOTE_THRESHOLD"},
}

def ask(prompt):
    try:
        r = requests.post(OLLAMA, json={
            "model": MODEL,
            "messages": [
                {"role": "system", "content": SYSTEM},
                {"role": "user", "content": prompt}
            ],
            "stream": False,
            "options": {"temperature": 0.05}  # Very low for code changes
        }, timeout=600)
        return r.json().get("message",{}).get("content","")
    except Exception as e:
        return f"ERROR: {e}"

def load_shadow_performance():
    """Load shadow bot metrics."""
    result = {"status": "no_data"}
    try:
        state = json.load(open(f"{SHADOW}/logs/agent_state.json"))
        agents = [(k,v) for k,v in state.items() if not k.startswith('__')]
        total_pnl = sum(v.get('total_pnl',0) for _,v in agents)
        total_trades = sum(v.get('trades',0) for _,v in agents)
        total_wins = sum(v.get('wins',0) for _,v in agents)
        win_rate = total_wins / max(total_trades, 1) * 100
        result = {
            "total_pnl": round(total_pnl, 2),
            "total_trades": total_trades,
            "win_rate_pct": round(win_rate, 1),
            "n_agents": len(agents),
            "top_agents": sorted(
                [(k, round(v.get('total_pnl',0),2)) for k,v in agents],
                key=lambda x: x[1], reverse=True
            )[:5],
            "bottom_agents": sorted(
                [(k, round(v.get('total_pnl',0),2)) for k,v in agents],
                key=lambda x: x[1]
            )[:3],
        }
    except FileNotFoundError:
        result = {"status": "fresh_start_no_trades_yet"}
    except Exception as e:
        result = {"status": f"error: {e}"}
    return result

def load_real_performance():
    """Load real bot metrics for comparison."""
    result = {"status": "no_data"}
    try:
        state = json.load(open(f"{REAL_BOT}/logs/agent_state.json"))
        agents = [(k,v) for k,v in state.items() if not k.startswith('__')]
        total_pnl = sum(v.get('total_pnl',0) for _,v in agents)
        total_trades = sum(v.get('trades',0) for _,v in agents)
        total_wins = sum(v.get('wins',0) for _,v in agents)
        win_rate = total_wins / max(total_trades,1) * 100
        result = {
            "total_pnl": round(total_pnl, 2),
            "total_trades": total_trades,
            "win_rate_pct": round(win_rate, 1),
        }
    except: pass
    return result

def load_recent_research():
    """Load last 5 research findings."""
    findings = []
    for path in sorted(glob.glob(f"{BASE}/knowledge/finding_*.json"))[-5:]:
        try:
            findings.append(json.load(open(path)))
        except: pass
    return findings

def load_current_constants():
    """Extract current values of safe-zone constants from shadow bot."""
    constants = {}
    try:
        src = open(f"{SHADOW}/StockTrading.py").read()
        lines = src.splitlines()
        for name, info in SAFE_ZONES.items():
            pattern = info["current_line_pattern"]
            for l in lines[:350]:
                if pattern in l and '=' in l:
                    try:
                        val = l.split('=')[1].strip().split()[0].rstrip(',#')
                        constants[name] = {"current": float(val), "bounds": info}
                    except: pass
    except: pass
    return constants

def validate_proposal(proposal):
    """Check proposal is within safe bounds."""
    errors = []
    change = proposal.get("change", {})
    param = change.get("parameter", "")
    value = change.get("new_value")

    if param in SAFE_ZONES:
        bounds = SAFE_ZONES[param]
        try:
            v = float(value)
            if v < bounds["min"] or v > bounds["max"]:
                errors.append(f"{param}={v} outside safe bounds [{bounds['min']}, {bounds['max']}]")
        except:
            errors.append(f"Cannot parse value: {value}")
    elif change.get("type") == "new_feature":
        # Features are allowed if they're derived from OHLCV
        allowed_bases = ["open","high","low","close","volume","vwap","atr"]
        desc = change.get("description","").lower()
        if not any(b in desc for b in allowed_bases):
            errors.append("Feature must be derivable from OHLCV data")
    else:
        errors.append(f"Parameter '{param}' not in approved safe zones")

    return errors

def generate_proposal():
    """Ask LLM to propose one improvement based on performance + research."""
    shadow_perf = load_shadow_performance()
    real_perf   = load_real_performance()
    research    = load_recent_research()
    constants   = load_current_constants()

    prompt = f"""
Shadow bot performance (what we're trying to improve):
{json.dumps(shadow_perf, indent=2)}

Real bot performance (target to beat):
{json.dumps(real_perf, indent=2)}

Current constant values in shadow bot:
{json.dumps(constants, indent=2)}

Recent research findings:
{json.dumps(research, indent=2)}

Based on this data, propose ONE specific improvement. Choose the change most likely to improve
win_rate or total_pnl without increasing risk significantly.

Output ONLY this JSON structure:
{{
    "proposal_id": "PROP_YYYYMMDD_N",
    "title": "short title",
    "rationale": "2-3 sentences explaining why this will help based on the data above",
    "change": {{
        "type": "hyperparameter|new_feature|training_param",
        "parameter": "exact parameter name",
        "current_value": "current value",
        "new_value": "proposed value",
        "description": "what this parameter does and why changing it helps"
    }},
    "expected_metric": "which metric should improve",
    "expected_improvement": "realistic estimate e.g. +2% win rate",
    "confidence": "low|medium|high",
    "risk": "what could go wrong if this change is wrong"
}}

Only output valid JSON. No preamble, no markdown.
"""

    response = ask(prompt)

    # Strip thinking block
    if "</think>" in response:
        response = response.split("</think>")[-1].strip()

    try:
        proposal = json.loads(response)
        # Validate against safe zones
        errors = validate_proposal(proposal)

        ts = datetime.now().strftime("%Y%m%d_%H%M%S")
        proposal["generated_at"]    = ts
        proposal["shadow_perf"]     = shadow_perf
        proposal["real_perf"]       = real_perf
        proposal["validation_errors"] = errors
        proposal["status"]          = "rejected_unsafe" if errors else "pending_review"

        path = f"{BASE}/proposals/proposal_{ts}.json"
        json.dump(proposal, open(path,'w'), indent=2)

        if errors:
            print(f"Proposal generated but REJECTED (unsafe): {errors}")
        else:
            print(f"Proposal saved: {proposal.get('title','?')[:60]}")
            print(f"  Change: {proposal.get('change',{}).get('parameter')} -> {proposal.get('change',{}).get('new_value')}")
            print(f"  Confidence: {proposal.get('confidence')}  Risk: {proposal.get('risk','?')[:60]}")

        return proposal
    except Exception as e:
        print(f"Parse error: {e}\nRaw: {response[:300]}")
        return None

def apply_proposal(proposal_path):
    """
    Apply an approved proposal to the shadow bot.
    Called manually by Jake after reviewing a proposal.
    Creates backup first, applies change, restarts shadow bot.
    """
    proposal = json.load(open(proposal_path))
    if proposal.get("status") != "approved":
        print(f"ERROR: Proposal status is '{proposal.get('status')}', must be 'approved'")
        return False

    # Backup current shadow script
    ts = datetime.now().strftime("%Y%m%d_%H%M%S")
    backup = f"{SHADOW}/StockTrading.py.backup_{ts}"
    import shutil
    shutil.copy(f"{SHADOW}/StockTrading.py", backup)
    print(f"Backup created: {backup}")

    src = open(f"{SHADOW}/StockTrading.py").read()
    change = proposal.get("change", {})
    change_type = change.get("type")
    param = change.get("parameter","")
    new_val = change.get("new_value")

    if change_type == "hyperparameter" and param in SAFE_ZONES:
        pattern = SAFE_ZONES[param]["current_line_pattern"]
        lines = src.splitlines()
        for i,l in enumerate(lines[:350]):
            if pattern in l and '=' in l:
                # Replace the value
                old_line = l
                parts = l.split('=')
                new_line = parts[0] + '= ' + str(new_val) + \
                          ('  # LLM improved' if '#' not in l else '')
                src = src.replace(old_line, new_line, 1)
                print(f"Applied: {old_line.strip()} -> {new_line.strip()}")
                break

    elif change_type == "training_param" and param in SAFE_ZONES:
        pattern = SAFE_ZONES[param]["current_line_pattern"]
        lines = src.splitlines()
        for i,l in enumerate(lines):
            if pattern in l and '=' in l:
                old_line = l
                parts = l.split('=')
                new_line = parts[0] + '= ' + str(new_val) + '  # LLM improved'
                src = src.replace(old_line, new_line, 1)
                print(f"Applied: {old_line.strip()} -> {new_line.strip()}")
                break

    elif change_type == "new_feature":
        print("New feature proposals require manual implementation — see proposal for details")
        return False

    # Save and syntax check
    import py_compile, tempfile
    tmp = tempfile.mktemp(suffix='.py')
    open(tmp,'w').write(src)
    try:
        py_compile.compile(tmp, doraise=True)
        open(f"{SHADOW}/StockTrading.py",'w').write(src)
        print("Syntax check: PASS")
        os.unlink(tmp)
    except py_compile.PyCompileError as e:
        print(f"SYNTAX ERROR: {e} — restoring backup")
        shutil.copy(backup, f"{SHADOW}/StockTrading.py")
        os.unlink(tmp)
        return False

    # Mark proposal applied
    proposal["status"] = "applied"
    proposal["applied_at"] = ts
    json.dump(proposal, open(proposal_path,'w'), indent=2)

    # Restart shadow bot
    subprocess.run(["pkill", "-f", "shadow_bot.*StockTrading"], capture_output=True)
    import time; time.sleep(3)
    log = open(f"{SHADOW}/trading.log","a")
    subprocess.Popen(
        ["/opt/services/bots/trading_bot/venv/bin/python3",
         f"{SHADOW}/StockTrading.py", "--mode", "both", "--max-workers", "8"],
        cwd=SHADOW, stdout=log, stderr=log, start_new_session=True
    )
    print("Shadow bot restarted with applied changes")
    return True

def load_last_metrics():
    """Load yesterday's shadow metrics for regression check."""
    history_path = f"{BASE}/performance_history.json"
    if os.path.exists(history_path):
        try:
            return json.load(open(history_path))
        except: pass
    return {}

def save_metrics_snapshot(metrics):
    """Save today's metrics for tomorrow's regression check."""
    history_path = f"{BASE}/performance_history.json"
    history = load_last_metrics()
    today = datetime.now().strftime("%Y-%m-%d")
    history[today] = metrics
    # Keep last 30 days only
    keys = sorted(history.keys())[-30:]
    history = {k: history[k] for k in keys}
    json.dump(history, open(history_path,'w'), indent=2)

def check_regression(current_metrics, previous_metrics):
    """
    Returns True if current is better than or equal to previous.
    Uses win_rate and total_pnl as primary metrics.
    """
    if not previous_metrics or previous_metrics.get("status") == "fresh_start":
        return True  # No baseline yet — always proceed

    prev_pnl   = previous_metrics.get("total_pnl", 0)
    curr_pnl   = current_metrics.get("total_pnl", 0)
    prev_wr    = previous_metrics.get("win_rate_pct", 50)
    curr_wr    = current_metrics.get("win_rate_pct", 50)

    # Consider improvement if either P&L improved OR win rate improved
    # without the other getting significantly worse
    pnl_ok = curr_pnl >= prev_pnl - 5.0   # allow $5 variance
    wr_ok  = curr_wr  >= prev_wr  - 1.0    # allow 1% variance
    return pnl_ok and wr_ok

def revert_last_change():
    """Revert shadow bot to most recent backup."""
    backups = sorted(glob.glob(f"{SHADOW}/StockTrading.py.backup_*"))
    if not backups:
        print("No backup found to revert to")
        return False
    latest = backups[-1]
    import shutil
    shutil.copy(latest, f"{SHADOW}/StockTrading.py")
    print(f"Reverted to: {latest}")
    # Restart shadow bot with reverted script
    subprocess.run(["pkill", "-f", "shadow_bot.*StockTrading"], capture_output=True)
    import time; time.sleep(3)
    log = open(f"{SHADOW}/trading.log","a")
    subprocess.Popen(
        ["/opt/services/bots/trading_bot/venv/bin/python3",
         f"{SHADOW}/StockTrading.py", "--mode", "both", "--max-workers", "8"],
        cwd=SHADOW, stdout=log, stderr=log, start_new_session=True
    )
    print("Shadow bot restarted with reverted script")
    return True

def autonomous_cycle():
    """
    Full autonomous improvement cycle:
    1. Load current shadow performance
    2. Check if last change caused regression — revert if so
    3. Generate new proposal
    4. If valid, apply automatically and restart shadow bot
    5. Save metrics snapshot for tomorrow's regression check
    6. Push report to GitHub
    """
    ts = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    print(f"\n[{ts}] === AUTONOMOUS IMPROVEMENT CYCLE ===")

    current  = load_shadow_performance()
    previous = load_last_metrics()

    # Step 1: Regression check
    last_date = sorted(previous.keys())[-1] if previous else None
    last_metrics = previous.get(last_date, {}) if last_date else {}

    if last_metrics and not check_regression(current, last_metrics):
        print(f"REGRESSION DETECTED")
        print(f"  Previous P&L: ${last_metrics.get('total_pnl',0):+.2f}  Current: ${current.get('total_pnl',0):+.2f}")
        print(f"  Previous WR:  {last_metrics.get('win_rate_pct',0):.1f}%  Current: {current.get('win_rate_pct',0):.1f}%")
        print("Reverting last change...")
        revert_last_change()
        # Mark last applied proposal as reverted
        applied = sorted(glob.glob(f"{BASE}/proposals/proposal_*.json"))
        for path in reversed(applied):
            try:
                p = json.load(open(path))
                if p.get("status") == "applied":
                    p["status"] = "reverted_regression"
                    p["reverted_at"] = ts
                    json.dump(p, open(path,'w'), indent=2)
                    print(f"Marked as reverted: {p.get('title','?')}")
                    break
            except: pass
    else:
        if last_metrics:
            print(f"Performance OK — P&L: ${current.get('total_pnl',0):+.2f}  WR: {current.get('win_rate_pct',0):.1f}%")
        else:
            print("No baseline yet — first cycle")

    # Step 2: Save today's snapshot
    save_metrics_snapshot(current)

    # Step 3: Generate proposal
    print("Generating improvement proposal...")
    proposal = generate_proposal()

    if proposal is None:
        print("Proposal generation failed — skipping this cycle")
        return

    if proposal.get("status") == "rejected_unsafe":
        print(f"Proposal rejected (unsafe): {proposal.get('validation_errors')}")
        print("Trying again next cycle")
        return

    # Step 4: Auto-apply valid proposals
    print(f"Auto-applying: {proposal.get('title','?')}")
    proposal["status"] = "approved"  # Auto-approve for shadow bot
    proposal_path = f"{BASE}/proposals/proposal_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json"
    json.dump(proposal, open(proposal_path,'w'), indent=2)

    success = apply_proposal(proposal_path)
    if success:
        print(f"Applied successfully. Shadow bot restarting with new parameters.")
        print(f"Change: {proposal.get('change',{}).get('parameter')} -> {proposal.get('change',{}).get('new_value')}")
    else:
        print("Apply failed — shadow bot running on previous parameters")

    print(f"[{ts}] === CYCLE COMPLETE ===\n")

if __name__ == "__main__":
    import sys
    if len(sys.argv) > 1 and sys.argv[1] == "apply":
        # Manual apply still supported
        if len(sys.argv) < 3:
            print("Usage: python3 llm_improver.py apply <proposal_path>")
        else:
            p = json.load(open(sys.argv[2]))
            p["status"] = "approved"
            json.dump(p, open(sys.argv[2],'w'), indent=2)
            apply_proposal(sys.argv[2])
    else:
        # Autonomous cycle
        print(f"[{datetime.now().strftime('%H:%M:%S')}] Running autonomous improvement cycle...")
        autonomous_cycle()