"""Optimiser endpoints: solve, status, apply, save, frontier, mlflow log."""

from __future__ import annotations

import json
import time
from pathlib import Path
from typing import Any

from fastapi import APIRouter, HTTPException

from haute._logging import get_logger
from haute._sandbox import _get_project_root
from haute._types import SolveResultLike
from haute.routes._helpers import _INTERNAL_ERROR_DETAIL, validate_safe_path
from haute.routes._job_store import JobStore
from haute.routes._optimiser_service import (
    _DEFAULT_CHUNK_SIZE,
    _DEFAULT_TIMEOUT,
    OptimiserSolveService,
    _compute_scenario_value_stats,
)
from haute.schemas import (
    OptimiserApplyRequest,
    OptimiserApplyResponse,
    OptimiserFrontierRequest,
    OptimiserFrontierResponse,
    OptimiserFrontierSelectRequest,
    OptimiserFrontierSelectResponse,
    OptimiserMlflowLogRequest,
    OptimiserMlflowLogResponse,
    OptimiserSaveRequest,
    OptimiserSaveResponse,
    OptimiserSolveRequest,
    OptimiserSolveResponse,
    OptimiserStatusResponse,
)

logger = get_logger(component="server.optimiser")

router = APIRouter(prefix="/api/optimiser", tags=["optimiser"])

# In-memory job store — same pattern as modelling.
_store = JobStore()
_solve_service = OptimiserSolveService(_store)

_APPLY_PREVIEW_ROWS = 100  # max rows returned in the apply preview payload


@router.post("/solve", response_model=OptimiserSolveResponse)
def solve(body: OptimiserSolveRequest) -> OptimiserSolveResponse:
    """Start optimisation for an optimiser node.

    Executes the pipeline up to the optimiser node to materialise the
    scored DataFrame, then runs the solver in a background thread.
    """
    return _solve_service.start(body)


@router.get("/solve/status/{job_id}", response_model=OptimiserStatusResponse)
async def solve_status(job_id: str) -> OptimiserStatusResponse:
    """Poll optimisation job progress."""
    job = _store.require_job(job_id)

    # Check for timeout on running jobs
    if job.get("status") == "running":
        start = job.get("start_time")
        timeout = job.get("timeout", _DEFAULT_TIMEOUT)
        if start and (time.monotonic() - start) > timeout:
            # P7: Atomic update — only if still running (avoids overwriting a
            # completed result with a timeout error).
            _store.atomic_update(
                job_id,
                {
                    "status": "error",
                    "message": f"Solve timed out after {timeout}s. "
                    "Increase timeout or simplify the problem.",
                    "elapsed_seconds": time.monotonic() - start,
                },
                expected_status="running",
            )
            job = _store.require_job(job_id)

    frontier_resp = None
    if job.get("status") == "completed":
        fd = job.get("frontier_data")
        if fd:
            frontier_resp = OptimiserFrontierResponse(**fd)

    return OptimiserStatusResponse(
        status=job.get("status", "unknown"),
        progress=job.get("progress", 0.0),
        message=job.get("message", ""),
        elapsed_seconds=job.get("elapsed_seconds", 0.0),
        result=job.get("result"),
        frontier=frontier_resp,
    )


@router.post("/apply", response_model=OptimiserApplyResponse)
def apply_lambdas(body: OptimiserApplyRequest) -> OptimiserApplyResponse:
    """Apply solved lambdas to the scored data."""
    logger.info("apply_requested", job_id=body.job_id)
    job = _store.require_completed_job(body.job_id)

    solve_result = job.get("solve_result")
    if solve_result is None:
        raise HTTPException(status_code=400, detail="Job has no solve result")

    try:
        df = solve_result.dataframe
        return OptimiserApplyResponse(
            status="ok",
            total_objective=solve_result.total_objective,
            constraints=solve_result.total_constraints,
            preview=df.head(_APPLY_PREVIEW_ROWS).to_dicts(),
            row_count=len(df),
        )
    except Exception as exc:
        logger.error("apply_failed", error=str(exc), job_id=body.job_id)
        raise HTTPException(status_code=500, detail=_INTERNAL_ERROR_DETAIL)


