"""
OpOp Telemetry — Crowdsourced optimizer intelligence.

Every training run teaches the brain something. Telemetry collects that
experience so the optimizer gets smarter for everyone.

What we collect:
    - Brain observations (loss trends, gradient statistics — just numbers)
    - Brain decisions (scale, clip, dampen per parameter group)
    - Training rewards (did loss improve? by how much?)
    - Trained brain checkpoints (~6KB each — the distilled strategy)
    - Run metadata (n_groups, obs_dim, framework, step count)

What we DO NOT collect:
    - Model weights or architecture details
    - Training data or file paths
    - User identity, IP addresses, or system info
    - Anything that could identify you or your work

Disable: OPOP_TELEMETRY=0 or pass telemetry=False to OptBrain/TwinCrack.

The more people train with OpOp, the smarter it gets for everyone.
"""

import os
import time
import uuid
import json
import numpy as np
from pathlib import Path

# ═════════════════════════════════════════════════════════════════════════════
# CONSTANTS
# ═════════════════════════════════════════════════════════════════════════════

OPOP_DIR = Path.home() / ".opop"
TELEMETRY_DIR = OPOP_DIR / "telemetry"
BRAINS_DIR = TELEMETRY_DIR / "brains"
EXPERIENCE_DIR = TELEMETRY_DIR / "experience"
NOTICE_FLAG = OPOP_DIR / ".telemetry_notice_shown"
ENDPOINT_FILE = OPOP_DIR / "telemetry_endpoint.txt"

# Buffer size before flushing to disk
FLUSH_EVERY = 100

DISCLAIMER = """
╔══════════════════════════════════════════════════════════════════════════╗
║                     OpOp Telemetry — Enabled                           ║
╠══════════════════════════════════════════════════════════════════════════╣
║                                                                        ║
║  OpOp collects anonymized optimizer training data to improve the       ║
║  brain for all users. Every run makes OpOp smarter.                    ║
║                                                                        ║
║  Collected:                                                            ║
║    • Brain observations (loss trends, gradient statistics)             ║
║    • Brain decisions (scale, clip, dampen per group)                   ║
║    • Training rewards (did loss improve?)                              ║
║    • Trained brain checkpoints (~6KB — the learned strategy)           ║
║                                                                        ║
║  NOT collected:                                                        ║
║    • Model weights, training data, file paths, user identity           ║
║    • Nothing that could identify you or your work                      ║
║                                                                        ║
║  Stored locally: ~/.opop/telemetry/                                    ║
║  Share with: opop contribute                                           ║
║  Disable:    OPOP_TELEMETRY=0 or telemetry=False                      ║
║                                                                        ║
║  Don't like it? Turn it off. We're doing everyone a favor.             ║
║                                                                        ║
╚══════════════════════════════════════════════════════════════════════════╝
"""


# ═════════════════════════════════════════════════════════════════════════════
# HELPERS
# ═════════════════════════════════════════════════════════════════════════════

def _is_enabled():
    """Check if telemetry is enabled via environment variable."""
    val = os.environ.get("OPOP_TELEMETRY", "1").strip().lower()
    return val not in ("0", "false", "no", "off", "disable")


def _ensure_dirs():
    """Create telemetry directories if they don't exist."""
    for d in (OPOP_DIR, TELEMETRY_DIR, BRAINS_DIR, EXPERIENCE_DIR):
        d.mkdir(parents=True, exist_ok=True)


def _show_notice():
    """Print the disclaimer once per machine."""
    if NOTICE_FLAG.exists():
        return
    print(DISCLAIMER)
    try:
        _ensure_dirs()
        NOTICE_FLAG.write_text(str(time.time()))
    except OSError:
        pass  # can't write flag, will show again next time


def _get_endpoint():
    """Get the upload endpoint URL, if configured."""
    if ENDPOINT_FILE.exists():
        url = ENDPOINT_FILE.read_text().strip()
        if url:
            return url
    return os.environ.get("OPOP_ENDPOINT", "")


# ═════════════════════════════════════════════════════════════════════════════
# TELEMETRY COLLECTOR
# ═════════════════════════════════════════════════════════════════════════════

