"""
LQNS Solver handler - Native Python implementation.

Calls LQNS (Layered Queueing Network Solver) via subprocess, bypassing JPype.

Port from:


Note: This handler works with LayeredNetworkStruct and requires the LQNX file
to be generated by the model. For full LQNX generation, use the Java wrapper.
"""

import numpy as np
import subprocess
import tempfile
import os
import shutil
from dataclasses import dataclass
from typing import Optional, Dict, Any
from xml.etree import ElementTree as ET
import time
import concurrent.futures


@dataclass
class SolverLQNSOptions:
    """Options for LQNS solver."""
    method: str = 'lqns'
    multiserver: str = 'rolia'
    samples: int = 10000
    seed: int = 23000
    iter_max: int = 50
    verbose: bool = False
    keep: bool = False


@dataclass
class SolverLQNSReturn:
    """
    Result of LQNS solver handler.

    Attributes:
        Q: Queue lengths (nidx x 1)
        U: Utilizations (nidx x 1)
        R: Response times (nidx x 1)
        T: Throughputs (nidx x 1)
        P: Processor utilizations (nidx x 1)
        S: Service times (nidx x 1)
        A: Arrival rates (nidx x 1)
        W: Waiting times (nidx x 1)
        runtime: Runtime in seconds
        method: Method used
        iterations: Number of iterations
    """
    Q: Optional[np.ndarray] = None
    U: Optional[np.ndarray] = None
    R: Optional[np.ndarray] = None
    T: Optional[np.ndarray] = None
    P: Optional[np.ndarray] = None  # Processor utilization
    S: Optional[np.ndarray] = None  # Service time
    A: Optional[np.ndarray] = None
    W: Optional[np.ndarray] = None
    runtime: float = 0.0
    method: str = "lqns"
    iterations: int = 0


def is_lqns_available() -> bool:
    """Check if lqns and lqsim commands are available in PATH."""
    try:
        # Use --help instead of --version because lqsim --version can hang
        result = subprocess.run(
            ['lqns', '--help'],
            capture_output=True,
            timeout=5
        )
        output = result.stdout.decode('utf-8', errors='ignore').lower()
        return 'lqns' in output or 'usage' in output or result.returncode == 0
    except (FileNotFoundError, subprocess.TimeoutExpired, OSError):
        return False


