#!/usr/bin/env python3
"""AI Strategy Code Generator — DeepSeek-powered StrategyPlugin code synthesis.

Takes a natural language strategy description and produces a working
StrategyPlugin implementation.  All generated code goes through three
validation gates before it is persisted: syntax check, interface check,
and security audit.

Usage::

    from strategy_code_generator import StrategyCodeGenerator

    gen = StrategyCodeGenerator()
    result = gen.generate("Trend-following with RSI divergence confirmation")
    if result["valid"]:
        print(result["file_path"])
"""

from __future__ import annotations

import ast
import importlib.util
import json
import os
import re
import sys
import textwrap
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

SRC_DIR = Path(__file__).resolve().parent
# In frozen mode (PyInstaller), _MEIPASS is read-only — use config_manager's writable dir
try:
    import config_manager
    PROJECT = config_manager.get_project_dir()
except Exception:
    PROJECT = SRC_DIR.parent
sys.path.insert(0, str(PROJECT))

# -- security: blocked imports in AI-generated code -----------------------

_BLOCKED_IMPORTS: set[str] = {
    "os", "subprocess", "shutil", "sys", "socket", "ctypes",
    "pickle", "marshal", "code", "codeop", "compileall",
    "importlib", "inspect", "traceback", "gc",
    "signal", "multiprocessing", "threading", "concurrent.futures",
    "pty", "fcntl", "resource", "mmap",
}

_BLOCKED_CALLS: set[str] = {
    "eval", "exec", "compile", "__import__", "open",
    "getattr", "setattr", "delattr", "hasattr",
    "globals", "locals", "vars",
}

_STRATEGY_PLUGIN_TEMPLATE = '''"""
{description}
Auto-generated by StrategyCodeGenerator.
Strategy ID: {strategy_id}
"""
from strategy_sdk.interfaces import StrategyPlugin
from strategy_sdk.types import (
    MarketSnapshot,
    PortfolioStateView,
    StrategyInitContext,
    StrategyIntent,
    StrategyManifest,
    StrategyRiskOpinion,
)


class {class_name}(StrategyPlugin):
    """{description}"""

    def __init__(self, params=None):
        self._params = params or {{}}
        self._ctx = None

    def manifest(self):
        return StrategyManifest(
            strategy_id="{strategy_id}",
            name="{display_name}",
            version="1.0.0",
            kind="signal",
            entrypoint="strategies.{strategy_id}",
            description="{description}",
            author="{author}",
            tags=[{tags}],
            capabilities=[{capabilities}],
            supported_symbols=["XAUUSD"],
            supported_timeframes=["M5", "M15", "H1"],
            params_schema={params_schema},
            default_mode="plugin",
            trust_level="sandbox",
        )

    def initialize(self, ctx: StrategyInitContext) -> None:
        self._ctx = ctx
        self._params.update(ctx.config)

    def on_market_snapshot(self, snapshot: MarketSnapshot) -> StrategyIntent:
{signal_body}

    def on_risk_check(
        self,
        intent: StrategyIntent,
        portfolio: PortfolioStateView,
    ) -> StrategyRiskOpinion:
        # NOTE: allow=True by design. Per-strategy risk opinions are advisory only.
        # Hard blocking (circuit breaker, daily loss, session filter, spread, position
        # limits) is enforced centrally by bridge/risk_manager.py.
        reasons = []
        if intent.action == "HOLD":
            return StrategyRiskOpinion(allow=True, reasons=["no action"])
        if intent.confidence < 0.5:
            reasons.append("low confidence")
        if portfolio.position_side != "FLAT" and portfolio.position_side == intent.side:
            reasons.append("already positioned")
        if not reasons:
            reasons.append("signal OK")
        return StrategyRiskOpinion(
            allow=True,
            severity="WARN" if reasons else "INFO",
            score=intent.confidence,
            reasons=reasons,
        )

    def shutdown(self) -> None:
        pass
'''


@dataclass
class GeneratorConfig:
    model: str = "deepseek-chat"
    temperature: float = 0.3
    max_tokens: int = 4096
    author: str = "AI"
    tags: list[str] = field(default_factory=lambda: ["ai-generated"])
    capabilities: list[str] = field(default_factory=lambda: ["intraday"])
    params_schema: dict[str, Any] = field(default_factory=lambda: {
        "ema_short": {"type": "int", "default": 9, "range": [5, 20]},
        "ema_long": {"type": "int", "default": 21, "range": [15, 50]},
        "atr_period": {"type": "int", "default": 14, "range": [7, 21]},
    })


