"""
=========================================
CV6 AI Trading OS
AI Consensus Engine
Providers: GPT, Claude, DeepSeek, Gemini,
           Llama (Groq), Ollama
=========================================
"""

import asyncio
import json
import os
import pathlib
import time
from typing import List, Dict, Any, Optional, Set
from dataclasses import dataclass, field, asdict
from loguru import logger

from app.core.settings import settings


# ─── Read enabled models from JSON config (honours UI toggles) ────────────────

def _enabled_from_config() -> Set[str]:
    """Return the set of model names enabled in cv6_autonomous_config.json.
    Empty set means the file is missing / has no list → fall back to env vars."""
    try:
        cfg_file = pathlib.Path(__file__).resolve().parents[2] / "cv6_autonomous_config.json"
        if cfg_file.exists():
            data = json.loads(cfg_file.read_text(encoding="utf-8"))
            models = data.get("enabled_ai_models", [])
            if isinstance(models, list) and models:
                return {m.lower() for m in models}
    except Exception:
        pass
    return set()


# ─── Per-model signal ─────────────────────────────────────────────

@dataclass
class ModelSignal:
    model:      str
    signal:     str        # BUY | SELL | HOLD
    confidence: float      # 0.0 – 1.0
    reason:     str
    latency_ms: float = 0.0
    error:      Optional[str] = None


# ─── Consensus result ─────────────────────────────────────────────

@dataclass
class ConsensusResult:
    signal:     str
    confidence: float      # 0–100 %
    reason:     str
    models:     List[dict] = field(default_factory=list)
    symbol:     str = ""
    timestamp:  float = field(default_factory=time.time)

    def to_dict(self) -> dict:
        return asdict(self)


# ─── Shared system prompt ─────────────────────────────────────────

_SYS = (
    "You are an expert algorithmic trading analyst for Indian markets. "
    "All technical indicators are pre-calculated. Do NOT recalculate them. "
    "Only analyse the complete picture and give your final verdict.\n"
    "Respond ONLY in this exact format:\n"
    "SIGNAL: BUY|SELL|WAIT|REJECT\n"
    "CONFIDENCE: 0-100\n"
    "REASON: <one sentence>\n"
    "RISK_LEVEL: LOW|MEDIUM|HIGH\n"
    "No other text."
)


def _parse(model: str, raw: str, elapsed: float) -> ModelSignal:
    signal, confidence, reason = "HOLD", 0.5, "No reason"
    try:
        lines = {}
        for line in raw.splitlines():
            if ":" in line:
                k, _, v = line.partition(":")
                lines[k.strip().upper()] = v.strip()
        sig = lines.get("SIGNAL", "HOLD").upper()
        # Accept WAIT/REJECT as HOLD (backward compatible with existing vote logic)
        if sig in ("BUY", "SELL"):
            signal = sig
        elif sig in ("WAIT", "REJECT"):
            signal = "HOLD"
        conf_str = lines.get("CONFIDENCE", "50").replace("%", "")
        confidence = min(1.0, float(conf_str) / 100)
        reason = lines.get("REASON", reason)
        # Append risk level to reason if present
        risk = lines.get("RISK_LEVEL", "")
        if risk:
            reason = f"[{risk}] {reason}"
    except Exception as e:
        logger.warning(f"[{model}] Parse error: {e} | raw={raw[:80]}")
    return ModelSignal(model=model, signal=signal, confidence=confidence,
                       reason=reason, latency_ms=round(elapsed * 1000, 1))


# ─── Model Adapters ───────────────────────────────────────────────

class _GPT:
    name = "gpt"

    async def get(self, prompt: str) -> ModelSignal:
        t = time.time()
        try:
            import openai
            client = openai.AsyncOpenAI(api_key=settings.OPENAI_API_KEY)
            r = await client.chat.completions.create(
                model=settings.GPT_MODEL,
                messages=[{"role": "system", "content": _SYS},
                          {"role": "user",   "content": prompt}],
                temperature=settings.GPT_TEMPERATURE,
                max_tokens=settings.GPT_MAX_TOKENS,
                timeout=settings.GPT_TIMEOUT,
            )
            return _parse(self.name, r.choices[0].message.content.strip(), time.time() - t)
        except Exception as e:
            logger.error(f"[GPT] {e}")
            return ModelSignal(self.name, "HOLD", 0.0, "GPT unavailable",
                               error=str(e))


