"""
reproducibility_proof/verify.py

Core reproducibility verification engine for Marco.

Provides:
  - verify(version_id, repo_path)    -> verifies a single version
  - verify_all(repo_path)            -> verifies every version in the repo

Features:
  1. Step-by-step intermediate hash verification — pinpoints WHICH DAG step broke.
  2. Byte-exact canonical enforcer via canonicalizer.py — prevents false FAILs
     due to OS line-ending differences (CRLF vs LF).
  3. Row-level delta on failure — outputs exactly which rows changed.
  4. Verification audit log — appended to .marco/verification_log.jsonl on every run.
  5. Environment fingerprinting — warns if Python/Numpy/Pandas version changed.
  6. Stochastic operation detector — warns if pipeline has non-deterministic steps.
  7. Sets manifest["verified"] = True on PASS.
"""

import hashlib
import json
import os
import sys
import platform
from datetime import datetime, timezone
from pathlib import Path
from typing import Optional, List, Tuple

from .canonicalizer import compute_canonical_hash, to_canonical_bytes, detect_stochastic_ops
from .audit_log import append_audit_entry

# ---------------------------------------------------------------------------
# Internal helpers — replicate core pipeline logic (stdlib + numpy/pandas only)
# ---------------------------------------------------------------------------

def _read_raw_rows(raw_path: Path) -> List[Tuple[str, str]]:
    text = raw_path.read_text(encoding='utf-8', errors='replace')
    text = text.replace('\r\n', '\n').replace('\r', '\n')
    lines = [ln for ln in text.split('\n') if ln.strip()]
    rows = []
    for ln in lines:
        if '\t' in ln:
            lbl, txt = ln.split('\t', 1)
        else:
            lbl, txt = '', ln
        rows.append((lbl.strip(), txt.strip()))
    return rows


def _run_step(func_name: str, params: dict, data):
    """Execute a single preprocessing step by function name."""
    # Import lazily so this module stays decoupled from core
    from marco.core.preprocessor import _NODE_REGISTRY
    if func_name not in _NODE_REGISTRY:
        raise ValueError(f"Unknown preprocessing function: '{func_name}'")
    return _NODE_REGISTRY[func_name](data, params)


def _rerun_pipeline_step_by_step(raw_rows, pipeline_spec: dict):
    """
    Re-execute the pipeline step by step in topological order.

    Returns:
        final_rows:              Output of the full pipeline.
        step_hashes:             {step_name: sha256_of_output_at_that_step}
    """
    from marco.core.preprocessor import topological_sort_order, validate_dag

    validate_dag(pipeline_spec)
    order = topological_sort_order(pipeline_spec)

    results = {'raw': raw_rows}
    step_hashes = {}

    for node in order:
        node_spec = pipeline_spec[node]
        func_key = node_spec.get('func')
        params = node_spec.get('params', {})
        depends = node_spec.get('depends_on', [])

        if depends:
            inputs = []
            for d in depends:
                inputs.extend(results[d])
        else:
            inputs = results['raw']

        out_rows = _run_step(func_key, params, inputs)
        results[node] = out_rows
        step_hashes[node] = compute_canonical_hash(out_rows)

    # Find final nodes (not depended upon by anyone)
    depended = set()
    for n, s in pipeline_spec.items():
        for d in s.get('depends_on', []):
            depended.add(d)
    final_nodes = sorted([n for n in pipeline_spec if n not in depended])

    final_rows = []
    for fn in final_nodes:
        final_rows.extend(results[fn])

    return final_rows, step_hashes


def _compute_step_hashes_from_tsv(version_dir: Path, pipeline_spec: dict) -> dict:
    """
    Recompute per-step hashes stored at creation time from metrics_per_step.
    Since we only have final output_hash stored per step, we rely on the stored
    metrics for count-level sanity but use our recomputed step hashes for comparison.
    NOTE: Returns empty dict if per-step output hashes weren't stored (older versions).
    """
    # Marco stores metrics_per_step but NOT per-step output hashes.
    # We return an empty dict here; the step-by-step comparison
    # is done between the fresh reruns at each step vs the fresh final.
    return {}


def _row_level_delta(stored_rows: List, computed_rows: List) -> dict:
    """
    Compute which rows differ between stored and freshly computed output.
    Works on raw (label, text) tuples.

    Returns a dict with:
      - added:   rows in computed but not in stored
      - removed: rows in stored but not in computed
      - count_stored / count_computed
    """
    stored_set = set(stored_rows)
    computed_set = set(computed_rows)
    return {
        'count_stored': len(stored_rows),
        'count_computed': len(computed_rows),
        'added': list(computed_set - stored_set)[:20],    # cap at 20 examples
        'removed': list(stored_set - computed_set)[:20],
    }