@router.post("/frontier", response_model=OptimiserFrontierResponse)
def run_frontier(body: OptimiserFrontierRequest) -> OptimiserFrontierResponse:
    """Compute efficient frontier for a completed optimisation job."""
    job = _store.require_completed_job(body.job_id)

    solver = job.get("solver")
    quote_grid = job.get("quote_grid")
    if solver is None or quote_grid is None:
        raise HTTPException(
            status_code=400,
            detail=(
                "Solver and quote grid have been released to free memory. "
                "Use the pre-computed frontier from the solve result, or re-run the solve."
            ),
        )

    try:
        # Convert threshold ranges from lists to tuples for Rust binding
        ranges = {k: tuple(v) for k, v in body.threshold_ranges.items()}
        frontier_result = solver.frontier(
            quote_grid,
            threshold_ranges=ranges,
            n_points_per_dim=body.n_points_per_dim,
        )
        points_df = frontier_result.points
        return OptimiserFrontierResponse(
            status="ok",
            points=points_df.to_dicts(),
            n_points=len(points_df),
            constraint_names=list(body.threshold_ranges.keys()),
        )
    except Exception as exc:
        logger.error("frontier_failed", error=str(exc), job_id=body.job_id)
        raise HTTPException(status_code=500, detail=_INTERNAL_ERROR_DETAIL)


@router.post("/frontier/select", response_model=OptimiserFrontierSelectResponse)
def select_frontier_point(body: OptimiserFrontierSelectRequest) -> OptimiserFrontierSelectResponse:
    """Re-solve at a specific frontier point's lambdas and swap the job's active result."""
    job = _store.require_completed_job(body.job_id)

    # M10: Short-circuit if this point is already selected (idempotent)
    if job.get("selected_frontier_point") == body.point_index:
        result = job.get("result", {})
        return OptimiserFrontierSelectResponse(
            status="ok",
            total_objective=result.get("total_objective", 0.0),
            constraints=result.get("constraints", {}),
            baseline_objective=result.get("baseline_objective", 0.0),
            baseline_constraints=result.get("baseline_constraints", {}),
            lambdas=result.get("lambdas", {}),
            converged=result.get("converged", True),
        )

    solver = job.get("solver")
    quote_grid = job.get("quote_grid")
    frontier_data = job.get("frontier_data")

    if solver is None or quote_grid is None:
        raise HTTPException(
            status_code=400,
            detail=(
                "Solver and quote grid have been released to free memory. "
                "Use the pre-computed frontier from the solve result, or re-run the solve."
            ),
        )
    if not frontier_data or not frontier_data.get("points"):
        raise HTTPException(status_code=400, detail="Job has no frontier data")

    points = frontier_data["points"]
    if body.point_index < 0 or body.point_index >= len(points):
        raise HTTPException(
            status_code=400,
            detail=f"Point index {body.point_index} out of range [0, {len(points)})",
        )

    point = points[body.point_index]
    # Extract lambdas from frontier point (columns are lambda_{constraint_name})
    new_lambdas = {
        k.removeprefix("lambda_"): v
        for k, v in point.items()
        if k.startswith("lambda_") and isinstance(v, (int, float))
    }

    if not new_lambdas:
        raise HTTPException(status_code=400, detail="Frontier point has no lambda values")

    try:
        # Re-solve with the selected point's lambdas (warm-start = fast)
        new_result = solver.solve(quote_grid, lambdas=new_lambdas)

        # Swap the solve_result; preserve original for revert
        result_dict = dict(job.get("result", {}))
        result_dict.update(
            {
                "total_objective": new_result.total_objective,
                "baseline_objective": new_result.baseline_objective,
                "constraints": new_result.total_constraints,
                "baseline_constraints": new_result.baseline_constraints,
                "lambdas": new_result.lambdas,
                "converged": new_result.converged,
                "selected_frontier_point": body.point_index,
            }
        )

        # M1: Recompute scenario value stats for the new solve result
        scenario_stats, scenario_histogram = _compute_scenario_value_stats(new_result)
        result_dict["scenario_value_stats"] = scenario_stats
        result_dict["scenario_value_histogram"] = scenario_histogram

        # M2: Update convergence warning for the new solve result
        if not new_result.converged:
            result_dict["warning"] = (
                "Solver did not converge. Consider increasing max_iter or relaxing tolerance."
            )
        else:
            result_dict.pop("warning", None)

        _store.atomic_update(
            body.job_id,
            {
                "solve_result": new_result,
                "selected_frontier_point": body.point_index,
                "result": result_dict,
            },
        )

        return OptimiserFrontierSelectResponse(
            status="ok",
            total_objective=new_result.total_objective,
            constraints=new_result.total_constraints,
            baseline_objective=new_result.baseline_objective,
            baseline_constraints=new_result.baseline_constraints,
            lambdas=new_result.lambdas,
            converged=new_result.converged,
        )
    except Exception as exc:
        logger.error("frontier_select_failed", error=str(exc), job_id=body.job_id)
        raise HTTPException(status_code=500, detail=_INTERNAL_ERROR_DETAIL)


