#!/usr/bin/env python3
"""
Planner Router - Batch scheduling via QUBO optimization.

Composes BatchSchedulerQUBO with solver auto-selection for multi-job planning.
This is where QPU infrastructure becomes meaningful (64+ jobs = 192+ vars).

Usage:
    from flowtorch.planner_router import PlannerRouter, PlannerJob, PlannerConfig

    planner = PlannerRouter()
    planner.queue_job(PlannerJob(id=0, vram_mb=100, deadline_ms=500, se=0.3))
    planner.queue_job(PlannerJob(id=1, vram_mb=200, deadline_ms=300, se=0.7))
    result = planner.plan()
"""

import json
import time
import hashlib
import itertools
from dataclasses import dataclass, field
from typing import List, Dict, Optional, Tuple, Any
from pathlib import Path
from datetime import datetime, timezone

from .batch_scheduler_qubo import BatchSchedulerQUBO, Job, BACKENDS, N_BACKENDS


@dataclass
class PlannerJob:
    """
    A job to be scheduled by the planner.
    Bridges user input to BatchSchedulerQUBO.Job.
    """
    id: int
    vram_mb: float = 0.0
    deadline_ms: float = 1000.0
    se: float = 0.5
    expected_runtime_cpu_ms: float = 100.0
    expected_runtime_gpu_ms: float = 50.0
    expected_runtime_qpu_ms: float = 200.0
    priority: float = 1.0

    def to_scheduler_job(self) -> Job:
        """Convert to BatchSchedulerQUBO.Job."""
        return Job(
            id=self.id,
            vram_mb=self.vram_mb,
            deadline_ms=self.deadline_ms,
            se=self.se,
            expected_runtime_cpu_ms=self.expected_runtime_cpu_ms,
            expected_runtime_gpu_ms=self.expected_runtime_gpu_ms,
            expected_runtime_qpu_ms=self.expected_runtime_qpu_ms,
            priority=self.priority,
        )


@dataclass
class PlannerConfig:
    """Configuration for the planner."""
    # Solver selection thresholds
    brute_force_max_jobs: int = 8  # Use brute force for ≤8 jobs (24 vars)
    sa_sweeps: int = 1000          # SA sweeps for larger problems
    sa_num_reads: int = 10         # Number of SA samples
    sa_seed: int = 42              # SA random seed

    # Auto-plan trigger
    auto_plan_batch_size: int = 16  # Trigger auto-plan when queue reaches this

    # Hardware
    gpu_vram_mb: float = 4000.0
    gpu_temp_c: float = 50.0
    cpu_load: float = 0.5
    qpu_available: bool = False    # Set True if real QPU access exists

    # Receipt output
    receipt_dir: Optional[Path] = None


@dataclass
class PlannerResult:
    """Result of batch planning."""
    assignments: List[Dict[str, Any]]
    backend_counts: Dict[str, int]
    energy: float
    solve_method: str
    solve_time_ms: float
    constraint_violations: int
    n_jobs: int
    n_vars: int
    verification: Optional[Dict[str, Any]] = None
    receipt_path: Optional[Path] = None