class _Claude:
    name = "claude"

    async def get(self, prompt: str) -> ModelSignal:
        t = time.time()
        try:
            import anthropic
            client = anthropic.AsyncAnthropic(api_key=settings.ANTHROPIC_API_KEY)
            r = await client.messages.create(
                model=settings.CLAUDE_MODEL,
                max_tokens=settings.CLAUDE_MAX_TOKENS,
                system=_SYS,
                messages=[{"role": "user", "content": prompt}],
            )
            return _parse(self.name, r.content[0].text.strip(), time.time() - t)
        except Exception as e:
            logger.error(f"[Claude] {e}")
            return ModelSignal(self.name, "HOLD", 0.0, "Claude unavailable",
                               error=str(e))


class _DeepSeek:
    name = "deepseek"

    async def get(self, prompt: str) -> ModelSignal:
        t = time.time()
        try:
            import openai
            client = openai.AsyncOpenAI(
                api_key=os.getenv("DEEPSEEK_API_KEY", ""),
                base_url="https://api.deepseek.com/v1",
            )
            r = await client.chat.completions.create(
                model="deepseek-chat",
                messages=[{"role": "system", "content": _SYS},
                          {"role": "user",   "content": prompt}],
                temperature=0.2, max_tokens=300,
            )
            return _parse(self.name, r.choices[0].message.content.strip(), time.time() - t)
        except Exception as e:
            logger.error(f"[DeepSeek] {e}")
            return ModelSignal(self.name, "HOLD", 0.0, "DeepSeek unavailable",
                               error=str(e))


class _Gemini:
    name = "gemini"

    async def get(self, prompt: str) -> ModelSignal:
        t = time.time()
        try:
            import google.generativeai as genai
            genai.configure(api_key=os.getenv("GEMINI_API_KEY", ""))
            model = genai.GenerativeModel("gemini-1.5-flash")
            r = await asyncio.to_thread(model.generate_content, f"{_SYS}\n\n{prompt}")
            return _parse(self.name, r.text.strip(), time.time() - t)
        except Exception as e:
            logger.error(f"[Gemini] {e}")
            return ModelSignal(self.name, "HOLD", 0.0, "Gemini unavailable",
                               error=str(e))


class _Llama:
    name = "llama"

    async def get(self, prompt: str) -> ModelSignal:
        t = time.time()
        try:
            import openai
            client = openai.AsyncOpenAI(
                api_key=os.getenv("GROQ_API_KEY", ""),
                base_url="https://api.groq.com/openai/v1",
            )
            r = await client.chat.completions.create(
                model="llama-3.3-70b-versatile",
                messages=[{"role": "system", "content": _SYS},
                          {"role": "user",   "content": prompt}],
                temperature=0.2, max_tokens=300,
            )
            return _parse(self.name, r.choices[0].message.content.strip(), time.time() - t)
        except Exception as e:
            logger.error(f"[Llama] {e}")
            return ModelSignal(self.name, "HOLD", 0.0, "Llama unavailable",
                               error=str(e))


class _Ollama:
    name = "ollama"

    async def get(self, prompt: str) -> ModelSignal:
        t = time.time()
        try:
            import httpx
            url = os.getenv("OLLAMA_URL", "http://localhost:11434")
            model = os.getenv("OLLAMA_MODEL", "mistral")
            async with httpx.AsyncClient(timeout=30) as c:
                r = await c.post(f"{url}/api/generate", json={
                    "model": model,
                    "prompt": f"{_SYS}\n\n{prompt}",
                    "stream": False,
                })
                r.raise_for_status()
            return _parse(self.name, r.json().get("response", "").strip(), time.time() - t)
        except Exception as e:
            logger.error(f"[Ollama] {e}")
            return ModelSignal(self.name, "HOLD", 0.0, "Ollama unavailable",
                               error=str(e))