class StrategyCodeGenerator:
    """Generates StrategyPlugin implementations from natural language.

    Pipeline:
      1. Build a prompt with the template and the user's description.
      2. Call DeepSeek API to fill in the signal logic body.
      3. Stamp the body into the template.
      4. Validate: syntax → interface → security.
      5. Write the .py file and .manifest.json to strategies/.
    """

    def __init__(self, config: GeneratorConfig | None = None):
        self._config = config or GeneratorConfig()

    # -- public API -------------------------------------------------------

    def generate(
        self,
        description: str,
        strategy_id: str = "",
        *,
        display_name: str = "",
    ) -> dict[str, Any]:
        """Generate a strategy from a natural language description.

        Returns a dict with ``valid``, ``file_path``, ``manifest_path``,
        ``strategy_id``, and any ``errors``.
        """
        sid = strategy_id or _slugify(description)
        name = display_name or description[:60]

        signal_body = self._call_ai(description)
        code = _STAMP_TEMPLATE.format(
            strategy_id=sid,
            class_name=_to_class_name(sid),
            display_name=name,
            description=description,
            author=self._config.author,
            tags=json.dumps(self._config.tags),
            capabilities=json.dumps(self._config.capabilities),
            params_schema=json.dumps(self._config.params_schema),
            signal_body=textwrap.indent(signal_body, "        "),
        )

        errors = self._validate(code, sid)
        if errors:
            return {
                "valid": False,
                "strategy_id": sid,
                "errors": errors,
                "code": code,
            }

        py_path = self._write_code(sid, code)
        manifest_path = self._write_manifest(sid, name, signal_body)
        return {
            "valid": True,
            "strategy_id": sid,
            "display_name": name,
            "file_path": str(py_path),
            "manifest_path": str(manifest_path),
            "code": code,
            "errors": [],
        }

    # -- AI call ----------------------------------------------------------

    def _call_ai(self, description: str) -> str:
        api_key = _get_api_key()
        if not api_key:
            return self._fallback_body(description)

        prompt = _build_prompt(description)
        try:
            import requests
            resp = requests.post(
                "https://api.deepseek.com/v1/chat/completions",
                headers={
                    "Authorization": f"Bearer {api_key}",
                    "Content-Type": "application/json",
                },
                json={
                    "model": self._config.model,
                    "messages": [
                        {"role": "system", "content": _SYSTEM_PROMPT},
                        {"role": "user", "content": prompt},
                    ],
                    "temperature": self._config.temperature,
                    "max_tokens": self._config.max_tokens,
                },
                timeout=60,
            )
            resp.raise_for_status()
            data = resp.json()
            content = data["choices"][0]["message"]["content"].strip()
            return _extract_code_block(content)
        except Exception as e:
            import logging
            logging.warning(f"[strategy_code_generator] AI call failed, using fallback: {e}")
            return self._fallback_body(description)

    def _fallback_body(self, description: str) -> str:
        """Template-based fallback when no API key or API fails."""
        desc_lower = description.lower()
        if "rsi" in desc_lower and "divergence" in desc_lower:
            return _RSI_DIVERGENCE_BODY
        if "breakout" in desc_lower:
            return _BREAKOUT_BODY
        if "volume" in desc_lower or "wyckoff" in desc_lower:
            return _VOLUME_BODY
        return _DEFAULT_BODY

    # -- validation -------------------------------------------------------

    # Required methods for StrategyPlugin
    _REQUIRED_METHODS = {
        "manifest", "initialize", "on_market_snapshot",
        "on_risk_check", "shutdown",
    }

    def _validate(self, code: str, strategy_id: str) -> list[str]:
        errors: list[str] = []

        # Gate 1: syntax
        try:
            tree = ast.parse(code)
        except SyntaxError as e:
            errors.append(f"syntax error: {e}")
            return errors

        # Gate 2: interface (must have a StrategyPlugin subclass)
        plugin_class = None
        for node in ast.walk(tree):
            if isinstance(node, ast.ClassDef):
                for base in node.bases:
                    if isinstance(base, ast.Name) and base.id == "StrategyPlugin":
                        plugin_class = node
                        break
        if plugin_class is None:
            errors.append("interface: no StrategyPlugin subclass found")
            return errors

        # Gate 2b: Pre-flight — check all 5 required methods exist
        class_methods = {
            n.name for n in ast.walk(plugin_class)
            if isinstance(n, ast.FunctionDef)
        }
        missing_methods = self._REQUIRED_METHODS - class_methods
        if missing_methods:
            errors.append(f"preflight: missing methods: {missing_methods}")

        # Gate 2c: Pre-flight — on_market_snapshot must return StrategyIntent
        for node in ast.walk(plugin_class):
            if isinstance(node, ast.FunctionDef) and node.name == "on_market_snapshot":
                # Check return annotation
                if node.returns and isinstance(node.returns, ast.Name):
                    if node.returns.id != "StrategyIntent":
                        errors.append(f"preflight: on_market_snapshot returns {node.returns.id}, expected StrategyIntent")
                # Check for self-referencing call
                for child in ast.walk(node):
                    if isinstance(child, ast.Call) and isinstance(child.func, ast.Attribute):
                        if (isinstance(child.func.value, ast.Name)
                                and child.func.value.id == "self"
                                and child.func.attr == "on_market_snapshot"):
                            errors.append("preflight: self-referencing call detected (on_market_snapshot calls itself)")

        # Gate 3: security audit
        for node in ast.walk(tree):
            if isinstance(node, ast.Import) or isinstance(node, ast.ImportFrom):
                imported = (
                    node.names[0].name if node.names else ""
                )
                if isinstance(node, ast.ImportFrom):
                    imported = node.module or ""
                base = imported.split(".")[0]
                if base in _BLOCKED_IMPORTS:
                    errors.append(f"security: blocked import '{base}'")
            if isinstance(node, ast.Call):
                if isinstance(node.func, ast.Name):
                    if node.func.id in _BLOCKED_CALLS:
                        errors.append(f"security: blocked call '{node.func.id}'")

        return errors

    # -- persistence ------------------------------------------------------

    def _write_code(self, strategy_id: str, code: str) -> Path:
        strategies_dir = PROJECT / "strategies"
        strategies_dir.mkdir(parents=True, exist_ok=True)
        file_path = strategies_dir / f"{strategy_id}.py"
        file_path.write_text(code, encoding="utf-8")
        return file_path

    def _write_manifest(
        self, strategy_id: str, display_name: str, signal_body: str
    ) -> Path:
        strategies_dir = PROJECT / "strategies"
        manifest_path = strategies_dir / f"{strategy_id}.manifest.json"
        manifest = {
            "strategy_id": strategy_id,
            "name": display_name,
            "version": "1.0.0",
            "kind": "signal",
            "entrypoint": f"strategies.{strategy_id}",
            "description": display_name,
            "author": self._config.author,
            "tags": self._config.tags,
            "capabilities": self._config.capabilities,
            "supported_symbols": ["XAUUSD"],
            "supported_timeframes": ["M5", "M15", "H1"],
            "params_schema": self._config.params_schema,
            "default_mode": "plugin",
            "trust_level": "sandbox",
        }
        manifest_path.write_text(
            json.dumps(manifest, indent=2, ensure_ascii=False),
            encoding="utf-8",
        )
        return manifest_path


