"""run_loop() — one function that turns any Python script into a reporting agent.

Usage:
    import wzrd
    wzrd.run_loop()  # uses default keypair, runs forever

    # Or with options:
    wzrd.run_loop(
        keypair="path/to/id.json",
        tasks=["code", "chat", "reasoning"],
        cycle_seconds=300,
        stake=True,
    )

The loop:
    1. Authenticate (Ed25519 challenge/verify)
    2. Pick top model per task via the live velocity signal
    3. Run server-witnessed inference (server calls the LLM and grades it)
       → If server 503s, falls back to client-side inference using local API keys
    4. Report picks with execution_id for verified reward eligibility
    5. Check reward status + claim if available (gasless relay)
    6. Sleep + repeat

Client-side inference fallback:
    Set any of these env vars and the client will use them when the server is overloaded:
        GROQ_API_KEY      → Groq (free tier: 30 RPM)
        OPENAI_API_KEY    → OpenAI
        ANTHROPIC_API_KEY → Anthropic

No SOL required for claims — the relay pays gas.
CCM rewards accrue when the backend marks them available.
"""

from __future__ import annotations

import logging
import signal
import time
import uuid
from typing import Optional, Sequence

import os
from dataclasses import dataclass

from .agent import SLOTS_PER_DAY, WZRDAgent, load_keypair
from .client import PickResult, WZRDError, pick_details, shortlist

logger = logging.getLogger("wzrd.loop")

# Defaults
DEFAULT_TASKS = ("code", "chat", "reasoning")
DEFAULT_CYCLE_SECONDS = 300  # 5 minutes
MIN_CYCLE_SECONDS = 30
SESSION_REFRESH_CYCLES = 100  # re-auth every ~8 hours at 5min cycles
_TREND_BOOST = {
    "surging": 3.0, "accelerating": 2.0, "stable": 1.0,
    "decelerating": 0.5, "cooling": 0.25,
}
# Backoff config (matches swarm-runner production patterns)
_BACKOFF_BASE_S = 30       # initial backoff
_BACKOFF_MAX_S = 600       # 10 min cap
_CLAIM_COOLDOWN_S = 120    # min seconds between claim attempts after failure


def _backoff(failures: int) -> float:
    """Exponential backoff: 30s, 60s, 120s, 240s, 480s, 600s (cap)."""
    return min(_BACKOFF_BASE_S * (2 ** failures), _BACKOFF_MAX_S)


def _is_auth_error(exc: WZRDError) -> bool:
    """Check if the error is a 401 (expired/invalid session)."""
    return getattr(exc, "status_code", None) == 401


def _is_rate_limited(exc: WZRDError) -> bool:
    """Check if the error is a 429 (rate limited)."""
    return getattr(exc, "status_code", None) == 429


def _fmt_rewards(raw: int) -> str:
    """Format raw reward amount (9 decimals) as human-readable."""
    if not raw:
        return "0"
    return f"{raw / 1e9:,.0f}"


def _fetch_ccm_price() -> Optional[float]:
    """Fetch CCM/USD price from Switchboard oracle (3s timeout)."""
    try:
        from .oracle import read_price as _read_price
        return _read_price("ccm", timeout=3.0)
    except Exception:
        return None


def _ccm_balance_from_status(status: dict) -> int:
    """Extract spendable CCM balance from agent.status()."""
    try:
        return int(status.get("solana", {}).get("ccm_balance", 0) or 0)
    except Exception:
        return 0


def _poll_new_ccm_balance(agent: WZRDAgent, previous_balance: int, attempts: int = 3) -> Optional[int]:
    """Poll for a post-claim CCM balance update and return the observed delta."""
    observed_balance = previous_balance
    for _ in range(max(1, attempts)):
        observed_balance = _ccm_balance_from_status(agent.status())
        if observed_balance > previous_balance:
            break
    return max(0, observed_balance - previous_balance)