def _parse_lqxo_results(result_path: str, names: Dict[int, str], nidx: int, ncalls: int) -> tuple:
    """
    Parse LQNS XML output (.lqxo file).

    Args:
        result_path: Path to the .lqxo result file
        names: Dictionary mapping indices to element names
        nidx: Number of node indices
        ncalls: Number of calls

    Returns:
        Tuple of (U, P, S, T, W, iterations)
    """
    # Initialize result arrays
    U = np.full(nidx, np.nan)  # Utilization
    P = np.full(nidx, np.nan)  # Processor utilization
    S = np.full(nidx, np.nan)  # Service time (phase1)
    T = np.full(nidx, np.nan)  # Throughput
    W = np.full(nidx, np.nan)  # Proc waiting
    edges_W = np.full(ncalls, np.nan)  # Edge waiting
    iterations = 0

    if not os.path.exists(result_path):
        return U, P, S, T, W, edges_W, iterations

    try:
        tree = ET.parse(result_path)
        root = tree.getroot()

        # Parse iterations from solver-params
        for solver_params in root.iter('solver-params'):
            for result_general in solver_params.iter('result-general'):
                iter_str = result_general.get('iterations', '0')
                try:
                    iterations = int(iter_str)
                except ValueError:
                    pass

        # Create reverse lookup: name -> index
        name_to_idx = {v: k for k, v in names.items()}

        # Parse processor results
        for proc in root.iter('processor'):
            proc_name = proc.get('name', '')
            proc_idx = name_to_idx.get(proc_name, -1)

            for result_proc in proc.iter('result-processor'):
                util_str = result_proc.get('utilization', '')
                if util_str and proc_idx > 0:
                    try:
                        P[proc_idx - 1] = float(util_str)
                    except ValueError:
                        pass

            # Parse task results
            for task in proc.iter('task'):
                task_name = task.get('name', '')
                task_idx = name_to_idx.get(task_name, -1)

                for result_task in task.iter('result-task'):
                    if task_idx > 0:
                        # Utilization
                        util_str = result_task.get('utilization', '')
                        if util_str:
                            try:
                                U[task_idx - 1] = float(util_str)
                            except ValueError:
                                pass

                        # Throughput
                        tput_str = result_task.get('throughput', '')
                        if tput_str:
                            try:
                                T[task_idx - 1] = float(tput_str)
                            except ValueError:
                                pass

                        # Processor utilization
                        proc_util_str = result_task.get('proc-utilization', '')
                        if proc_util_str:
                            try:
                                P[task_idx - 1] = float(proc_util_str)
                            except ValueError:
                                pass

                # Parse entry results
                for entry in task.iter('entry'):
                    entry_name = entry.get('name', '')
                    entry_idx = name_to_idx.get(entry_name, -1)

                    for result_entry in entry.iter('result-entry'):
                        if entry_idx > 0:
                            # Utilization
                            util_str = result_entry.get('utilization', '')
                            if util_str:
                                try:
                                    U[entry_idx - 1] = float(util_str)
                                except ValueError:
                                    pass

                            # Phase 1 service time
                            st_str = result_entry.get('phase1-service-time', '')
                            if st_str:
                                try:
                                    S[entry_idx - 1] = float(st_str)
                                except ValueError:
                                    pass

                            # Throughput
                            tput_str = result_entry.get('throughput', '')
                            if tput_str:
                                try:
                                    T[entry_idx - 1] = float(tput_str)
                                except ValueError:
                                    pass

                            # Processor utilization
                            proc_util_str = result_entry.get('proc-utilization', '')
                            if proc_util_str:
                                try:
                                    P[entry_idx - 1] = float(proc_util_str)
                                except ValueError:
                                    pass

        # Parse activity results from task-activities
        for task_acts in root.iter('task-activities'):
            for activity in task_acts.iter('activity'):
                # Only process direct children of task-activities
                if activity.getparent() is None or activity.getparent().tag != 'task-activities':
                    # Check parent name if etree supports it
                    parent = None
                    for parent_elem in task_acts:
                        if parent_elem.tag == 'activity' and parent_elem.get('name') == activity.get('name'):
                            parent = parent_elem
                            break
                    if parent is None:
                        continue

                act_name = activity.get('name', '')
                act_idx = name_to_idx.get(act_name, -1)

                for result_act in activity.iter('result-activity'):
                    if act_idx > 0:
                        # Utilization
                        util_str = result_act.get('utilization', '')
                        if util_str:
                            try:
                                U[act_idx - 1] = float(util_str)
                            except ValueError:
                                pass

                        # Service time
                        st_str = result_act.get('service-time', '')
                        if st_str:
                            try:
                                S[act_idx - 1] = float(st_str)
                            except ValueError:
                                pass

                        # Throughput
                        tput_str = result_act.get('throughput', '')
                        if tput_str:
                            try:
                                T[act_idx - 1] = float(tput_str)
                            except ValueError:
                                pass

                        # Proc waiting
                        pw_str = result_act.get('proc-waiting', '')
                        if pw_str:
                            try:
                                W[act_idx - 1] = float(pw_str)
                            except ValueError:
                                pass

                        # Processor utilization
                        proc_util_str = result_act.get('proc-utilization', '')
                        if proc_util_str:
                            try:
                                P[act_idx - 1] = float(proc_util_str)
                            except ValueError:
                                pass

    except Exception as e:
        # Return partially parsed results on error
        pass

    return U, P, S, T, W, edges_W, iterations


