"""
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 AutoMLConfig
from .utils import ensure_dir, get_logger

logger = get_logger(__name__)


def export_model(
    model: nn.Module,
    input_shape: tuple,
    config: AutoMLConfig,
) -> 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 = "automl_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="AutoML 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