class PlannerRouter:
    """
    Batch scheduler that composes BatchSchedulerQUBO with solver auto-selection.

    Provides:
    - queue_job(): Add jobs to queue
    - plan(): Execute batch optimization
    - generate_receipt(): Create verifiable receipt

    Solver auto-selection:
    - ≤8 jobs (≤24 vars): Brute force enumeration of valid one-hot combinations
    - >8 jobs (>24 vars): Simulated Annealing via solve_qubo_neal()
    - >64 jobs (future): D-Wave QPU if authenticated
    """

    def __init__(self, config: Optional[PlannerConfig] = None):
        self.config = config or PlannerConfig()
        self.queue: List[PlannerJob] = []
        self._last_result: Optional[PlannerResult] = None

    def queue_job(self, job: PlannerJob) -> bool:
        """
        Add a job to the queue.

        Returns True if queue size reached auto_plan_batch_size (trigger hint).
        """
        self.queue.append(job)
        return len(self.queue) >= self.config.auto_plan_batch_size

    def queue_jobs(self, jobs: List[PlannerJob]) -> bool:
        """Add multiple jobs to the queue."""
        self.queue.extend(jobs)
        return len(self.queue) >= self.config.auto_plan_batch_size

    def clear_queue(self):
        """Clear the job queue."""
        self.queue = []

    def plan(self, telemetry: Optional[Dict[str, Any]] = None) -> PlannerResult:
        """
        Execute batch planning on current queue.

        Args:
            telemetry: Optional dict with 'gpu_temp_c', 'cpu_load', etc.

        Returns:
            PlannerResult with assignments and diagnostics.
        """
        if not self.queue:
            raise ValueError("No jobs in queue to plan")

        start_time = time.time()

        # Apply telemetry overrides
        gpu_temp_c = self.config.gpu_temp_c
        cpu_load = self.config.cpu_load
        if telemetry:
            gpu_temp_c = telemetry.get('gpu_temp_c', gpu_temp_c)
            cpu_load = telemetry.get('cpu_load', cpu_load)

        # Convert queue to scheduler jobs
        scheduler_jobs = [j.to_scheduler_job() for j in self.queue]
        n_jobs = len(scheduler_jobs)

        # Build QUBO (this determines actual n_vars including slack bits)
        scheduler = BatchSchedulerQUBO(
            jobs=scheduler_jobs,
            gpu_vram_mb=self.config.gpu_vram_mb,
            gpu_temp_c=gpu_temp_c,
            cpu_load=cpu_load,
            qpu_available=self.config.qpu_available,
        )
        qubo = scheduler.build_qubo()

        # Get actual n_vars from QUBO (includes slack bits)
        n_vars = qubo['n']

        # Select solver and solve
        if n_jobs <= self.config.brute_force_max_jobs:
            bits, energy = self._solve_brute_force(qubo, n_jobs)
            solve_method = "brute_force"
        else:
            bits, energy, solve_method = self._solve_sa(qubo)

        solve_time_ms = (time.time() - start_time) * 1000

        # Decode solution
        decoded = scheduler.decode_solution(bits)

        # Build result
        result = PlannerResult(
            assignments=decoded['assignments'],
            backend_counts=decoded['summary']['backend_counts'],
            energy=energy,
            solve_method=solve_method,
            solve_time_ms=solve_time_ms,
            constraint_violations=decoded['summary']['constraint_violations'],
            n_jobs=n_jobs,
            n_vars=n_vars,
            verification=decoded.get('verification'),
        )

        # Generate receipt if configured
        if self.config.receipt_dir:
            result.receipt_path = self._generate_receipt(result, qubo)

        self._last_result = result
        return result

    def _solve_brute_force(
        self,
        qubo: Dict,
        n_jobs: int
    ) -> Tuple[List[int], float]:
        """
        Solve QUBO by brute force enumeration of valid one-hot combinations.

        For n_jobs, there are 3^n_jobs valid assignments (each job picks one backend).
        This is tractable for small n (3^8 = 6561 combinations).

        For each job assignment, we deterministically set slack bits to satisfy
        the VRAM constraint: slack = cap - used_vram_chunks.
        """
        n_vars = qubo['n']
        metadata = qubo['metadata']
        n_assignment_vars = metadata['n_assignment_vars']
        n_vram_slack = metadata['n_vram_slack']
        vram_cap_chunks = metadata['vram_cap_chunks']
        vram_chunk_mb = metadata['vram_chunk_mb']

        # Sanity check: assignment var layout must match our indexing assumption
        if n_assignment_vars != n_jobs * N_BACKENDS:
            raise ValueError(
                f"Layout mismatch: n_assignment_vars={n_assignment_vars} != "
                f"n_jobs*N_BACKENDS={n_jobs * N_BACKENDS}"
            )
        if n_vars != n_assignment_vars + n_vram_slack:
            raise ValueError(
                f"Total vars mismatch: n_vars={n_vars} != "
                f"n_assignment_vars + n_vram_slack={n_assignment_vars + n_vram_slack}"
            )

        # Parse QUBO coefficients
        linear = {int(k): float(v) for k, v in qubo['linear'].items()}
        quadratic = {}
        for k, v in qubo['quadratic'].items():
            i, j = map(int, k.split(','))
            quadratic[(i, j)] = float(v)

        # Precompute VRAM chunks per job (for slack calculation)
        job_vram_chunks = []
        for job in self.queue:
            chunks = max(1, int((job.vram_mb + vram_chunk_mb - 1) // vram_chunk_mb))
            job_vram_chunks.append(chunks)

        gpu_backend_idx = 1  # gpu is backend index 1

        min_energy = float('inf')
        best_bits = [0] * n_vars

        # Enumerate all valid one-hot combinations
        # Each job picks exactly one backend (0, 1, or 2)
        for combo in itertools.product(range(N_BACKENDS), repeat=n_jobs):
            # Build bit vector for assignments
            bits = [0] * n_vars
            used_vram_chunks = 0

            for job_idx, backend_idx in enumerate(combo):
                var_idx = job_idx * N_BACKENDS + backend_idx
                bits[var_idx] = 1

                # Track GPU VRAM usage
                if backend_idx == gpu_backend_idx:
                    used_vram_chunks += job_vram_chunks[job_idx]

            # Skip infeasible: VRAM over capacity
            if used_vram_chunks > vram_cap_chunks:
                continue

            # Set slack bits to satisfy constraint: used + slack = cap
            slack_value = vram_cap_chunks - used_vram_chunks
            for k in range(n_vram_slack):
                slack_var_idx = n_assignment_vars + k
                bits[slack_var_idx] = (slack_value >> k) & 1

            # Compute energy: E = Σ h_i x_i + Σ J_ij x_i x_j
            energy = 0.0

            # Linear terms
            for i, h in linear.items():
                energy += h * bits[i]

            # Quadratic terms
            for (i, j), J in quadratic.items():
                energy += J * bits[i] * bits[j]

            if energy < min_energy:
                min_energy = energy
                best_bits = bits.copy()

        return best_bits, min_energy

    def _solve_sa(self, qubo: Dict) -> Tuple[List[int], float, str]:
        """
        Solve QUBO using simulated annealing via solve_qubo_neal.
        Returns (bits, energy, method_name).
        """
        import sys
        from pathlib import Path

        # Add scripts to path for qubo_solve
        base_dir = Path(__file__).parent.parent
        scripts_dir = base_dir / "scripts"
        if str(scripts_dir) not in sys.path:
            sys.path.insert(0, str(scripts_dir))

        try:
            from qubo_solve import solve_qubo_neal
        except ImportError:
            # Fallback: random restart greedy (not real SA)
            bits, energy = self._solve_random_restart_greedy(qubo)
            return bits, energy, "random_restart_greedy"

        result = solve_qubo_neal(
            qubo,
            seed=self.config.sa_seed,
            sweeps=self.config.sa_sweeps,
            num_reads=self.config.sa_num_reads,
        )

        return result['bits'], result['energy'], "simulated_annealing"

    def _solve_random_restart_greedy(self, qubo: Dict) -> Tuple[List[int], float]:
        """
        Random restart greedy solver (fallback if dwave-samplers not available).
        NOT real simulated annealing - just generates random valid solutions
        and keeps the best one.
        """
        import random

        n_vars = qubo['n']
        metadata = qubo['metadata']
        n_assignment_vars = metadata['n_assignment_vars']
        n_vram_slack = metadata['n_vram_slack']
        vram_cap_chunks = metadata['vram_cap_chunks']
        vram_chunk_mb = metadata['vram_chunk_mb']
        n_jobs = metadata['n_jobs']

        # Sanity check: assignment var layout must match our indexing assumption
        if n_assignment_vars != n_jobs * N_BACKENDS:
            raise ValueError(
                f"Layout mismatch: n_assignment_vars={n_assignment_vars} != "
                f"n_jobs*N_BACKENDS={n_jobs * N_BACKENDS}"
            )
        if n_vars != n_assignment_vars + n_vram_slack:
            raise ValueError(
                f"Total vars mismatch: n_vars={n_vars} != "
                f"n_assignment_vars + n_vram_slack={n_assignment_vars + n_vram_slack}"
            )

        linear = {int(k): float(v) for k, v in qubo['linear'].items()}
        quadratic = {}
        for k, v in qubo['quadratic'].items():
            i, j = map(int, k.split(','))
            quadratic[(i, j)] = float(v)

        # Precompute VRAM chunks per job
        job_vram_chunks = []
        for job in self.queue:
            chunks = max(1, int((job.vram_mb + vram_chunk_mb - 1) // vram_chunk_mb))
            job_vram_chunks.append(chunks)

        gpu_backend_idx = 1

        def compute_energy(bits):
            energy = 0.0
            for i, h in linear.items():
                energy += h * bits[i]
            for (i, j), J in quadratic.items():
                energy += J * bits[i] * bits[j]
            return energy

        def generate_valid_solution():
            """Generate random valid one-hot assignment with proper slack."""
            bits = [0] * n_vars
            used_vram_chunks = 0

            for job_idx in range(n_jobs):
                backend = random.randint(0, N_BACKENDS - 1)
                bits[job_idx * N_BACKENDS + backend] = 1
                if backend == gpu_backend_idx:
                    used_vram_chunks += job_vram_chunks[job_idx]

            # If over capacity, move some GPU jobs to CPU
            while used_vram_chunks > vram_cap_chunks:
                # Find a GPU job to move
                for job_idx in range(n_jobs):
                    gpu_var = job_idx * N_BACKENDS + gpu_backend_idx
                    if bits[gpu_var] == 1:
                        bits[gpu_var] = 0
                        bits[job_idx * N_BACKENDS + 0] = 1  # Move to CPU
                        used_vram_chunks -= job_vram_chunks[job_idx]
                        break

            # Set slack bits
            slack_value = vram_cap_chunks - used_vram_chunks
            for k in range(n_vram_slack):
                bits[n_assignment_vars + k] = (slack_value >> k) & 1

            return bits

        random.seed(self.config.sa_seed)

        best_bits = generate_valid_solution()
        best_energy = compute_energy(best_bits)

        # Multiple random restarts
        for _ in range(self.config.sa_num_reads * 100):
            bits = generate_valid_solution()
            energy = compute_energy(bits)
            if energy < best_energy:
                best_energy = energy
                best_bits = bits.copy()

        return best_bits, best_energy

    def _generate_receipt(
        self,
        result: PlannerResult,
        qubo: Dict
    ) -> Path:
        """Generate verifiable receipt for the planning run."""
        receipt_dir = self.config.receipt_dir or Path("receipts/planner")
        receipt_dir.mkdir(parents=True, exist_ok=True)

        timestamp = datetime.now(timezone.utc)
        timestamp_str = timestamp.strftime("%Y%m%d_%H%M%S")

        receipt = {
            "receipt_type": "QUBO_BATCH_PLANNER",
            "schema": "qubo_batch_planner:v1",
            "timestamp_utc": timestamp.isoformat(),
            "inputs": {
                "n_jobs": result.n_jobs,
                "n_vars": result.n_vars,
                "gpu_vram_mb": self.config.gpu_vram_mb,
                "gpu_temp_c": self.config.gpu_temp_c,
                "qpu_available": self.config.qpu_available,
            },
            "solver": {
                "method": result.solve_method,
                "solve_time_ms": round(result.solve_time_ms, 3),
                "config": {
                    "brute_force_max_jobs": self.config.brute_force_max_jobs,
                    "sa_sweeps": self.config.sa_sweeps,
                    "sa_num_reads": self.config.sa_num_reads,
                    "sa_seed": self.config.sa_seed,
                }
            },
            "results": {
                "energy": round(result.energy, 4),
                "backend_counts": result.backend_counts,
                "constraint_violations": result.constraint_violations,
                "assignments_sample": result.assignments[:10],  # First 10 for brevity
            },
            "verification": result.verification,
            "qubo_metadata": {
                "linear_terms": len(qubo['linear']),
                "quadratic_terms": len(qubo['quadratic']),
                "penalty_weights": qubo['metadata'].get('penalty_weights', {}),
            }
        }

        # Add receipt hash
        receipt_json = json.dumps(receipt, sort_keys=True)
        receipt["receipt_sha256"] = hashlib.sha256(receipt_json.encode()).hexdigest()

        # Write receipt
        receipt_path = receipt_dir / f"planner_{timestamp_str}_{result.n_jobs}jobs.json"
        receipt_path.write_text(json.dumps(receipt, indent=2))

        return receipt_path

    def generate_receipt(self) -> Optional[Path]:
        """Generate receipt for last planning run."""
        if self._last_result is None:
            return None

        # Build minimal QUBO for receipt metadata
        scheduler_jobs = [j.to_scheduler_job() for j in self.queue]
        scheduler = BatchSchedulerQUBO(
            jobs=scheduler_jobs,
            gpu_vram_mb=self.config.gpu_vram_mb,
            gpu_temp_c=self.config.gpu_temp_c,
            qpu_available=self.config.qpu_available,
        )
        qubo = scheduler.build_qubo()

        return self._generate_receipt(self._last_result, qubo)

    def get_last_result(self) -> Optional[PlannerResult]:
        """Get the result of the last planning run."""
        return self._last_result

    def info(self) -> Dict:
        """Return planner state information."""
        return {
            "queue_size": len(self.queue),
            "config": {
                "brute_force_max_jobs": self.config.brute_force_max_jobs,
                "sa_sweeps": self.config.sa_sweeps,
                "auto_plan_batch_size": self.config.auto_plan_batch_size,
                "qpu_available": self.config.qpu_available,
            },
            "last_result": {
                "n_jobs": self._last_result.n_jobs,
                "solve_method": self._last_result.solve_method,
                "energy": self._last_result.energy,
            } if self._last_result else None,
        }


def generate_synthetic_planner_jobs(
    n_jobs: int,
    seed: int = 42
) -> List[PlannerJob]:
    """Generate synthetic jobs for testing."""
    import random
    random.seed(seed)

    jobs = []
    for i in range(n_jobs):
        job_type = random.choice(['light', 'medium', 'heavy', 'urgent'])

        if job_type == 'light':
            se = random.uniform(0.05, 0.3)
            vram = random.uniform(0, 100)
            deadline = random.uniform(500, 2000)
            cpu_ms, gpu_ms, qpu_ms = 50, 30, 250
        elif job_type == 'medium':
            se = random.uniform(0.3, 0.6)
            vram = random.uniform(100, 500)
            deadline = random.uniform(200, 1000)
            cpu_ms, gpu_ms, qpu_ms = 200, 80, 300
        elif job_type == 'heavy':
            se = random.uniform(0.6, 0.95)
            vram = random.uniform(500, 2000)
            deadline = random.uniform(100, 500)
            cpu_ms, gpu_ms, qpu_ms = 500, 100, 350
        else:  # urgent
            se = random.uniform(0.4, 0.8)
            vram = random.uniform(200, 800)
            deadline = random.uniform(50, 200)
            cpu_ms, gpu_ms, qpu_ms = 150, 50, 400

        jobs.append(PlannerJob(
            id=i,
            vram_mb=vram,
            deadline_ms=deadline,
            se=se,
            expected_runtime_cpu_ms=cpu_ms * random.uniform(0.8, 1.2),
            expected_runtime_gpu_ms=gpu_ms * random.uniform(0.8, 1.2),
            expected_runtime_qpu_ms=qpu_ms * random.uniform(0.8, 1.2),
            priority=random.uniform(0.5, 2.0) if job_type == 'urgent' else 1.0,
        ))

    return jobs


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="Planner Router Demo")
    parser.add_argument('--n-jobs', '-n', type=int, default=16, help='Number of jobs')
    parser.add_argument('--seed', type=int, default=42, help='Random seed')
    parser.add_argument('--receipt-dir', type=str, default=None, help='Receipt output dir')

    args = parser.parse_args()

    print(f"Planner Router Demo")
    print(f"=" * 50)
    print(f"Jobs: {args.n_jobs}")
    print()

    # Configure planner
    config = PlannerConfig(
        receipt_dir=Path(args.receipt_dir) if args.receipt_dir else Path("receipts/planner"),
        sa_seed=args.seed,
    )
    planner = PlannerRouter(config)

    # Generate and queue jobs
    jobs = generate_synthetic_planner_jobs(args.n_jobs, args.seed)
    planner.queue_jobs(jobs)

    # Plan
    print(f"Planning {args.n_jobs} jobs...")
    result = planner.plan()

    # Results
    print(f"\nResults:")
    print(f"  Solve method: {result.solve_method}")
    print(f"  Solve time: {result.solve_time_ms:.2f} ms")
    print(f"  Energy: {result.energy:.2f}")
    print(f"  Backend distribution: {result.backend_counts}")
    print(f"  Constraint violations: {result.constraint_violations}")

    print(f"\nAssignments (first 10):")
    for a in result.assignments[:10]:
        print(f"  Job {a['job_id']:2d}: {a['backend']:4s} (Se={a['se']:.2f})")
    if len(result.assignments) > 10:
        print(f"  ... and {len(result.assignments) - 10} more")

    if result.receipt_path:
        print(f"\nReceipt: {result.receipt_path}")