# -- template stamping ----------------------------------------------------

_STAMP_TEMPLATE = _STRATEGY_PLUGIN_TEMPLATE


# -- helpers --------------------------------------------------------------

def _slugify(description: str) -> str:
    slug = re.sub(r"[^a-z0-9]+", "_", description.lower()).strip("_")
    return slug[:40] or "custom_strategy"


def _to_class_name(strategy_id: str) -> str:
    return "".join(p.capitalize() for p in strategy_id.split("_"))


def _get_api_key() -> str:
    key = os.environ.get("DEEPSEEK_API_KEY", "")
    if key:
        return key
    try:
        cfg = config_manager.load_app_config()
        key = cfg.get("DEEPSEEK_API_KEY", "")
    except Exception:
        pass
    return key


def _build_prompt(description: str) -> str:
    # 转义花括号防止 f-string KeyError（用户描述可能含 {breakout} 等）
    safe_desc = str(description or "").replace("{", "{{").replace("}", "}}")
    return textwrap.dedent(f"""\
        Write the body of the `on_market_snapshot` method for a gold (XAUUSD)
        trading strategy with this description:

        {safe_desc}

        Only output the method body (from `bars = snapshot.bars` to the
        `return StrategyIntent(...)` statement).  Wrap your code in a single
        ```python ... ``` block.  No other text.
    """)


_SYSTEM_PROMPT = textwrap.dedent("""\
    You are an expert quantitative trading strategy developer.
    You write clean, safe, correct Python code for gold (XAUUSD) trading
    strategies.  Your code:

    - Extracts bars, indicators, and regime from MarketSnapshot.
    - Computes a trading signal (OPEN_BUY, OPEN_SELL, or HOLD).
    - MUST generate both BUY and SELL signal paths — never BUY-only.
      If your strategy only handles one direction, it is INCOMPLETE.
    - Sets entry price, stop-loss, and take-profit levels using ATR-based
      dynamic values (never hardcoded fixed points like sl=-10, tp=+20).
    - Uses snapshot.bars, snapshot.indicators, snapshot.regime for data.
    - NEVER uses eval, exec, subprocess, os, sys or any os-level operations.
    - Is self-contained: no file I/O, no network calls.
    - Returns StrategyIntent with action, side, confidence, entry_plan, exit_plan.
    - All values are float-safe.
""")