def _build_artifact_payload(
    job: dict[str, Any],
    solve_result: SolveResultLike,
    version_override: str = "",
) -> dict[str, Any]:
    """Build the JSON payload for an optimiser artifact.

    Shared by both file-save and MLflow-log paths to avoid duplication.
    """
    from datetime import datetime, timezone

    node_label = job.get("node_label", "optimiser")
    label_slug = node_label.lower().replace(" ", "_")
    ts = datetime.now(timezone.utc).strftime("%Y%m%d_%H%M%S")  # noqa: UP017
    auto_version = f"{label_slug}_{ts}"
    job_config = job.get("config", {})

    payload: dict[str, Any] = {
        "version": version_override or auto_version,
        "created_at": datetime.now(timezone.utc).isoformat(),  # noqa: UP017
        "mode": job_config.get("mode", "online"),
        "lambdas": solve_result.lambdas,
        "total_objective": solve_result.total_objective,
        "baseline_objective": getattr(solve_result, "baseline_objective", None),
        "total_constraints": solve_result.total_constraints,
        "baseline_constraints": getattr(solve_result, "baseline_constraints", None),
        "constraints": job_config.get("constraints"),
        "objective": job_config.get("objective"),
        "quote_id": job_config.get("quote_id", "quote_id"),
        "scenario_index": job_config.get("scenario_index", "scenario_index"),
        "scenario_value": job_config.get("scenario_value", "scenario_value"),
        "chunk_size": job_config.get("chunk_size", _DEFAULT_CHUNK_SIZE),
        "converged": solve_result.converged,
        "iterations": getattr(solve_result, "iterations", None),
        "cd_iterations": getattr(solve_result, "cd_iterations", None),
    }
    # Frontier provenance — record which point was selected (if any)
    selected_idx = job.get("selected_frontier_point")
    frontier_data = job.get("frontier_data")
    if selected_idx is not None and frontier_data:
        payload["frontier_selection"] = {
            "selected_from_frontier": True,
            "point_index": selected_idx,
            "n_frontier_points": frontier_data.get("n_points", 0),
        }
    if job_config.get("mode") == "ratebook":
        payload["factor_tables"] = job.get("result", {}).get("factor_tables")
        payload["clamp_rate"] = getattr(solve_result, "clamp_rate", None)
    return payload


@router.post("/save", response_model=OptimiserSaveResponse)
def save_result(body: OptimiserSaveRequest) -> OptimiserSaveResponse:
    """Save the optimisation result to disk."""
    job = _store.require_completed_job(body.job_id)

    solve_result = job.get("solve_result")
    if solve_result is None:
        raise HTTPException(status_code=400, detail="Job has no solve result")

    from haute.routes._helpers import pipeline_dir

    # Absolute paths resolve against the project root (security boundary).
    # Relative paths resolve against the pipeline directory (e.g. rating/)
    # so outputs land next to the pipeline file, not at the project root.
    user_path = Path(body.output_path)
    if user_path.is_absolute():
        base = _get_project_root()
    else:
        base = pipeline_dir()
    out = validate_safe_path(base, body.output_path)

    try:
        out.parent.mkdir(parents=True, exist_ok=True)

        payload = _build_artifact_payload(job, solve_result, version_override=body.version)
        out.write_text(json.dumps(payload, indent=2, default=str))
        logger.info("result_saved", path=str(out), job_id=body.job_id)

        return OptimiserSaveResponse(
            status="ok",
            path=str(out),
            message=f"Saved optimisation result to {out}",
        )
    except HTTPException:
        raise
    except OSError as exc:
        logger.error("save_failed", error=str(exc), job_id=body.job_id)
        raise HTTPException(
            status_code=500,
            detail="Filesystem error saving optimiser result. Check the server logs for details.",
        )
    except Exception as exc:
        logger.error("save_failed", error=str(exc), job_id=body.job_id)
        raise HTTPException(status_code=500, detail=_INTERNAL_ERROR_DETAIL)


