"""Admiral Solver SDK — Python client for the Admiral Optimization Platform.

Supports all 11 problem types, solver racing, AI advisor, NL builder, and auto-detection.
"""
from __future__ import annotations
import os, time
import httpx
import numpy as np
from admiral.models.result import SolveResult, AnalysisResult, FormulationResult, DetectionResult


class Solver:
    """Admiral Solver client.

    Args:
        api_key: Your Admiral API key (adm_sk_...). Falls back to ADMIRAL_API_KEY env var.
        base_url: API base URL. Defaults to https://api.admiral-platform.tech
        timeout: HTTP request timeout in seconds.

    Example:
        >>> from admiral import Solver
        >>> solver = Solver(api_key="adm_sk_...")
        >>> result = solver.solve_qubo([[-5, 2], [0, -3]])
        >>> print(result.energy)
    """

    def __init__(self, api_key=None, base_url=None, timeout=120):
        self.api_key = api_key or os.getenv("ADMIRAL_API_KEY")
        if not self.api_key:
            raise ValueError("API key required. Pass api_key= or set ADMIRAL_API_KEY env var.")
        self.base_url = (base_url or os.getenv("ADMIRAL_API_URL", "https://api.admiral-platform.tech")).rstrip("/")
        self.timeout = timeout
        self._client = httpx.Client(timeout=timeout)

    @property
    def _headers(self):
        return {"Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json"}

    # ══════════════════════════════════════════════
    # SOLVE METHODS — 11 problem types
    # ══════════════════════════════════════════════

    def solve_qubo(self, Q, timeout=30, num_reads=100, race=False):
        """Solve a QUBO problem: min x^T Q x, x in {0,1}^n."""
        return self._submit("qubo", {"Q": _np(Q)}, timeout, num_reads, race)

    def solve_ising(self, h, J, timeout=30, num_reads=100, race=False):
        """Solve an Ising model: min sum(h_i * s_i) + sum(J_ij * s_i * s_j), s in {-1,+1}^n."""
        return self._submit("ising", {"h": _np(h), "J": _np(J)}, timeout, num_reads, race)

    def solve_hubo(self, terms, num_variables, timeout=30, num_reads=100, race=False):
        """Solve a higher-order binary optimization problem."""
        data = {"terms": [{"indices": list(idx), "coefficient": c} for idx, c in terms], "num_variables": num_variables}
        return self._submit("hubo", data, timeout, num_reads, race)

    def solve_dqm(self, domains, linear=None, quadratic=None, timeout=30, num_reads=100, race=False):
        """Solve a discrete quadratic model."""
        data = {"domains": domains}
        if linear: data["linear"] = linear
        if quadratic:
            data["quadratic"] = {}
            for (v1, v2), mat in quadratic.items():
                data["quadratic"][f"{v1},{v2}"] = _np(mat)
        return self._submit("dqm", data, timeout, num_reads, race)

    def solve_cqm(self, c, Q, A, b, var_types=None, bounds=None, timeout=30, num_reads=100, race=False):
        """Solve a constrained quadratic model."""
        data = {"c": _np(c), "Q": _np(Q), "A": _np(A), "b": _np(b)}
        if var_types: data["var_types"] = var_types
        if bounds: data["bounds"] = bounds
        return self._submit("cqm", data, timeout, num_reads, race)

    def solve_milp(self, c, A, b, integer_vars=None, bounds=None, timeout=30, race=False):
        """Solve a mixed-integer linear program."""
        data = {"c": _np(c), "A": _np(A), "b": _np(b)}
        if integer_vars is not None: data["integer_vars"] = list(integer_vars)
        if bounds: data["bounds"] = bounds
        return self._submit("milp", data, timeout, 1, race)

    def solve_miqp(self, c, Q, A, b, var_types=None, bounds=None, timeout=30, num_reads=100, race=False):
        """Solve a mixed-integer quadratic program."""
        data = {"c": _np(c), "Q": _np(Q), "A": _np(A), "b": _np(b)}
        if var_types: data["var_types"] = var_types
        if bounds: data["bounds"] = bounds
        return self._submit("miqp", data, timeout, num_reads, race)

    def solve_maxsat(self, clauses, num_variables=None, timeout=30, num_reads=100, race=False):
        """Solve a weighted MAX-SAT problem."""
        cl = [{"literals": c[0], "weight": c[1] if len(c) > 1 else 1, "hard": c[2] if len(c) > 2 else False} for c in clauses]
        data = {"clauses": cl}
        if num_variables: data["num_variables"] = num_variables
        return self._submit("maxsat", data, timeout, num_reads, race)

    def solve_pbo(self, objective, constraints, num_variables, timeout=30, num_reads=100, race=False):
        """Solve a pseudo-boolean optimization problem."""
        obj = [{"coefficient": c, "literals": lits} for c, lits in objective]
        cons = [{"terms": [{"coefficient": c, "literals": l} for c, l in con[0]], "rhs": con[1]} for con in constraints]
        return self._submit("pbo", {"objective": obj, "constraints": cons, "num_variables": num_variables}, timeout, num_reads, race)

    def solve_wcsp(self, num_variables, domains, functions, timeout=30, num_reads=100, race=False):
        """Solve a weighted constraint satisfaction problem."""
        funcs = [{"scope": list(s), "entries": [{"values": list(k), "cost": v} for k, v in c.items()]} for s, c in functions]
        return self._submit("wcsp", {"num_variables": num_variables, "domains": domains, "functions": funcs}, timeout, num_reads, race)

    def solve_potts(self, num_nodes, num_states, edges, fields=None, timeout=30, num_reads=100, race=False):
        """Solve a Potts model optimization problem."""
        e = [{"i": i, "j": j, "coupling": c} for (i, j), c in edges.items()]
        f = [{"node": n, "costs": c} for n, c in (fields or {}).items()]
        return self._submit("potts", {"num_nodes": num_nodes, "num_states": num_states, "edges": e, "fields": f}, timeout, num_reads, race)

    # ══════════════════════════════════════════════
    # AI FEATURES — v2.0
    # ══════════════════════════════════════════════

    def analyze(self, problem_type, data):
        """AI Advisor: analyze problem structure and get solver recommendations.

        Args:
            problem_type: One of qubo, ising, hubo, dqm, cqm, milp, miqp, maxsat, pbo, wcsp, potts
            data: Problem data dict (e.g. {"Q": [[...]]})

        Returns:
            AnalysisResult with structure, complexity, recommendations, suggested_config
        """
        r = self._client.post(f"{self.base_url}/v1/analyze", headers=self._headers,
                              json={"problem_type": problem_type, "data": data})
        r.raise_for_status()
        return AnalysisResult(r.json())

    def formulate(self, description, target_type="auto"):
        """NL Builder: convert natural language to optimization formulation.

        Args:
            description: Plain English description of the optimization problem
            target_type: Target problem type (auto, qubo, milp, etc.)

        Returns:
            FormulationResult with problem_type, data, explanation
        """
        r = self._client.post(f"{self.base_url}/v1/formulate", headers=self._headers,
                              json={"description": description, "target_type": target_type})
        r.raise_for_status()
        return FormulationResult(r.json())

    def detect(self, data):
        """Auto-detect: identify problem type from data structure.

        Args:
            data: Problem data dict (e.g. {"Q": [[...]]})

        Returns:
            DetectionResult with detected_type, confidence, reason
        """
        r = self._client.post(f"{self.base_url}/v1/detect", headers=self._headers,
                              json={"data": data})
        r.raise_for_status()
        return DetectionResult(r.json())

    def health(self):
        """Check platform health and version."""
        r = self._client.get(f"{self.base_url}/v1/health")
        r.raise_for_status()
        return r.json()

    # ══════════════════════════════════════════════
    # INTERNAL
    # ══════════════════════════════════════════════

    def _submit(self, problem_type, data, timeout, num_reads, race):
        body = {
            "problem_type": problem_type,
            "data": data,
            "config": {"timeout": timeout, "num_reads": num_reads, "race": race},
        }
        r = self._client.post(f"{self.base_url}/v1/solve", headers=self._headers, json=body)
        r.raise_for_status()
        jid = r.json()["job_id"]

        deadline = time.time() + timeout + 60
        while time.time() < deadline:
            time.sleep(1)
            r = self._client.get(f"{self.base_url}/v1/jobs/{jid}", headers=self._headers)
            r.raise_for_status()
            j = r.json()
            if j["status"] == "completed":
                return SolveResult(j.get("result", j))
            if j["status"] == "failed":
                raise RuntimeError(f"Job failed: {j.get('error')}")
        raise TimeoutError(f"Job {jid} did not complete within {timeout + 60}s")

    def __del__(self):
        try:
            self._client.close()
        except Exception:
            pass


def _np(x):
    """Convert numpy arrays to lists for JSON serialization."""
    if isinstance(x, np.ndarray):
        return x.tolist()
    return x