_EXTRACT_CODE_RE = re.compile(r"```(?:python)?\s*\n(.*?)```", re.DOTALL)


def _extract_code_block(content: str) -> str:
    m = _EXTRACT_CODE_RE.search(content)
    if m:
        return m.group(1).strip()
    # Remove markdown fences manually if regex fails
    for fence in ("```python", "```"):
        content = content.replace(fence, "")
    return content.strip()


# -- fallback signal bodies -----------------------------------------------

_DEFAULT_BODY = textwrap.dedent("""\
        bars = snapshot.bars
        if not bars or len(bars) < 20:
            return StrategyIntent(
                strategy_id=self.manifest().strategy_id,
                strategy_version="1.0.0",
                config_revision="sandbox",
                instance_id="default",
                action="HOLD",
                side="BUY",
                confidence=0.0,
                reason="insufficient bars",
            )

        close = [float(b.get("close", 0)) for b in bars[-20:]]
        highs = [float(b.get("high", 0)) for b in bars[-20:]]
        lows = [float(b.get("low", 0)) for b in bars[-20:]]
        ema_short = sum(close[-9:]) / 9
        ema_long = sum(close[-21:]) / min(21, len(close))

        # ATR-based dynamic SL/TP
        trs = [max(highs[i]-lows[i], abs(highs[i]-close[i-1]), abs(lows[i]-close[i-1]))
               for i in range(1, len(close))]
        atr = sum(trs[-14:]) / min(14, len(trs)) if trs else 5.0

        if ema_short > ema_long:
            return StrategyIntent(
                strategy_id=self.manifest().strategy_id,
                strategy_version="1.0.0",
                config_revision="sandbox",
                instance_id="default",
                action="BUY",
                side="BUY",
                confidence=0.65,
                reason="EMA crossover bullish",
                entry_plan={"entry": close[-1]},
                exit_plan={"sl": round(close[-1] - atr * 1.5, 2),
                           "tp": round(close[-1] + atr * 2.5, 2)},
            )
        elif ema_short < ema_long:
            return StrategyIntent(
                strategy_id=self.manifest().strategy_id,
                strategy_version="1.0.0",
                config_revision="sandbox",
                instance_id="default",
                action="SELL",
                side="SELL",
                confidence=0.65,
                reason="EMA crossover bearish",
                entry_plan={"entry": close[-1]},
                exit_plan={"sl": round(close[-1] + atr * 1.5, 2),
                           "tp": round(close[-1] - atr * 2.5, 2)},
            )
        return StrategyIntent(
            strategy_id=self.manifest().strategy_id,
            strategy_version="1.0.0",
            config_revision="sandbox",
            instance_id="default",
            action="HOLD",
            side="BUY",
            confidence=0.0,
            reason="no setup",
        )""")

_RSI_DIVERGENCE_BODY = textwrap.dedent("""\
        bars = snapshot.bars
        if not bars or len(bars) < 30:
            return StrategyIntent(
                strategy_id=self.manifest().strategy_id,
                strategy_version="1.0.0",
                config_revision="sandbox",
                instance_id="rsi",
                action="HOLD",
                side="BUY",
                confidence=0.0,
                reason="insufficient bars",
            )

        closes = [float(b.get("close", 0)) for b in bars]
        highs = [float(b.get("high", 0)) for b in bars]
        lows = [float(b.get("low", 0)) for b in bars]
        n = len(closes)

        # Simple RSI14
        gains = sum(max(closes[i] - closes[i-1], 0) for i in range(max(n-14, 0), n))
        losses = sum(max(closes[i-1] - closes[i], 0) for i in range(max(n-14, 0), n))
        rs = gains / losses if losses > 0 else 100
        rsi = 100 - 100 / (1 + rs) if losses > 0 else 50

        entry = closes[-1]
        atr = sum(h - l for h, l in zip(highs[-14:], lows[-14:])) / 14

        if rsi < 35 and closes[-1] > closes[-2]:
            return StrategyIntent(
                strategy_id=self.manifest().strategy_id,
                strategy_version="1.0.0",
                config_revision="sandbox",
                instance_id="rsi",
                action="BUY",
                side="BUY",
                confidence=0.70,
                reason=f"RSI divergence bullish (RSI={rsi:.0f})",
                entry_plan={"entry": entry},
                exit_plan={"sl": entry - 1.5 * atr, "tp": entry + 3.0 * atr},
            )
        return StrategyIntent(
            strategy_id=self.manifest().strategy_id,
            strategy_version="1.0.0",
            config_revision="sandbox",
            instance_id="rsi",
            action="HOLD",
            side="BUY",
            confidence=0.0,
            reason=f"no divergence (RSI={rsi:.0f})",
        )""")