@router.post("/mlflow/log", response_model=OptimiserMlflowLogResponse)
def mlflow_log(body: OptimiserMlflowLogRequest) -> OptimiserMlflowLogResponse:
    """Log optimisation results to MLflow."""
    job = _store.require_completed_job(body.job_id)

    solve_result = job.get("solve_result")
    solver = job.get("solver")
    if solve_result is None:
        raise HTTPException(status_code=400, detail="Job has no solve result")
    if solver is None:
        raise HTTPException(
            status_code=400,
            detail=(
                "Solver has been released to free memory. "
                "MLflow logging must happen before the solver is garbage-collected. "
                "Re-run the solve to log results."
            ),
        )

    try:
        import mlflow
    except ImportError:
        raise HTTPException(
            status_code=400,
            detail="MLflow is not installed. Install with: pip install mlflow",
        )

    try:
        from haute.modelling._mlflow_log import (
            build_run_url,
            configure_mlflow_tracking,
            resolve_experiment_name,
        )

        tracking_uri, backend = configure_mlflow_tracking()

        summary = solver.summary(solve_result)

        node_label = job.get("node_label", "optimiser")
        job_config = job.get("config", {})

        experiment_name = resolve_experiment_name(
            explicit=body.experiment_name,
            config_value=job_config.get("mlflow_experiment"),
            node_label=node_label,
            backend=backend,
        )
        mlflow.set_experiment(experiment_name)

        with mlflow.start_run(run_name=node_label) as run:
            mlflow.log_params(summary["params"])
            mlflow.log_metrics(summary["metrics"])

            # Log artifacts as JSON files
            import tempfile

            with tempfile.TemporaryDirectory() as tmpdir:
                artifacts = summary.get("artifacts", {})
                for name, data in artifacts.items():
                    if data is None:
                        continue
                    artifact_path = Path(tmpdir) / f"{name}.json"
                    artifact_path.write_text(json.dumps(data, indent=2, default=str))
                    mlflow.log_artifact(str(artifact_path))

                # Also log the complete artifact used by OPTIMISER_APPLY
                complete_payload = _build_artifact_payload(job, solve_result)
                complete_path = Path(tmpdir) / "optimiser_result.json"
                complete_path.write_text(json.dumps(complete_payload, indent=2, default=str))
                mlflow.log_artifact(str(complete_path))

                # Log frontier CSV artifact + provenance tags
                frontier_data = job.get("frontier_data")
                if frontier_data and frontier_data.get("points"):
                    import csv
                    import io

                    points = frontier_data["points"]
                    buf = io.StringIO()
                    writer = csv.DictWriter(buf, fieldnames=points[0].keys())
                    writer.writeheader()
                    writer.writerows(points)
                    frontier_path = Path(tmpdir) / "frontier.csv"
                    frontier_path.write_text(buf.getvalue())
                    mlflow.log_artifact(str(frontier_path))
                    mlflow.set_tag("frontier.n_points", str(frontier_data["n_points"]))

                selected_idx = job.get("selected_frontier_point")
                if selected_idx is not None:
                    mlflow.set_tag("frontier.selected_point_index", str(selected_idx))

            run_id = run.info.run_id
            run_url = build_run_url(backend, experiment_name, run_id)

        return OptimiserMlflowLogResponse(
            status="ok",
            backend=backend,
            experiment_name=experiment_name,
            run_id=run_id,
            run_url=run_url,
            tracking_uri=tracking_uri,
        )
    except Exception as exc:
        logger.error("mlflow_log_failed", error=str(exc), job_id=body.job_id)
        raise HTTPException(status_code=500, detail=_INTERNAL_ERROR_DETAIL)