_EVAL_PROMPTS = {
    "code": "Write a Python function that flattens a nested list. Return only code.",
    "chat": "Explain what a Merkle tree is in two sentences.",
    "reasoning": "If all roses are flowers and some flowers fade quickly, can we conclude some roses fade quickly? Explain.",
    "math": "What is the derivative of x^3 * ln(x)?",
}


@dataclass
class _ClientInferResult:
    content: str
    provider: str
    model: str
    latency_ms: int
    cost_usd: float


def _try_client_inference(task: str) -> Optional[_ClientInferResult]:
    """Try inference using locally-configured API keys. Returns None if no keys available."""
    prompt = _EVAL_PROMPTS.get(task, _EVAL_PROMPTS["chat"])
    import json as _json
    from urllib.request import Request, urlopen
    from urllib.error import URLError

    providers = [
        ("GROQ_API_KEY", "https://api.groq.com/openai/v1/chat/completions", "llama-3.3-70b-versatile", "groq"),
        ("OPENAI_API_KEY", "https://api.openai.com/v1/chat/completions", "gpt-4o-mini", "openai"),
        ("ANTHROPIC_API_KEY", "https://api.anthropic.com/v1/messages", "claude-haiku-4-5-20251001", "anthropic"),
    ]

    for env_var, url, model, provider_name in providers:
        key = os.environ.get(env_var, "")
        if not key:
            continue

        start = time.monotonic()
        try:
            if provider_name == "anthropic":
                body = _json.dumps({
                    "model": model,
                    "max_tokens": 256,
                    "messages": [{"role": "user", "content": prompt}],
                }).encode()
                req = Request(url, data=body, headers={
                    "Content-Type": "application/json",
                    "x-api-key": key,
                    "anthropic-version": "2023-06-01",
                })
                resp = _json.loads(urlopen(req, timeout=30).read())
                content = resp.get("content", [{}])[0].get("text", "")
            else:
                body = _json.dumps({
                    "model": model,
                    "messages": [{"role": "user", "content": prompt}],
                    "max_tokens": 256,
                }).encode()
                req = Request(url, data=body, headers={
                    "Content-Type": "application/json",
                    "Authorization": f"Bearer {key}",
                })
                resp = _json.loads(urlopen(req, timeout=30).read())
                content = resp.get("choices", [{}])[0].get("message", {}).get("content", "")

            latency = int((time.monotonic() - start) * 1000)
            if content:
                return _ClientInferResult(
                    content=content,
                    provider=provider_name,
                    model=f"{provider_name}/{model}",
                    latency_ms=latency,
                    cost_usd=0.0,
                )
        except Exception as exc:
            logger.debug("client-side %s failed: %s", provider_name, exc)
            continue

    return None