_VOLUME_BODY = textwrap.dedent("""\
        bars = snapshot.bars
        if not bars or len(bars) < 20:
            return StrategyIntent(
                strategy_id=self.manifest().strategy_id,
                strategy_version="1.0.0",
                config_revision="sandbox",
                instance_id="vol",
                action="HOLD",
                side="BUY",
                confidence=0.0,
                reason="insufficient bars",
            )

        closes = [float(b.get("close", 0)) for b in bars]
        highs = [float(b.get("high", 0)) for b in bars]
        lows = [float(b.get("low", 0)) for b in bars]
        n = len(closes)

        # Volume spike detection (using range as volume proxy)
        ranges = [highs[i] - lows[i] for i in range(n)]
        avg_range = sum(ranges[-20:]) / 20
        current_range = ranges[-1]
        range_ratio = current_range / avg_range if avg_range > 0 else 1.0

        entry = closes[-1]
        atr = sum(h - l for h, l in zip(highs[-14:], lows[-14:])) / 14

        if range_ratio > 1.5 and closes[-1] > closes[-2]:
            return StrategyIntent(
                strategy_id=self.manifest().strategy_id,
                strategy_version="1.0.0",
                config_revision="sandbox",
                instance_id="vol",
                action="BUY",
                side="BUY",
                confidence=min(0.75, range_ratio * 0.4),
                reason=f"volume spike bullish (ratio={range_ratio:.1f})",
                entry_plan={"entry": entry},
                exit_plan={"sl": entry - 1.5 * atr, "tp": entry + 3.0 * atr},
            )
        return StrategyIntent(
            strategy_id=self.manifest().strategy_id,
            strategy_version="1.0.0",
            config_revision="sandbox",
            instance_id="vol",
            action="HOLD",
            side="BUY",
            confidence=0.0,
            reason="no volume spike",
        )""")

_BREAKOUT_BODY = textwrap.dedent("""\
        bars = snapshot.bars
        if not bars or len(bars) < 20:
            return StrategyIntent(
                strategy_id=self.manifest().strategy_id,
                strategy_version="1.0.0",
                config_revision="sandbox",
                instance_id="brk",
                action="HOLD",
                side="BUY",
                confidence=0.0,
                reason="insufficient bars",
            )

        highs = [float(b.get("high", 0)) for b in bars]
        lows = [float(b.get("low", 0)) for b in bars]
        closes = [float(b.get("close", 0)) for b in bars]
        n = len(closes)

        highest_20 = max(highs[-21:-1]) if n > 21 else max(highs[:-1])
        lowest_20 = min(lows[-21:-1]) if n > 21 else min(lows[:-1])
        entry = closes[-1]
        atr = sum(h - l for h, l in zip(highs[-14:], lows[-14:])) / 14

        if entry > highest_20:
            return StrategyIntent(
                strategy_id=self.manifest().strategy_id,
                strategy_version="1.0.0",
                config_revision="sandbox",
                instance_id="brk",
                action="BUY",
                side="BUY",
                confidence=0.72,
                reason=f"breakout above 20-bar high {highest_20:.1f}",
                entry_plan={"entry": entry},
                exit_plan={"sl": entry - 1.5 * atr, "tp": entry + 3.0 * atr},
            )
        if entry < lowest_20:
            return StrategyIntent(
                strategy_id=self.manifest().strategy_id,
                strategy_version="1.0.0",
                config_revision="sandbox",
                instance_id="brk",
                action="SELL",
                side="SELL",
                confidence=0.72,
                reason=f"breakdown below 20-bar low {lowest_20:.1f}",
                entry_plan={"entry": entry},
                exit_plan={"sl": entry + 1.5 * atr, "tp": entry - 3.0 * atr},
            )
        return StrategyIntent(
            strategy_id=self.manifest().strategy_id,
            strategy_version="1.0.0",
            config_revision="sandbox",
            instance_id="brk",
            action="HOLD",
            side="BUY",
            confidence=0.0,
            reason="no breakout",
        )""")
