from __future__ import annotations

from collections.abc import Mapping, MutableMapping
import dataclasses
import enum
import json
import logging
import os
from typing import Any, Dict, Optional, cast

from ..common import Json, Uuid, format_percent, json_dict
from ..data.outcomes import FeasibleOutcome
from ..data.outlines import Label, ProblemOutline, outline_from_json
from ..data.solves import SolveInputs, SolveOptions, SolveStrategy
from ..data.tensors import DimensionArgument, Tensor, TensorArgument
from ..executors import Executor, JsonExecutorResult
from ..specifications import FormulationSpecification, Specification
from ..transformations import (
    ProblemTransformation,
    ProblemTransformationContext,
)


DEMO_ENDPOINT = "https://api.try.opvious.io"


class ClientSetting(enum.Enum):
    """Client configuration environment variables"""

    ENDPOINT = "OPVIOUS_ENDPOINT"
    TOKEN = "OPVIOUS_TOKEN"
    EXECUTOR = "OPVIOUS_EXECUTOR"

    def read(self, env: Optional[Mapping[str, str]] = None) -> str:
        """Read the setting's current value or default if missing

        Args:
            env: Environment, defaults to `os.environ`.
        """
        if env is None:
            env = cast(Any, os.environ)
        return env.get(self.value, "").strip()


def log_progress(logger: logging.Logger, progress: Json) -> None:
    kind = progress["kind"]
    if kind == "activity":
        iter_count = progress.get("lpIterationCount")
        gap = progress.get("relativeGap")
        if iter_count is not None:
            logger.info(
                "Solve in progress... [iterations=%s, gap=%s]",
                iter_count,
                "n/a" if gap is None else format_percent(gap),
            )
    elif kind == "epsilonConstraint":
        logger.info(
            "Added epsilon constraint. [objective_value=%s]",
            progress["objectiveValue"],
        )
    else:
        raise Exception(f"Unsupported progress kind: {kind}")


class SolveInputsBuilder:
    def __init__(self, outline: ProblemOutline):
        self._outline = outline
        self._dimensions: Dict[Label, Any] = {}
        self._parameters: Dict[Label, Any] = {}
        self.parameter_entry_count = 0

    def set_dimension(self, label: Label, arg: DimensionArgument) -> None:
        outline = self._outline.dimensions.get(label)
        if not outline:
            raise Exception(f"Unknown dimension: {label}")
        if label in self._dimensions:
            raise Exception(f"Duplicate dimension: {label}")
        items = list(arg)
        self._dimensions[label] = {"label": label, "items": items}

    def set_parameter(self, label: Label, arg: Any) -> None:
        outline = self._outline.parameters.get(label)
        if not outline:
            raise Exception(f"Unknown parameter: {label}")
        if label in self._parameters:
            raise Exception(f"Duplicate parameter: {label}")
        try:
            tensor = Tensor.from_argument(
                arg,
                rank=len(outline.bindings),
                is_indicator=outline.is_indicator,
                is_pin=outline.derivation_kind == "pinnedVariable",
            )
        except Exception as exc:
            raise ValueError(f"Invalid  parameter: {label}") from exc
        self._parameters[label] = json_dict(
            label=label,
            entries=tensor.entries,
            default_value=tensor.default_value,
        )
        self.parameter_entry_count += len(tensor.entries)

    def build(self) -> SolveInputs:
        missing_labels = set()

        for label in self._outline.parameters:
            if label not in self._parameters:
                missing_labels.add(label)

        if self._dimensions:
            for label in self._outline.dimensions:
                if label not in self._dimensions:
                    missing_labels.add(label)

        if missing_labels:
            raise Exception(f"Missing label(s): {missing_labels}")

        return SolveInputs(
            problem_outline=self._outline,
            raw_parameters=list(self._parameters.values()),
            raw_dimensions=list(self._dimensions.values()) or None,
        )


async def _generate_outline(
    executor: Executor, outline_data: Json, transformation_data: Json
) -> ProblemOutline:
    if not transformation_data:
        return outline_from_json(outline_data)
    async with executor.execute(
        result_type=JsonExecutorResult,
        url="/outlines/transform",
        method="POST",
        json_data=json_dict(
            outline=outline_data,
            transformations=transformation_data,
        ),
    ) as res:
        data = res.json_data()
    return outline_from_json(data["outline"])