def solver_lqns(
    lqnx_path: str,
    names: Dict[int, str],
    nidx: int,
    ncalls: int = 0,
    options: Optional[SolverLQNSOptions] = None
) -> SolverLQNSReturn:
    """
    LQNS solver handler - calls lqns/lqsim via subprocess.

    This function expects an LQNX file to already exist. For full model
    serialization, use the Java-based wrapper.

    Args:
        lqnx_path: Path to the LQNX input file
        names: Dictionary mapping indices to element names
        nidx: Number of node indices
        ncalls: Number of calls (edges)
        options: Solver options

    Returns:
        SolverLQNSReturn with all performance metrics

    Raises:
        RuntimeError: If LQNS is not available or fails
    """
    start_time = time.time()

    if options is None:
        options = SolverLQNSOptions()

    if not is_lqns_available():
        raise RuntimeError(
            "SolverLQNS requires the 'lqns' and 'lqsim' commands to be available in your system PATH.\n"
            "You can install them from: http://www.sce.carleton.ca/rads/lqns/\n\n"
            "Alternatively, use remote execution via Docker:\n"
            "  1. Pull and run: docker run -d -p 8080:8080 imperialqore/lqns-rest:latest\n"
            "  2. Configure remote execution:\n"
            "     options = SolverOptions()\n"
            "     options.config.remote = True\n"
            "     options.config.remote_url = 'http://localhost:8080'\n"
            "     solver = SolverLQNS(model, options)"
        )

    if not os.path.exists(lqnx_path):
        raise RuntimeError(f"LQNX file not found: {lqnx_path}")

    # Build command
    verbose_flag = '-a -w' if not options.verbose else ''
    praqma_flag = f'-Pmultiserver={options.multiserver}' if options.method.lower() != 'lqsim' else ''
    common = f'{verbose_flag} {praqma_flag} -Pstop-on-message-loss=false -x {lqnx_path}'

    method = options.method.lower()
    if method == 'srvn':
        cmd = f'lqns {common} -Playering=srvn'
    elif method == 'exactmva':
        cmd = f'lqns {common} -Pmva=exact'
    elif method == 'srvn.exactmva':
        cmd = f'lqns {common} -Playering=srvn -Pmva=exact'
    elif method in ('sim', 'lqsim'):
        cmd = f'lqsim {common} -A {options.samples}'
    elif method == 'lqnsdefault':
        cmd = f'lqns {verbose_flag} {praqma_flag} -x {lqnx_path}'
    else:
        cmd = f'lqns {common}'

    if options.verbose:
        print(f"LQNS command: {cmd}")

    # Execute command
    try:
        result = subprocess.run(
            cmd.split(),
            capture_output=True,
            timeout=300
        )

        if result.returncode != 0:
            stderr = result.stderr.decode('utf-8', errors='ignore')
            if stderr:
                raise RuntimeError(f"LQNS failed: {stderr}")
    except subprocess.TimeoutExpired:
        raise RuntimeError("LQNS timed out after 5 minutes")

    # Parse results
    result_path = lqnx_path.replace('.lqnx', '.lqxo')
    U, P, S, T, W, edges_W, iterations = _parse_lqxo_results(result_path, names, nidx, ncalls)

    # Clean up if not keeping files
    if not options.keep:
        try:
            if os.path.exists(result_path):
                os.remove(result_path)
        except OSError:
            pass

    runtime = time.time() - start_time

    return SolverLQNSReturn(
        Q=np.full(nidx, np.nan),  # LQNS doesn't provide queue lengths
        U=U.reshape(-1, 1),
        R=W.reshape(-1, 1),  # Response time from waiting
        T=T.reshape(-1, 1),
        P=P.reshape(-1, 1),
        S=S.reshape(-1, 1),
        A=np.full((nidx, 1), np.nan),
        W=W.reshape(-1, 1),
        runtime=runtime,
        method=options.method,
        iterations=iterations
    )


def solver_lqns_from_model(
    model,
    options: Optional[SolverLQNSOptions] = None
) -> SolverLQNSReturn:
    """
    LQNS solver handler that generates LQNX from a LayeredNetwork model.

    This is a convenience wrapper that handles LQNX generation for models
    that support it (requires the model to have writeXML capability).

    Args:
        model: LayeredNetwork model with writeXML method
        options: Solver options

    Returns:
        SolverLQNSReturn with all performance metrics
    """
    if options is None:
        options = SolverLQNSOptions()

    # Create temporary directory
    temp_dir = tempfile.mkdtemp(prefix='lqns_')

    try:
        lqnx_path = os.path.join(temp_dir, 'model.lqnx')

        # Try to get model structure for parsing
        try:
            lqn = model.getStruct()
            names = dict(lqn.names) if hasattr(lqn, 'names') else {}
            nidx = lqn.nidx if hasattr(lqn, 'nidx') else 0
            ncalls = lqn.ncalls if hasattr(lqn, 'ncalls') else 0
        except Exception:
            names = {}
            nidx = 0
            ncalls = 0

        # Generate LQNX file
        if hasattr(model, 'writeXML'):
            model.writeXML(lqnx_path, False)
        elif hasattr(model, 'obj') and hasattr(model.obj, 'writeXML'):
            model.obj.writeXML(lqnx_path, False)
        else:
            raise RuntimeError("Model does not support LQNX export")

        # Run solver
        result = solver_lqns(lqnx_path, names, nidx, ncalls, options)

        return result

    finally:
        if not options.keep:
            shutil.rmtree(temp_dir, ignore_errors=True)
