"""

automl_lib.exporter

~~~~~~~~~~~~~~~~~~~

Export trained models to ONNX / TorchScript and scaffold a FastAPI

inference server.
"""


from __future__ import annotations

import os

from typing import Any

import torch
import torch.nn as nn


from .config import trainiqConfig

from .utils import ensure_dir, get_logger


logger = get_logger(__name__)



def export_model(

    model: nn.Module,

    input_shape: tuple,

    config: trainiqConfig,

) -> dict[str, str]:

    """Export the model to TorchScript and/or ONNX.


    Returns a dict of ``{format: filepath}``.
    """

    out_dir = ensure_dir(config.export_path or os.path.join(config.output_dir, "exported"))

    model.eval()
    model.cpu()

    results: dict[str, str] = {}


    fmt = config.export_format.lower()


    # ── TorchScript ───────────────────────────────────────────────────

    if fmt in ("torchscript", "both"):

        path = os.path.join(out_dir, "model_scripted.pt")

        try:

            dummy = torch.randn(1, *input_shape)

            traced = torch.jit.trace(model, dummy)

            traced.save(path)

            results["torchscript"] = path

            logger.info(f"TorchScript saved -> {path}")

        except Exception as e:

            logger.warning(f"TorchScript export failed: {e}")

            # Fallback: try scripting

            try:

                scripted = torch.jit.script(model)

                scripted.save(path)

                results["torchscript"] = path

                logger.info(f"TorchScript (scripted) saved -> {path}")

            except Exception as e2:

                logger.error(f"TorchScript scripting also failed: {e2}")


    # ── ONNX ──────────────────────────────────────────────────────────

    if fmt in ("onnx", "both"):

        path = os.path.join(out_dir, "model.onnx")

        try:

            dummy = torch.randn(1, *input_shape)

            torch.onnx.export(

                model, dummy, path,

                input_names=["input"],

                output_names=["output"],

                dynamic_axes={

                    "input": {0: "batch_size"},

                    "output": {0: "batch_size"},

                },

                opset_version=14,
            )

            results["onnx"] = path

            logger.info(f"ONNX saved -> {path}")

        except Exception as e:

            logger.error(f"ONNX export failed: {e}")

    return results



def scaffold_api(

    model_path: str,

    output_dir: str = "trainiq_output/api",

) -> str:

    """Generate a minimal FastAPI app that serves predictions.


    Creates:

    - ``app.py`` – FastAPI application

    - ``Dockerfile`` – ready-to-build container

    - ``requirements.txt`` – API dependencies


    Returns the directory path.
    """

    ensure_dir(output_dir)


    # ── app.py ────────────────────────────────────────────────────────

    # Normalize path for cross-platform compatibility

    normalized_path = model_path.replace("\\", "/")
    

    app_code = f'''\

"""Auto-generated FastAPI inference server."""
import torch

import numpy as np

from fastapi import FastAPI

from pydantic import BaseModel


app = FastAPI(title="trainiq Inference API", version="1.0.0")


# Load model on startup

MODEL_PATH = "{normalized_path}"

model = torch.jit.load(MODEL_PATH, map_location="cpu")

model.eval()



class PredictRequest(BaseModel):

    data: list[list[float]]



class PredictResponse(BaseModel):

    predictions: list



@app.post("/predict", response_model=PredictResponse)

def predict(req: PredictRequest):

    tensor = torch.tensor(req.data, dtype=torch.float32)

    with torch.no_grad():

        output = model(tensor)

    if output.dim() > 1 and output.size(1) > 1:

        preds = output.argmax(dim=1).tolist()

    else:

        preds = output.squeeze().tolist()

    if not isinstance(preds, list):

        preds = [preds]

    return PredictResponse(predictions=preds)



@app.get("/health")

def health():

    return {{"status": "ok"}}

'''


    # ── Dockerfile ────────────────────────────────────────────────────

    dockerfile = """\

FROM python:3.11-slim

WORKDIR /app

COPY requirements.txt .

RUN pip install --no-cache-dir -r requirements.txt

COPY . .

EXPOSE 8000

CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "8000"]
"""


    # ── requirements.txt ──────────────────────────────────────────────

    reqs = "fastapi>=0.100\\nuvicorn[standard]\\ntorch\\nnumpy\\n"


    for name, content in [("app.py", app_code),

                          ("Dockerfile", dockerfile),

                          ("requirements.txt", reqs)]:

        path = os.path.join(output_dir, name)

        with open(path, "w") as f:

            f.write(content)


    logger.info(f"FastAPI scaffold created -> {output_dir}")

    logger.info(f"   Run: cd {output_dir} && uvicorn app:app --reload")

    return output_dir