def _compare_environments(stored_env: dict, current_env: dict) -> List[str]:
    """Return a list of human-readable warnings for environment differences."""
    import importlib.metadata
    warnings = []
    for key in ('python_version', 'platform'):
        if stored_env.get(key) and stored_env.get(key) != current_env.get(key):
            warnings.append(
                f"  {key} changed: '{stored_env[key]}' → '{current_env[key]}'"
            )
    for pkg in ('numpy', 'pandas'):
        k = f'{pkg}_version'
        if stored_env.get(k) and stored_env.get(k) != current_env.get(k):
            warnings.append(
                f"  {pkg} version changed: '{stored_env[k]}' → '{current_env[k]}'"
            )
    return warnings


def _get_current_env() -> dict:
    import importlib.metadata
    env = {
        'python_version': sys.version,
        'platform': platform.platform(),
        'os': os.name,
    }
    for pkg in ('numpy', 'pandas'):
        try:
            env[f'{pkg}_version'] = importlib.metadata.version(pkg)
        except Exception:
            env[f'{pkg}_version'] = 'unknown'
    return env


# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------

def verify(version_id: str, repo_path='.') -> bool:
    """
    Verify the reproducibility of a single dataset version.

    Steps:
      1. Warn if pipeline has stochastic (non-seeded) operations.
      2. Re-run the full pipeline step by step from the stored raw.txt.
      3. Compute canonical hash of fresh output.
      4. Compare to stored output_hash in manifest.json.
      5. On FAIL: write verification_report.json with row-level delta.
      6. On PASS: set manifest["verified"] = True.
      7. Append audit log entry regardless of outcome.

    Returns True on PASS, False on FAIL.
    """
    repo_path = Path(repo_path)
    marco_dir = repo_path / '.marco'
    version_dir = marco_dir / 'versions' / version_id

    if not version_dir.exists():
        print(f"❌ ERROR: Version '{version_id}' not found at {version_dir}")
        return False

    manifest_path = version_dir / 'manifest.json'
    config_path = version_dir / 'config.json'
    raw_path = version_dir / 'raw.txt'

    if not manifest_path.exists() or not config_path.exists() or not raw_path.exists():
        print(f"❌ ERROR: Version '{version_id}' is missing required files.")
        return False

    manifest = json.loads(manifest_path.read_text(encoding='utf-8'))
    config = json.loads(config_path.read_text(encoding='utf-8'))
    stored_hash = manifest.get('output_hash')
    pipeline_spec = config.get('dag', {})

    print(f"\n🔬 Verifying version: {version_id[:16]}...")

    # --- Step 0: Stochastic operation detection ---
    stochastic_warnings = detect_stochastic_ops(pipeline_spec)
    if stochastic_warnings:
        print("  ⚠️  Stochastic operations detected (results may not be reproducible):")
        for w in stochastic_warnings:
            print(f"     {w}")

    # --- Step 1: Environment comparison ---
    stored_env = manifest.get('environment', {})
    current_env = _get_current_env()
    env_warnings = _compare_environments(stored_env, current_env)
    if env_warnings:
        print("  ⚠️  Environment has changed since this version was created:")
        for w in env_warnings:
            print(w)

    # --- Step 2: Re-run pipeline step by step ---
    raw_rows = _read_raw_rows(raw_path)
    try:
        computed_rows, step_hashes = _rerun_pipeline_step_by_step(raw_rows, pipeline_spec)
    except Exception as e:
        print(f"  ❌ FAIL: Pipeline re-run crashed: {e}")
        append_audit_entry(marco_dir, version_id, 'FAIL', {'error': str(e)})
        return False

    # --- Step 3: Step-by-step hash report ---
    print("\n  Pipeline Step Verification:")
    stored_step_hashes = manifest.get('step_hashes', {})
    for step_name, fresh_hash in step_hashes.items():
        if step_name in stored_step_hashes:
            stored_step_hash = stored_step_hashes[step_name]
            icon = '  ✅' if fresh_hash == stored_step_hash else '  ❌'
            print(f"  {icon}  {step_name}")
        else:
            # No stored step hash; just report it was executed
            print(f"    ⬜  {step_name}  (no stored step hash to compare)")

    # --- Step 4: Final hash comparison ---
    computed_hash = compute_canonical_hash(computed_rows)
    print(f"\n  Stored  output_hash: {stored_hash}")
    print(f"  Recomputed   hash: {computed_hash}")

    passed = (computed_hash == stored_hash)

    if passed:
        print(f"\n  ✅ VERIFIED — output hash matches. Version is reproducible.\n")
        # Update manifest with verified flag and record environment
        manifest['verified'] = True
        manifest['verified_at'] = datetime.now(timezone.utc).isoformat().replace('+00:00', 'Z')
        manifest['environment'] = current_env
        manifest['step_hashes'] = step_hashes
        # Write back (chmod relaxed for this manifest only)
        _write_manifest_safe(manifest_path, manifest)
        append_audit_entry(marco_dir, version_id, 'PASS', {
            'stored_hash': stored_hash,
            'computed_hash': computed_hash,
        })
        return True
    else:
        print(f"\n  ❌ FAIL — hash mismatch! This version is NOT reproducible.\n")

        # --- Step 5: Row-level delta ---
        stored_tsv = version_dir / 'preprocessed.tsv'
        stored_rows = []
        if stored_tsv.exists():
            import csv
            with open(stored_tsv, 'r', encoding='utf-8', newline='') as f:
                reader = csv.DictReader(f, delimiter='\t')
                for row in reader:
                    stored_rows.append((row.get('label', ''), row.get('text', '')))

        delta = _row_level_delta(stored_rows, computed_rows)

        # Write verification report
        report = {
            'version_id': version_id,
            'verified_at': datetime.now(timezone.utc).isoformat().replace('+00:00', 'Z'),
            'result': 'FAIL',
            'stored_hash': stored_hash,
            'computed_hash': computed_hash,
            'step_hashes': step_hashes,
            'stochastic_warnings': stochastic_warnings,
            'environment_warnings': env_warnings,
            'row_delta': delta,
        }
        report_path = version_dir / 'verification_report.json'
        # Write report regardless of file permissions (use temp + replace)
        import tempfile
        tmp_fd, tmp_path = tempfile.mkstemp(
            prefix='vreport_', suffix='.json', dir=version_dir
        )
        try:
            with os.fdopen(tmp_fd, 'w', encoding='utf-8') as f:
                json.dump(report, f, indent=2)
            os.replace(tmp_path, str(report_path))
        except Exception:
            try:
                os.unlink(tmp_path)
            except Exception:
                pass

        print(f"  📄 Diagnostic report written to: {report_path}")
        print(f"     Rows in stored output:     {delta['count_stored']}")
        print(f"     Rows in recomputed output: {delta['count_computed']}")
        if delta['removed']:
            print(f"     Sample rows MISSING in recomputed output ({len(delta['removed'])} shown):")
            for lbl, txt in delta['removed'][:5]:
                preview = txt[:80] + '...' if len(txt) > 80 else txt
                print(f"       [{lbl}] {preview}")
        if delta['added']:
            print(f"     Sample rows NEW in recomputed output ({len(delta['added'])} shown):")
            for lbl, txt in delta['added'][:5]:
                preview = txt[:80] + '...' if len(txt) > 80 else txt
                print(f"       [{lbl}] {preview}")
        print()

        append_audit_entry(marco_dir, version_id, 'FAIL', {
            'stored_hash': stored_hash,
            'computed_hash': computed_hash,
            'stochastic_warnings': stochastic_warnings,
            'environment_warnings': env_warnings,
        })
        return False