def run_loop(
    keypair: "str | bytes | None" = None,
    *,
    tasks: Sequence[str] = DEFAULT_TASKS,
    cycle_seconds: int = DEFAULT_CYCLE_SECONDS,
    max_cycles: Optional[int] = None,
    claim: bool = True,
    stake: bool = False,
    quiet: bool = False,
) -> None:
    """Run the WZRD reward loop. Blocks forever (or for max_cycles).

    Args:
        keypair: Solana keypair — path, base58 string, bytes, or None (auto-detect).
        tasks: Task types to report on. Each task picks and reports one model per cycle.
        cycle_seconds: Seconds between cycles (min 30).
        max_cycles: Stop after N cycles. None = run forever.
        claim: If True, attempt gasless relay claim when CCM is claimable.
        stake: If True, auto-stake claimed CCM on-chain (gasless, locks for 7 days, 1.25x boost).
        quiet: If True, suppress info logging (errors still print).
    """
    if not quiet:
        logging.basicConfig(
            level=logging.INFO,
            format="%(asctime)s [wzrd] %(message)s",
            datefmt="%H:%M:%S",
        )

    cycle_seconds = max(cycle_seconds, MIN_CYCLE_SECONDS)

    # Load keypair
    signer = load_keypair(keypair)
    pubkey = str(signer.pubkey())
    logger.info("agent: %s", pubkey)
    logger.info("tasks: %s | cycle: %ds", ", ".join(tasks), cycle_seconds)

    # Set up graceful shutdown
    _shutdown = False

    def _handle_signal(signum, frame):
        nonlocal _shutdown
        _shutdown = True
        logger.info("shutting down...")

    try:
        signal.signal(signal.SIGINT, _handle_signal)
        signal.signal(signal.SIGTERM, _handle_signal)
    except ValueError:
        pass  # not in main thread — signal handlers unavailable

    agent = WZRDAgent()
    cycle = 0
    total_reports = 0
    verified_reports = 0
    total_inferences = 0
    claims_made = 0
    session_cost_usd: float = 0.0
    session_ccm_start: Optional[int] = None
    quality_by_task: dict[str, list[float]] = {}
    shortlist_shown: set[str] = set()
    prev_pipeline_state: Optional[str] = None
    consecutive_claim_failures: int = 0
    last_claim_attempt: float = 0.0       # monotonic timestamp of last claim try
    consecutive_cycle_errors: int = 0     # for backoff on repeated failures
    force_reauth: bool = False            # set True on 401 to force re-auth next cycle

    while not _shutdown:
        if max_cycles is not None and cycle >= max_cycles:
            logger.info("max_cycles reached (%d), stopping", max_cycles)
            break

        cycle += 1
        cycle_start = time.monotonic()
        cycle_picks: list[PickResult] = []  # [Change 6] track picks for same-model detection

        try:
            # Health check before cycle (skip cycle if backend is down)
            if cycle > 1 and not agent.health_check():
                logger.warning("backend unhealthy — skipping cycle %d", cycle)
                consecutive_cycle_errors += 1
                backoff = _backoff(consecutive_cycle_errors)
                logger.info("→ backing off %.0fs", backoff)
                time.sleep(backoff)
                continue

            # Re-auth: on schedule, on first cycle, or forced by 401
            if cycle == 1 or force_reauth or cycle % SESSION_REFRESH_CYCLES == 0:
                session = agent.authenticate(keypair)
                force_reauth = False
                logger.info("authenticated (session %d cycles)", SESSION_REFRESH_CYCLES)
                # Snapshot starting CCM on first auth
                if session_ccm_start is None:
                    try:
                        _e = agent.earned()
                        session_ccm_start = _e.get("total_earned_ccm", 0)
                    except WZRDError:
                        session_ccm_start = 0
                    # Check SOL balance on first auth
                    try:
                        _s = agent.status()
                        sol_lamports = _s.get("solana", {}).get("sol_lamports", 0)
                        sol = sol_lamports / 1e9
                        if sol < 0.01:
                            logger.warning("LOW SOL: %.4f SOL — claims are free but deposits need SOL", sol)
                        else:
                            logger.info("SOL balance: %.4f", sol)
                    except WZRDError:
                        pass

            # Pick + infer + report per task
            for task in tasks:
                try:
                    choice = pick_details(task)
                    if choice is None:
                        continue

                    # [Change 3] Show signal ranking on first cycle per task
                    if task not in shortlist_shown:
                        shortlist_shown.add(task)
                        try:
                            candidates = shortlist(task, limit=3)
                            if candidates:
                                parts = []
                                for i, c in enumerate(candidates[:3]):
                                    name = c.model.split("/")[-1] if "/" in c.model else c.model
                                    parts.append(f"#{i+1} {name} ({c.score:.2f})")
                                logger.info(
                                    "cycle %d | %s signal: %s",
                                    cycle, task, " > ".join(parts),
                                )
                        except WZRDError:
                            pass  # signal display is best-effort

                    model_id = choice.signal_model or choice.model

                    # Server-witnessed inference: server calls the LLM,
                    # grades the response, and returns execution_id.
                    execution_id = None
                    infer_quality = None
                    infer_latency = None
                    infer_cost = None
                    infer_provider = None
                    infer_preview = None
                    try:
                        pick_id = str(uuid.uuid4())
                        infer_resp = agent.infer(
                            model_id,
                            task_type=task,
                            pick_id=pick_id,
                        )
                        execution_id = infer_resp.get("execution_id")
                        infer_quality = infer_resp.get("quality_score")
                        infer_latency = infer_resp.get("latency_ms")
                        infer_cost = infer_resp.get("cost_usd")
                        infer_provider = infer_resp.get("provider")
                        infer_preview = infer_resp.get("response_preview", "")
                        if execution_id:
                            total_inferences += 1
                            # [Change 1] Accumulate cost
                            if infer_cost is not None:
                                session_cost_usd += infer_cost
                            # [Change 3] Track quality per task
                            if infer_quality is not None:
                                quality_by_task.setdefault(task, []).append(infer_quality)
                            # [Change 2] Show provider, model, quality, latency, cost
                            cost_str = f", ${infer_cost:.4f}" if infer_cost else ""
                            logger.info(
                                "cycle %d | %s | infer via %s → %s (q=%.2f, %dms%s)",
                                cycle, task,
                                infer_provider or "?",
                                infer_resp.get("executed_model", "?"),
                                infer_quality if infer_quality is not None else 0,
                                infer_latency if infer_latency is not None else 0,
                                cost_str,
                            )
                            # [Change 2] Show response preview on first cycle only
                            if cycle == 1 and infer_preview:
                                preview = infer_preview[:80]
                                logger.info('  └─ "%s%s"', preview, "..." if len(infer_preview) > 80 else "")
                    except WZRDError as exc:
                        # Server inference failed — try client-side fallback
                        if not execution_id:
                            local = _try_client_inference(task)
                            if local:
                                infer_quality = 0.5  # self-attested floor
                                infer_latency = local.latency_ms
                                infer_cost = local.cost_usd
                                infer_provider = f"{local.provider}(local)"
                                infer_preview = local.content
                                total_inferences += 1
                                logger.info(
                                    "cycle %d | %s | fallback via %s → %s (%dms)",
                                    cycle, task, local.provider, local.model, local.latency_ms,
                                )
                            else:
                                logger.warning("infer %s failed (no local keys): %s", task, exc)

                    _used_fallback = infer_provider is not None and "(local)" in (infer_provider or "")
                    result = agent.report_pick(
                        choice,
                        execution_id=execution_id,
                        quality_score=infer_quality,
                        latency_ms=infer_latency,
                        cost_usd=infer_cost,
                        source="wzrd-client",
                        policy_version="1.0",
                        candidate_rank=0,
                        exploration=False,
                        fallback_used=_used_fallback or None,
                        run_id=pubkey[:12],
                        cycle_id=cycle,
                    )
                    total_reports += 1
                    if execution_id:
                        verified_reports += 1
                    tag = "✓" if execution_id else "○"
                    logger.info(
                        "cycle %d | %s %s → %s",
                        cycle, tag, task, choice.model,
                    )
                    cycle_picks.append(choice)
                except WZRDError as exc:
                    if _is_auth_error(exc):
                        logger.warning("401 — session expired, will re-auth next cycle")
                        agent.invalidate_session()
                        force_reauth = True
                        break  # skip remaining tasks this cycle
                    if _is_rate_limited(exc):
                        backoff = _backoff(consecutive_cycle_errors)
                        logger.warning("429 rate limited — backing off %.0fs", backoff)
                        time.sleep(backoff)
                        consecutive_cycle_errors += 1
                        break
                    logger.warning("report %s failed: %s", task, exc)

            # [Change 6] Explain same-model picks
            if len(cycle_picks) > 1:
                models = {c.model for c in cycle_picks}
                if len(models) == 1:
                    pick = cycle_picks[0]
                    name = pick.model.split("/")[-1] if "/" in pick.model else pick.model
                    boost = _TREND_BOOST.get(pick.trend, 1.0)
                    logger.info(
                        "→ same pick across tasks: %s is %s on %s (%.1fx trend boost)",
                        name, pick.trend, pick.platform, boost,
                    )

            # Claim (with cooldown backoff after failures)
            if claim and not force_reauth:
                since_last = time.monotonic() - last_claim_attempt if last_claim_attempt else float("inf")
                cooldown = _CLAIM_COOLDOWN_S * (2 ** min(consecutive_claim_failures, 3)) if consecutive_claim_failures > 0 else 0
                if cooldown > 0 and since_last < cooldown:
                    pass  # cooldown active, skip claim
                else:
                    try:
                        earned = agent.earned()
                        pending = earned.get("pending_ccm", 0)
                        pipeline = earned.get("pipeline", {})

                        state = pipeline.get("state")
                        if prev_pipeline_state and state and state != prev_pipeline_state:
                            logger.info("→ pipeline: %s → %s", prev_pipeline_state, state)
                        prev_pipeline_state = state

                        eta_secs = pipeline.get("next_root_eta_secs", 0)
                        if eta_secs and eta_secs > 30:
                            logger.info("→ next reward batch in ~%d min", round(eta_secs / 60))

                        claimable = 0
                        relay_ready = bool(pipeline.get("relay_ready"))
                        try:
                            claim_status = agent.claim_status()
                            claimable = int(claim_status.get("claimable_ccm", 0) or 0)
                            relay_ready = bool(claim_status.get("relay_ready", relay_ready))
                        except WZRDError:
                            claim_status = None

                        if pending > 0 and claimable <= 0:
                            logger.info("→ pending rewards recorded, but nothing new is claimable from the latest root")

                        if claimable > 0 and relay_ready:
                            try:
                                pre_claim_ccm = None
                                if stake:
                                    try:
                                        pre_claim_ccm = _ccm_balance_from_status(agent.status())
                                    except WZRDError:
                                        pre_claim_ccm = None

                                last_claim_attempt = time.monotonic()
                                claim_resp = agent.claim()
                                tx = claim_resp.get("tx_sig", "?")
                                claim_state = str(claim_resp.get("status", "") or "").strip()
                                consecutive_claim_failures = 0
                                if tx and tx != "?":
                                    claims_made += 1
                                    logger.info(
                                        "💰 claimed %s | tx: %s",
                                        _fmt_rewards(claimable) + " rewards",
                                        tx[:20] + "..." if len(tx) > 20 else tx,
                                    )
                                    # Auto-stake only the spendable CCM that actually landed.
                                    # Claims use Token-2022 transfer fees, so claimable_ccm is
                                    # gross protocol accounting, not necessarily the wallet delta.
                                    if stake and claimable > 0:
                                        try:
                                            if pre_claim_ccm is None:
                                                logger.debug("auto-stake skipped: unable to read spendable CCM balance")
                                            else:
                                                stake_amount = _poll_new_ccm_balance(agent, pre_claim_ccm)
                                                if stake_amount <= 0:
                                                    logger.debug(
                                                        "auto-stake skipped: no new spendable CCM detected after claim"
                                                    )
                                                else:
                                                    stake_resp = agent.stake(
                                                        amount=stake_amount,
                                                        lock_duration=SLOTS_PER_DAY * 7,
                                                        keypair=signer,
                                                    )
                                                    stake_tx = stake_resp.get("tx_sig", "?")
                                                    logger.info(
                                                        "⛏ staked %s (7d lock, 1.25x) | tx: %s",
                                                        _fmt_rewards(stake_amount) + " rewards",
                                                        stake_tx[:20] + "..." if len(str(stake_tx)) > 20 else stake_tx,
                                                    )
                                        except WZRDError as stake_exc:
                                            logger.debug("auto-stake skipped: %s", stake_exc)
                                elif claim_state:
                                    logger.info("→ claim relay: %s", claim_state)
                                else:
                                    logger.info("→ claim relay accepted without transaction signature")
                            except WZRDError as exc:
                                consecutive_claim_failures += 1
                                if _is_auth_error(exc):
                                    agent.invalidate_session()
                                    force_reauth = True
                                elif consecutive_claim_failures >= 3:
                                    cd = _CLAIM_COOLDOWN_S * (2 ** min(consecutive_claim_failures - 1, 3))
                                    logger.warning("⚠ %d claim failures — cooldown %ds", consecutive_claim_failures, cd)
                    except WZRDError as exc:
                        if _is_auth_error(exc):
                            agent.invalidate_session()
                            force_reauth = True

            # Reset error counter on successful cycle
            consecutive_cycle_errors = 0

        except WZRDError as exc:
            if _is_auth_error(exc):
                logger.warning("401 — session expired, will re-auth next cycle")
                agent.invalidate_session()
                force_reauth = True
            elif _is_rate_limited(exc):
                consecutive_cycle_errors += 1
                backoff = _backoff(consecutive_cycle_errors)
                logger.warning("429 rate limited — backing off %.0fs", backoff)
                time.sleep(backoff)
            else:
                consecutive_cycle_errors += 1
                logger.error("cycle %d failed: %s", cycle, exc)
                if consecutive_cycle_errors >= 5:
                    backoff = _backoff(consecutive_cycle_errors)
                    logger.warning("→ %d consecutive errors, backing off %.0fs", consecutive_cycle_errors, backoff)
                    time.sleep(backoff)
        except Exception as exc:
            consecutive_cycle_errors += 1
            logger.error("cycle %d unexpected error: %s", cycle, exc)

        # Stop immediately after the final requested cycle instead of idling
        # for a full sleep interval before the next loop check.
        if max_cycles is not None and cycle >= max_cycles:
            logger.info("max_cycles reached (%d), stopping", max_cycles)
            break

        # Sleep until next cycle
        elapsed = time.monotonic() - cycle_start
        sleep_time = max(0, cycle_seconds - elapsed)
        if sleep_time > 0 and not _shutdown:
            time.sleep(sleep_time)

    # Session summary
    logger.info("─" * 48)
    logger.info("session complete: %d cycles", cycle)
    logger.info(
        "reports: %d total (%d verified, %d unverified)",
        total_reports, verified_reports, total_reports - verified_reports,
    )

    # [Change 3] Quality scores per task
    if quality_by_task:
        parts = []
        all_scores = []
        for t in tasks:
            scores = quality_by_task.get(t, [])
            if scores:
                avg = sum(scores) / len(scores)
                parts.append(f"{t}={avg:.2f}")
                all_scores.extend(scores)
        if parts and all_scores:
            session_avg = sum(all_scores) / len(all_scores)
            logger.info("quality: %s | session avg: %.2f", " ".join(parts), session_avg)

    # [Change 1] Session cost
    if total_inferences > 0:
        logger.info("session cost: $%.4f | %d inferences", session_cost_usd, total_inferences)

    if claims_made:
        logger.info("claims: %d", claims_made)

    # Show final earned status
    try:
        final = agent.earned()
        total_ccm = final.get("total_earned_ccm", 0)
        pending_ccm = final.get("pending_ccm", 0)
        rank = final.get("rank")
        streak = final.get("contribution_streak_days", 0)
        session_earned = total_ccm - (session_ccm_start or 0)
        parts = []
        if session_earned > 0:
            parts.append(f"+{_fmt_rewards(session_earned)} earned this session")
        parts.append(_fmt_rewards(total_ccm) + " lifetime")
        if pending_ccm > 0:
            parts.append(f"{_fmt_rewards(pending_ccm)} pending")
        if rank:
            parts.append(f"rank #{rank}")
        if streak > 1:
            parts.append(f"{streak}-day streak")
        logger.info(" | ".join(parts))

        try:
            claim_status = agent.claim_status()
            claimable = int(claim_status.get("claimable_ccm", 0) or 0)
            if claimable > 0 and claim_status.get("relay_ready"):
                logger.info("→ rewards are claimable now (gasless relay)")
            elif pending_ccm > 0:
                logger.info("→ rewards pending - claimable after next batch (usually within 10 min)")
        except WZRDError:
            hint = final.get("pipeline", {}).get("hint")
            if hint:
                logger.info("→ %s", hint)
    except WZRDError:
        pass