class Telemetry:
    """Collects and stores optimizer experience for crowdsourced improvement.

    Usage (automatic — wired into OptBrain/TwinCrack by default):

        # Telemetry is ON by default. Just use OpOp normally.
        optimizer = OptBrain(base_optimizer)

        # To disable:
        optimizer = OptBrain(base_optimizer, telemetry=False)
        # or: OPOP_TELEMETRY=0

    Manual usage:

        tel = Telemetry(mode='optbrain', n_groups=5, obs_dim=36)
        tel.record(obs_vector, decisions_dict, reward)
        tel.save_brain(brain.state_dict(), final_loss=2.1)
        tel.flush()

    CLI:
        opop telemetry status   — show collection stats
        opop contribute         — package & optionally upload
    """

    def __init__(self, mode='optbrain', n_groups=1, obs_dim=6,
                 framework='numpy', flush_every=FLUSH_EVERY):
        """
        Args:
            mode: 'optbrain' or 'twincrack'
            n_groups: number of parameter groups
            obs_dim: observation vector dimensionality
            framework: 'torch' or 'numpy'
            flush_every: buffer N experiences before writing to disk
        """
        self.enabled = _is_enabled()
        self._step_count = 0
        self._flush_count = 0
        self._obs_buf = []
        self._dec_buf = []
        self._rew_buf = []
        self.run_id = None

        if not self.enabled:
            return

        _show_notice()
        _ensure_dirs()

        self.run_id = uuid.uuid4().hex[:12]
        self.mode = mode
        self.n_groups = n_groups
        self.obs_dim = obs_dim
        self.framework = framework
        self.flush_every = flush_every
        self._start_time = time.time()

        # Run metadata
        self._meta = {
            'run_id': self.run_id,
            'mode': mode,
            'n_groups': n_groups,
            'obs_dim': obs_dim,
            'framework': framework,
            'start_time': self._start_time,
        }

    def record(self, obs, decisions, reward):
        """Record one experience step.

        Args:
            obs: observation vector (numpy array, shape [obs_dim])
            decisions: dict {group_idx: (scale, clip, dampen/momentum/...)}
            reward: float reward signal
        """
        if not self.enabled:
            return

        self._step_count += 1

        # Store observation
        if isinstance(obs, np.ndarray):
            self._obs_buf.append(obs.copy())
        else:
            self._obs_buf.append(np.array(obs, dtype=np.float32))

        # Flatten decisions to fixed-size array
        dec_flat = []
        for gi in range(self.n_groups):
            vals = decisions.get(gi, (1.0, 1.0, 0.0))
            dec_flat.extend([float(v) for v in vals])
        self._dec_buf.append(np.array(dec_flat, dtype=np.float32))

        self._rew_buf.append(float(reward))

        # Auto-flush
        if len(self._obs_buf) >= self.flush_every:
            self.flush()

    def flush(self):
        """Write buffered experience to disk."""
        if not self.enabled or not self._obs_buf:
            return

        try:
            obs = np.stack(self._obs_buf)
            dec = np.stack(self._dec_buf)
            rew = np.array(self._rew_buf, dtype=np.float32)

            fname = f"{self.run_id}_{self._flush_count:04d}.npz"
            path = EXPERIENCE_DIR / fname

            np.savez_compressed(str(path),
                obs=obs, decisions=dec, rewards=rew,
                step_offset=np.int32(self._step_count - len(self._obs_buf)),
                **{k: np.array(v) if not isinstance(v, (int, float, str))
                   else np.array(v) for k, v in self._meta.items()}
            )

            self._flush_count += 1
            self._obs_buf.clear()
            self._dec_buf.clear()
            self._rew_buf.clear()

        except (OSError, IOError):
            pass  # disk full, permissions, etc — silently skip

    def save_brain(self, brain_state_dict, final_loss=None, steps=None):
        """Save a trained brain checkpoint for crowdsourcing.

        This is the most valuable data point — a ~6KB brain that represents
        the distilled optimization strategy from an entire training run.

        Args:
            brain_state_dict: from brain.state_dict()
            final_loss: final training loss (optional)
            steps: total training steps (optional)
        """
        if not self.enabled:
            return

        try:
            save_dict = {}

            # Brain weights
            for k, v in brain_state_dict.items():
                if isinstance(v, np.ndarray):
                    save_dict[f'brain_{k}'] = v
                else:
                    save_dict[f'brain_{k}'] = np.array(v)

            # Metadata
            save_dict['run_id'] = np.array(self.run_id)
            save_dict['mode'] = np.array(self.mode)
            save_dict['n_groups'] = np.int32(self.n_groups)
            save_dict['obs_dim'] = np.int32(self.obs_dim)
            save_dict['framework'] = np.array(self.framework)
            save_dict['steps'] = np.int32(steps or self._step_count)
            save_dict['duration_sec'] = np.float32(time.time() - self._start_time)

            if final_loss is not None:
                save_dict['final_loss'] = np.float32(final_loss)

            fname = f"brain_{self.run_id}.npz"
            path = BRAINS_DIR / fname
            np.savez_compressed(str(path), **save_dict)

        except (OSError, IOError):
            pass

    def finalize(self, brain_state_dict=None, final_loss=None):
        """Call at end of training. Flushes buffer and saves brain."""
        if not self.enabled:
            return

        self.flush()

        if brain_state_dict is not None:
            self.save_brain(brain_state_dict, final_loss=final_loss,
                           steps=self._step_count)

        # Save run summary
        try:
            summary = {
                'run_id': self.run_id,
                'mode': self.mode,
                'n_groups': self.n_groups,
                'obs_dim': self.obs_dim,
                'framework': self.framework,
                'steps': self._step_count,
                'duration_sec': time.time() - self._start_time,
                'experience_files': self._flush_count,
            }
            if final_loss is not None:
                summary['final_loss'] = float(final_loss)

            summary_path = TELEMETRY_DIR / f"run_{self.run_id}.json"
            with open(str(summary_path), 'w') as f:
                json.dump(summary, f, indent=2)

        except (OSError, IOError):
            pass

    @property
    def stats(self):
        """Return collection statistics."""
        return {
            'run_id': self.run_id if self.enabled else None,
            'enabled': self.enabled,
            'steps': self._step_count if self.enabled else 0,
            'buffered': len(self._obs_buf) if self.enabled else 0,
            'flushed_files': self._flush_count if self.enabled else 0,
        }