# ─── Consensus Engine ─────────────────────────────────────────────

class AIConsensusEngine:
    """
    Runs all enabled AI models in parallel.
    Computes weighted majority vote → BUY / SELL / HOLD.
    """

    WEIGHTS: Dict[str, float] = {
        "gpt": 1.2, "claude": 1.2, "deepseek": 1.0,
        "gemini": 1.0, "llama": 0.9, "ollama": 0.8,
    }

    def __init__(self):
        self.adapters: List = []
        self._rebuild_adapters()

    def _rebuild_adapters(self) -> None:
        """Build adapter list — honours UI config JSON, falls back to ENABLE_* env vars."""
        self.adapters = []
        env = os.environ
        cfg_enabled = _enabled_from_config()   # set from JSON config; empty = use env vars

        def _on(env_key: str, model_name: str, env_default: str = "1") -> bool:
            if cfg_enabled:                    # UI config takes priority when present
                return model_name in cfg_enabled
            return env.get(env_key, env_default) == "1"

        if _on("ENABLE_GPT",      "gpt"):      self.adapters.append(_GPT())
        if _on("ENABLE_CLAUDE",   "claude"):   self.adapters.append(_Claude())
        if _on("ENABLE_DEEPSEEK", "deepseek"): self.adapters.append(_DeepSeek())
        if _on("ENABLE_GEMINI",   "gemini"):   self.adapters.append(_Gemini())
        if _on("ENABLE_LLAMA",    "llama"):    self.adapters.append(_Llama())
        if _on("ENABLE_OLLAMA",   "ollama", "0"): self.adapters.append(_Ollama())
        logger.info(f"[Consensus] Active adapters: {[a.name for a in self.adapters]}")

    def refresh_adapters(self) -> None:
        """Call this after UI config is saved so toggles take effect immediately."""
        old = [a.name for a in self.adapters]
        self._rebuild_adapters()
        new = [a.name for a in self.adapters]
        logger.info(f"[Consensus] Adapters refreshed: {old} → {new}")

    def _build_prompt(self, symbol: str, data: Dict[str, Any],
                      context_package: Optional[Dict[str, Any]] = None) -> str:
        """
        Build AI prompt.
        If context_package is provided (from AIContextBuilder), use the rich
        9-section format. Otherwise fall back to basic format.
        The context_package should contain a 'prompt_text' key with the
        pre-formatted prompt string from AIContextPackage.to_prompt().
        """
        if context_package and context_package.get("prompt_text"):
            return context_package["prompt_text"]

        # ── Enriched fallback prompt: fetch real FII/DII + news if available ──
        fii_line  = "FII/DII      : Data not available"
        news_line = "News         : Data not available"
        oi_line   = "Nifty OI/PCR : Data not available"

        try:
            from app.market.nse_data import get_fii_dii, get_india_vix, get_option_chain_oi
            fii = get_fii_dii()
            if fii:
                fii_line = (
                    f"FII          : {fii['fii_activity']} ₹{fii['fii_net']:+.0f}Cr  |  "
                    f"DII: {fii['dii_activity']} ₹{fii['dii_net']:+.0f}Cr  ({fii.get('date','')})"
                )
            vix = get_india_vix()
            vix_line = f"India VIX    : {vix:.1f}" if vix else "India VIX    : N/A"
            oi = get_option_chain_oi("NIFTY")
            if oi:
                oi_line = (
                    f"Nifty PCR    : {oi['pcr']:.3f}  OI Signal: {oi['oi_signal']}  "
                    f"({oi['sentiment']})"
                )
        except Exception:
            vix_line = "India VIX    : N/A"

        try:
            from app.market.news_feed import get_news_summary
            ns = get_news_summary(symbol)
            if ns.total_articles > 0:
                news_line = (
                    f"News         : {ns.sentiment} ({ns.total_articles} articles, "
                    f"score {ns.sentiment_score:+.2f}) | "
                    + (ns.headlines[0][:80] if ns.headlines else "")
                )
        except Exception:
            pass

        return (
            f"=== CV6 TRADING ANALYSIS ===\n"
            f"Symbol       : {symbol}\n"
            f"LTP          : ₹{data.get('ltp','N/A')}  Change: {data.get('change_pct','N/A')}%\n"
            f"OHLC         : O={data.get('open')} H={data.get('high')} "
            f"L={data.get('low')} C={data.get('close')}\n"
            f"Volume       : {data.get('volume','N/A')}\n"
            f"RSI          : {data.get('rsi','N/A')}  "
            f"MACD         : {data.get('macd','N/A')}  "
            f"EMA Trend    : {data.get('trend','N/A')}\n"
            f"{vix_line}\n"
            f"{fii_line}\n"
            f"{oi_line}\n"
            f"{news_line}\n"
            f"=== END DATA ===\n"
            f"Analyse the above: technical indicators + FII activity + OI + news.\n"
            f"Give your final trading signal."
        )

    async def consensus(
        self,
        symbol: str,
        market_data: Dict[str, Any],
        timeout: float = 15.0,
        context_package: Optional[Dict[str, Any]] = None,
    ) -> ConsensusResult:
        prompt = self._build_prompt(symbol, market_data, context_package)
        prompt_type = "rich-context" if (context_package and context_package.get("prompt_text")) else "basic"
        logger.info(f"[Consensus] Running {len(self.adapters)} models for {symbol} [{prompt_type}]")

        tasks = [asyncio.create_task(a.get(prompt)) for a in self.adapters]
        done, pending = await asyncio.wait(tasks, timeout=timeout)
        for t in pending:
            t.cancel()

        signals: List[ModelSignal] = []
        for t in done:
            try:
                signals.append(t.result())
            except Exception as e:
                logger.warning("AI model task failed: {}", e)

        if not signals:
            return ConsensusResult(
                signal="HOLD", confidence=0.0,
                reason="All models timed out or failed",
                symbol=symbol,
            )

        return self._vote(symbol, signals)

    def _vote(self, symbol: str, signals: List[ModelSignal]) -> ConsensusResult:
        scores: Dict[str, float] = {"BUY": 0.0, "SELL": 0.0, "HOLD": 0.0}
        total_w = 0.0
        breakdown = []

        for s in signals:
            w = 0.0 if s.error else self.WEIGHTS.get(s.model, 1.0)
            scores[s.signal] += s.confidence * w
            total_w += w
            breakdown.append({
                "model":      s.model,
                "signal":     s.signal,
                "confidence": round(s.confidence * 100, 1),
                "reason":     s.reason,
                "latency_ms": s.latency_ms,
                "error":      s.error,
            })

        # ── All models errored → safe default: HOLD, no trade ────────────────
        if total_w == 0.0:
            logger.warning(f"[Consensus] {symbol} — all models failed, defaulting to HOLD")
            return ConsensusResult(
                signal="HOLD", confidence=0.0,
                reason="All AI models failed — holding (no trade placed)",
                models=breakdown, symbol=symbol,
            )

        winner = max(scores, key=scores.__getitem__)

        # ── Tie-break: if top two signals are equal, prefer HOLD ──────────────
        sorted_scores = sorted(scores.values(), reverse=True)
        if len(sorted_scores) >= 2 and sorted_scores[0] == sorted_scores[1]:
            winner = "HOLD"

        raw_conf = scores[winner] / total_w if total_w else 0.0
        pct = round(raw_conf * 100, 1)
        agreeing = [s for s in signals if s.signal == winner and not s.error]
        reason = agreeing[0].reason if agreeing else f"Consensus: {winner}"

        logger.info(f"[Consensus] {symbol} → {winner} @ {pct}%")
        return ConsensusResult(
            signal=winner, confidence=pct,
            reason=reason, models=breakdown, symbol=symbol,
        )


# Singleton
consensus_engine = AIConsensusEngine()