def verify_all(repo_path='.') -> dict:
    """
    Verify every version in the Marco repository.

    Prints a summary table and returns a dict of {version_id: True/False}.
    """
    repo_path = Path(repo_path)
    versions_dir = repo_path / '.marco' / 'versions'

    if not versions_dir.exists():
        print("❌ No Marco repository found. Run 'marco init' first.")
        return {}

    version_ids = [
        d.name for d in versions_dir.iterdir()
        if d.is_dir() and (d / 'manifest.json').exists()
    ]

    if not version_ids:
        print("ℹ️  No versions found in this repository.")
        return {}

    print(f"\n{'─'*70}")
    print(f"  Marco Reproducibility Audit — {len(version_ids)} version(s) found")
    print(f"{'─'*70}")

    results = {}
    for vid in sorted(version_ids):
        passed = verify(vid, repo_path)
        results[vid] = passed

    total = len(results)
    passed_count = sum(1 for v in results.values() if v)
    print(f"{'─'*70}")
    print(f"  Summary: {passed_count}/{total} versions passed.")
    print(f"{'─'*70}\n")

    return results


# ---------------------------------------------------------------------------
# Internal helper
# ---------------------------------------------------------------------------

def _write_manifest_safe(manifest_path: Path, manifest: dict):
    """Write manifest JSON, handling read-only file permissions gracefully."""
    import tempfile
    import stat
    version_dir = manifest_path.parent
    tmp_fd, tmp_path = tempfile.mkstemp(prefix='manifest_', suffix='.json', dir=version_dir)
    try:
        with os.fdopen(tmp_fd, 'w', encoding='utf-8') as f:
            json.dump(manifest, f, indent=2)
        # Temporarily make original writable if read-only
        try:
            current_mode = os.stat(manifest_path).st_mode
            os.chmod(manifest_path, stat.S_IREAD | stat.S_IWRITE | stat.S_IRGRP | stat.S_IROTH)
            os.replace(tmp_path, str(manifest_path))
            os.chmod(manifest_path, stat.S_IREAD | stat.S_IRGRP | stat.S_IROTH)
        except PermissionError:
            os.replace(tmp_path, str(manifest_path))
    except Exception:
        try:
            os.unlink(tmp_path)
        except Exception:
            pass
        raise