# ═════════════════════════════════════════════════════════════════════════════
# GLOBAL STATS — for CLI
# ═════════════════════════════════════════════════════════════════════════════

def get_collection_stats():
    """Get statistics about all collected telemetry data."""
    stats = {
        'enabled': _is_enabled(),
        'telemetry_dir': str(TELEMETRY_DIR),
        'brains': 0,
        'experience_files': 0,
        'run_summaries': 0,
        'total_size_bytes': 0,
        'total_size_mb': 0.0,
    }

    if not TELEMETRY_DIR.exists():
        return stats

    # Count brains
    if BRAINS_DIR.exists():
        brain_files = list(BRAINS_DIR.glob("brain_*.npz"))
        stats['brains'] = len(brain_files)
        for f in brain_files:
            stats['total_size_bytes'] += f.stat().st_size

    # Count experience
    if EXPERIENCE_DIR.exists():
        exp_files = list(EXPERIENCE_DIR.glob("*.npz"))
        stats['experience_files'] = len(exp_files)
        for f in exp_files:
            stats['total_size_bytes'] += f.stat().st_size

    # Count run summaries
    run_files = list(TELEMETRY_DIR.glob("run_*.json"))
    stats['run_summaries'] = len(run_files)
    for f in run_files:
        stats['total_size_bytes'] += f.stat().st_size

    stats['total_size_mb'] = round(stats['total_size_bytes'] / (1024 * 1024), 2)

    return stats


def package_contribution(output_path=None):
    """Package all collected telemetry into a single .npz for sharing.

    Args:
        output_path: where to save (default: ~/.opop/contribution_<timestamp>.npz)

    Returns:
        str: path to the packaged file, or None if nothing to package
    """
    stats = get_collection_stats()
    if stats['brains'] == 0 and stats['experience_files'] == 0:
        return None

    if output_path is None:
        ts = int(time.time())
        output_path = str(OPOP_DIR / f"contribution_{ts}.npz")

    package = {}
    idx = 0

    # Collect brain checkpoints (most valuable)
    if BRAINS_DIR.exists():
        for f in sorted(BRAINS_DIR.glob("brain_*.npz")):
            try:
                data = np.load(str(f), allow_pickle=True)
                for k in data.files:
                    package[f'brain_{idx}_{k}'] = data[k]
                idx += 1
            except Exception:
                continue

    package['n_brains'] = np.int32(idx)

    # Collect experience trajectories
    exp_idx = 0
    if EXPERIENCE_DIR.exists():
        for f in sorted(EXPERIENCE_DIR.glob("*.npz")):
            try:
                data = np.load(str(f), allow_pickle=True)
                for k in data.files:
                    package[f'exp_{exp_idx}_{k}'] = data[k]
                exp_idx += 1
            except Exception:
                continue

    package['n_experience'] = np.int32(exp_idx)

    # Run summaries
    summaries = []
    for f in sorted(TELEMETRY_DIR.glob("run_*.json")):
        try:
            with open(str(f)) as fh:
                summaries.append(json.load(fh))
        except Exception:
            continue

    if summaries:
        package['run_summaries'] = np.array(json.dumps(summaries))

    np.savez_compressed(output_path, **package)
    return output_path


def upload_contribution(package_path, endpoint=None):
    """Upload a contribution package to the collection server.

    Args:
        package_path: path to the .npz package
        endpoint: URL to POST to (default: from config or env)

    Returns:
        bool: success
    """
    if endpoint is None:
        endpoint = _get_endpoint()

    if not endpoint:
        print("No upload endpoint configured.")
        print("Set endpoint: echo 'https://your-server.com/api/contribute' > ~/.opop/telemetry_endpoint.txt")
        print("Or: export OPOP_ENDPOINT=https://your-server.com/api/contribute")
        return False

    try:
        import urllib.request
        import urllib.error

        with open(package_path, 'rb') as f:
            data = f.read()

        req = urllib.request.Request(
            endpoint,
            data=data,
            headers={
                'Content-Type': 'application/octet-stream',
                'X-OpOp-Version': '0.3.0',
            },
            method='POST',
        )

        with urllib.request.urlopen(req, timeout=30) as resp:
            if resp.status == 200:
                return True
            else:
                print(f"Upload returned status {resp.status}")
                return False

    except ImportError:
        print("urllib not available")
        return False
    except urllib.error.URLError as e:
        print(f"Upload failed: {e}")
        return False
    except Exception as e:
        print(f"Upload error: {e}")
        return False


def clear_telemetry():
    """Delete all collected telemetry data."""
    import shutil
    count = 0

    for d in (BRAINS_DIR, EXPERIENCE_DIR):
        if d.exists():
            for f in d.iterdir():
                f.unlink()
                count += 1

    for f in TELEMETRY_DIR.glob("run_*.json"):
        f.unlink()
        count += 1

    return count