class ProblemOutlineCache:
    """Efficiently returns outlines for past queued solves"""

    def __init__(self, executor: Executor) -> None:
        self._executor = executor

        try:
            import lru
        except ImportError:
            self._by_solve: MutableMapping[Uuid, ProblemOutline] = {}
        else:
            self._by_solve = cast(Any, lru.LRU(100))

    async def get_solve_outline(self, uuid: Uuid) -> ProblemOutline:
        cached = self._by_solve.get(uuid)
        if cached:
            return cached

        data = await self._executor.execute_graphql_query(
            query="@FetchQueuedSolve",
            variables=json_dict(uuid=uuid),
        )
        solve_data = data["queuedSolve"]
        if not solve_data:
            raise ValueError(f"Unknown solve: {uuid}")
        outline = await _generate_outline(
            self._executor,
            solve_data["specification"]["outline"],
            solve_data["transformations"],
        )
        self._by_solve[uuid] = outline
        return outline


class ProblemOutlineGenerator:
    """Generates outlines from a formulation and transformations"""

    def __init__(self, executor: Executor, outline_data: Json):
        self._executor = executor
        self._pristine_outline_data = outline_data
        self._transformations = cast(list[ProblemTransformation], [])

    @classmethod
    async def formulation(
        cls, executor: Executor, specification: FormulationSpecification
    ) -> tuple[ProblemOutlineGenerator, str]:
        data = await executor.execute_graphql_query(
            query="@FetchOutline",
            variables={
                "formulationName": specification.formulation_name,
                "tagName": specification.tag_name,
            },
        )
        formulation = data.get("formulation")
        if not formulation:
            raise Exception("No matching formulation found")
        tag = formulation.get("tag")
        if not tag:
            raise Exception("No matching specification found")
        spec = tag["specification"]
        return (
            ProblemOutlineGenerator(executor, spec["outline"]),
            tag["name"],
        )

    @classmethod
    async def sources(
        cls, executor: Executor, sources: list[str]
    ) -> ProblemOutlineGenerator:
        async with executor.execute(
            result_type=JsonExecutorResult,
            url="/sources/parse",
            method="POST",
            json_data=json_dict(sources=sources, outline=True),
        ) as res:
            outline_data = res.json_data()
        errors = outline_data.get("errors")
        if errors:
            raise Exception(f"Invalid sources: {json.dumps(errors)}")
        return ProblemOutlineGenerator(executor, outline_data["outline"])

    def add_transformation(self, tf: ProblemTransformation) -> None:
        self._transformations.append(tf)

    async def generate(self) -> tuple[ProblemOutline, Json]:
        pristine_outline = outline_from_json(self._pristine_outline_data)
        if not self._transformations:
            return pristine_outline, []
        executor = self._executor
        pristine_outline_data = self._pristine_outline_data

        class Context(ProblemTransformationContext):
            async def fetch_outline(self) -> ProblemOutline:
                transformation_data = self.get_json()
                if not transformation_data:
                    return pristine_outline
                return await _generate_outline(
                    executor, pristine_outline_data, transformation_data
                )

        context = Context()
        for tf in self._transformations:
            await tf.register(context)
        outline = await context.fetch_outline()
        return outline, context.get_json()


def feasible_outcome_details(outcome: FeasibleOutcome) -> Optional[str]:
    details = []
    if outcome.objective_value:
        details.append(f"objective={outcome.objective_value}")
    if outcome.relative_gap:
        details.append(f"gap={format_percent(outcome.relative_gap)}")
    return ", ".join(details) if details else None


@dataclasses.dataclass(frozen=True)
class Problem:
    """An optimization problem instance"""

    specification: Specification
    """Model :ref:`specification <specifications>`"""

    parameters: Optional[Mapping[Label, TensorArgument]] = None
    """Input data, keyed by parameter label

    Values may be any value accepted by :meth:`.Tensor.from_argument` and must
    match the corresponding parameter's definition.
    """

    dimensions: Optional[Mapping[Label, DimensionArgument]] = None
    """Dimension items, keyed by dimension label

    If omitted, these will be automatically inferred from the parameters.
    """

    transformations: Optional[list[ProblemTransformation]] = None
    """Optional :ref:`transformations <transformations>`"""

    strategy: Optional[SolveStrategy] = None
    """Optional :ref:`multi-objective strategy <multi-objective strategies>`

    This argument is required if the problem has two or more objectives.
    """

    options: Optional[SolveOptions] = None
    """Optional solve options (gap thresholds, timeout, etc.)"""
