#!/usr/bin/env python3
"""
orchestrator.py - GGUF Layer-by-Layer Encoding Orchestrator
============================================================

Orchestrates full GGUF model encoding:
1. Scan GGUF for all tensors
2. For each tensor: extract → encode → store
3. Generate manifest with structural verification
"""

import json
from pathlib import Path
from datetime import datetime, timezone
from typing import Dict, Any, List, Optional
import numpy as np

from helix_cdc.gguf.scanner import GGUFScanner
from helix_cdc.gguf.encoder import LayerEncoder


class NumpyEncoder(json.JSONEncoder):
    """JSON encoder that handles numpy types."""
    def default(self, obj):
        if isinstance(obj, np.bool_):
            return bool(obj)
        if isinstance(obj, np.integer):
            return int(obj)
        if isinstance(obj, np.floating):
            return float(obj)
        if isinstance(obj, np.ndarray):
            return obj.tolist()
        return super().default(obj)


class GGUFOrchestrator:
    """Orchestrate full GGUF model encoding."""

    def __init__(self, gguf_path: Path, output_dir: Path,
                 k_effective: int = 64,
                 layer_filter: Optional[str] = None,
                 max_layers: Optional[int] = None):
        self.gguf_path = Path(gguf_path)
        self.output_dir = Path(output_dir)
        self.k_effective = k_effective
        self.layer_filter = layer_filter
        self.max_layers = max_layers

        self.scanner = GGUFScanner(gguf_path)
        self.encoder = LayerEncoder(k_effective=k_effective)

    def run(self, verbose: bool = True) -> Dict[str, Any]:
        """Run full orchestration pipeline."""
        timestamp = datetime.now(timezone.utc).strftime('%Y%m%dT%H%M%SZ')

        if verbose:
            print("=" * 70)
            print("GGUF LAYER-BY-LAYER ORCHESTRATOR")
            print("=" * 70)

        # Step 1: Scan GGUF
        if verbose:
            print("\n[Step 1] Scanning GGUF...")
        self.scanner.scan()

        # Filter tensors
        tensors = self.scanner.tensors
        if self.layer_filter:
            tensors = [t for t in tensors if self.layer_filter in t['name']]
            if verbose:
                print(f"  Filtered to {len(tensors)} tensors matching '{self.layer_filter}'")

        if self.max_layers:
            tensors = tensors[:self.max_layers]
            if verbose:
                print(f"  Limited to first {len(tensors)} tensors")

        # Create output directory
        self.output_dir.mkdir(parents=True, exist_ok=True)
        layers_dir = self.output_dir / 'layers'
        layers_dir.mkdir(exist_ok=True)

        # Step 2: Extract and encode each tensor
        if verbose:
            print(f"\n[Step 2] Processing {len(tensors)} tensors...")

        encoded_layers = []
        total_original_bytes = 0
        total_encoded_bytes = 0

        for i, tensor_info in enumerate(tensors):
            name = tensor_info['name']
            if verbose:
                print(f"\n  [{i+1}/{len(tensors)}] {name}")
                print(f"    Shape: {tensor_info['dims']}, Type: {tensor_info['dtype_name']}")

            # Extract
            tensor = self.scanner.extract_tensor(tensor_info)
            if tensor is None:
                if verbose:
                    print(f"    SKIP: Failed to extract")
                continue

            total_original_bytes += tensor.nbytes
            if verbose:
                print(f"    Extracted: {tensor.shape}, {tensor.nbytes:,} bytes")

            # Encode
            encoded = self.encoder.encode(tensor, name)

            if verbose:
                print(f"    SVD k={encoded['k_effective']}, error={encoded['recon_error']:.2e}, energy={encoded['energy_ratio']:.3f}")

            # Save layer
            layer_file = layers_dir / f"{name.replace('.', '_').replace('/', '_')}.npz"
            layer_bytes = self.encoder.save_encoded(encoded, str(layer_file))

            total_encoded_bytes += layer_bytes
            if verbose:
                print(f"    Saved: {layer_file.name} ({layer_bytes:,} bytes)")

            # Record without numpy arrays
            encoded_clean = {k: v for k, v in encoded.items() if not k.startswith('_')}
            encoded_clean['file'] = layer_file.name
            encoded_clean['file_bytes'] = layer_bytes
            encoded_layers.append(encoded_clean)

        # Step 3: Generate manifest
        if verbose:
            print(f"\n[Step 3] Generating manifest...")

        manifest = {
            'schema': 'GGUF_ORCHESTRATOR_V1',
            'protocol_version': 'helix_cdc:v0.3.0-FROZEN-SUBSTRATE',
            'timestamp_utc': timestamp,
            'source_gguf': str(self.gguf_path),
            'source_metadata': {k: v for k, v in self.scanner.metadata.items()
                               if isinstance(v, (str, int, float))},
            'encoding': {
                'method': 'truncated_SVD',
                'k_effective': self.k_effective,
            },
            'summary': {
                'total_tensors': len(self.scanner.tensors),
                'encoded_tensors': len(encoded_layers),
                'original_bytes': total_original_bytes,
                'encoded_bytes': total_encoded_bytes,
                'compression_ratio': total_original_bytes / total_encoded_bytes if total_encoded_bytes > 0 else 0,
            },
            'layers': encoded_layers,
        }

        manifest_path = self.output_dir / 'manifest.json'
        manifest_path.write_text(json.dumps(manifest, indent=2, cls=NumpyEncoder))

        if verbose:
            print(f"  Manifest: {manifest_path}")
            print("\n" + "=" * 70)
            print("ORCHESTRATION COMPLETE")
            print("=" * 70)
            print(f"  Source: {self.gguf_path}")
            print(f"  Output: {self.output_dir}")
            print(f"  Tensors: {len(encoded_layers)}/{len(self.scanner.tensors)}")
            print(f"  Original: {total_original_bytes:,} bytes")
            print(f"  Encoded:  {total_encoded_bytes:,} bytes")
            if total_encoded_bytes > 0:
                print(f"  Ratio:    {total_original_bytes/total_encoded_bytes:.2f}x")

        return manifest

    @staticmethod
    def verify_manifest(manifest_path: Path) -> Dict[str, Any]:
        """Verify a manifest and its layers."""
        manifest = json.loads(manifest_path.read_text())
        layers_dir = manifest_path.parent / 'layers'

        results = {
            'manifest_valid': True,
            'layers_found': 0,
            'layers_missing': [],
            'total_bytes': 0,
        }

        for layer in manifest.get('layers', []):
            layer_file = layers_dir / layer['file']
            if layer_file.exists():
                results['layers_found'] += 1
                results['total_bytes'] += layer_file.stat().st_size
            else:
                results['layers_missing'].append(layer['file'])
                results['manifest_valid'] = False

        return results
