"""
DeployIntegration: generates a FastAPI server.py with FRAMEWORM-COST
baked in automatically.

Used by FRAMEWORM-DEPLOY when generating the deployment server.
"""

from __future__ import annotations

from typing import Optional

# ── Architecture → parameter count lookup ────────────────────────────────────
ARCH_PARAMS = {
    "vae": 4.5,
    "dcgan": 6.8,
    "ddpm": 35.0,
    "cfg_ddpm": 40.0,
    "vqvae2": 12.0,
    "vit_gan": 86.0,
    "unknown": 10.0,
}

ARCH_MEMORY = {
    "vae": 0.5,
    "dcgan": 0.8,
    "ddpm": 2.2,
    "cfg_ddpm": 2.6,
    "vqvae2": 1.4,
    "vit_gan": 4.0,
    "unknown": 1.0,
}


def generate_server_py(
    model_name: str,
    architecture: str = "unknown",
    hardware: str = "default",
    slack_webhook: Optional[str] = None,
    monthly_threshold: float = 500.0,
    latency_threshold: float = 500.0,
    persist_costs: bool = True,
) -> str:
    """
    Generate a complete FastAPI server.py with FRAMEWORM-COST integrated.

    This is called by `frameworm deploy` to produce the deployment server.
    """
    params = ARCH_PARAMS.get(architecture.lower(), 10.0)
    memory = ARCH_MEMORY.get(architecture.lower(), 1.0)

    slack_line = (
        f'    slack_webhook="{slack_webhook}",'
        if slack_webhook
        else '    # slack_webhook="https://hooks.slack.com/YOUR_WEBHOOK",'
    )

    persist_line = (
        '    persist_path="costs.json",'
        if persist_costs
        else '    # persist_path="costs.json",  # enable to save to disk'
    )

    return f'''"""
Auto-generated by FRAMEWORM-DEPLOY
Model: {model_name}
Architecture: {architecture}
Generated with FRAMEWORM-COST for inference cost tracking.
"""

import time
import torch
import onnxruntime as ort
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import Any, List
import numpy as np

# FRAMEWORM-COST — auto-injected by FRAMEWORM-DEPLOY
from cost import CostMiddleware, CostTracker
from cost.alerter import CostAlerter

# ── App ───────────────────────────────────────────────────────────────────────
app = FastAPI(
    title="{model_name} API",
    description="Auto-generated by FRAMEWORM-DEPLOY with cost tracking",
    version="1.0.0",
)

# ── Cost Middleware (tracks every /predict request automatically) ─────────────
app.add_middleware(
    CostMiddleware,
    model_name="{model_name}",
    architecture="{architecture}",
    hardware="{hardware}",
    parameters_millions={params},
    memory_gb={memory},
    {persist_line}
)

# ── Cost Alerter ──────────────────────────────────────────────────────────────
alerter = CostAlerter(
{slack_line}
    monthly_threshold_usd={monthly_threshold},
    latency_threshold_ms={latency_threshold},
    model_name="{model_name}",
)

# ── Cost Tracker (for manual wrapping if needed) ──────────────────────────────
tracker = CostTracker(
    model_name="{model_name}",
    architecture="{architecture}",
    hardware="{hardware}",
    parameters_millions={params},
    memory_gb={memory},
)

# ── Model Loading ─────────────────────────────────────────────────────────────
print("[DEPLOY] Loading ONNX model...")
_session = ort.InferenceSession("model.onnx")
_input_name = _session.get_inputs()[0].name
print(f"[DEPLOY] Model loaded. Input: {{_input_name}}")


# ── Schemas ───────────────────────────────────────────────────────────────────
class PredictRequest(BaseModel):
    inputs: List[List[float]]
    batch_size: int = 1

class PredictResponse(BaseModel):
    outputs: List[Any]
    latency_ms: float
    cost_usd: float
    model: str = "{model_name}"


# ── Endpoints ─────────────────────────────────────────────────────────────────
@app.get("/health")
def health():
    return {{"status": "ok", "model": "{model_name}", "architecture": "{architecture}"}}


@app.get("/ready")
def ready():
    if _session is None:
        raise HTTPException(status_code=503, detail="Model not loaded")
    return {{"status": "ready"}}


@app.post("/predict", response_model=PredictResponse)
def predict(req: PredictRequest):
    start = time.perf_counter()

    inputs = np.array(req.inputs, dtype=np.float32)
    outputs = _session.run(None, {{_input_name: inputs}})

    latency_ms = (time.perf_counter() - start) * 1000

    # Track cost and check alerts
    cost = tracker.calculator.calculate(latency_ms, batch_size=req.batch_size)
    tracker.store.record(cost)
    alerter.check(cost)

    return PredictResponse(
        outputs=[o.tolist() for o in outputs],
        latency_ms=round(latency_ms, 2),
        cost_usd=cost.total_cost_usd,
    )


@app.post("/batch")
def batch_predict(req: PredictRequest):
    """High-throughput batch endpoint. More cost-efficient than /predict."""
    start = time.perf_counter()
    inputs = np.array(req.inputs, dtype=np.float32)
    outputs = _session.run(None, {{_input_name: inputs}})
    latency_ms = (time.perf_counter() - start) * 1000

    cost = tracker.calculator.calculate(latency_ms, batch_size=len(req.inputs))
    tracker.store.record(cost)

    return {{
        "outputs": [o.tolist() for o in outputs],
        "latency_ms": round(latency_ms, 2),
        "cost_usd": cost.total_cost_usd,
        "items": len(req.inputs),
        "cost_per_item_usd": cost.total_cost_usd / len(req.inputs),
    }}


@app.get("/metrics")
def metrics():
    """Live latency and cost metrics."""
    return tracker.summary()


@app.get("/cost/summary")
def cost_summary():
    """Full cost breakdown with monthly projections."""
    return tracker.store.summary()


@app.get("/cost/hints")
def cost_hints():
    """Optimization recommendations based on observed traffic."""
    from cost.report import CostReport
    report = CostReport(tracker.store)
    return report.to_dict()


@app.get("/cost/alerts")
def cost_alerts():
    """Alert history and current thresholds."""
    return alerter.summary()


# ── Entry point ───────────────────────────────────────────────────────────────
if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)
'''


def generate_dockerfile(model_name: str, include_cost: bool = True) -> str:
    """Generate a Dockerfile that includes frameworm-cost."""
    cost_line = "    frameworm-cost \\\n" if include_cost else ""
    return f"""# Auto-generated by FRAMEWORM-DEPLOY
FROM python:3.11-slim

WORKDIR /app

RUN apt-get update && apt-get install -y curl && rm -rf /var/lib/apt/lists/*

COPY requirements.txt .
RUN pip install --no-cache-dir \\
    fastapi uvicorn onnxruntime numpy \\
{cost_line}    -r requirements.txt

COPY model.onnx .
COPY server.py .
COPY config.json .

EXPOSE 8000

HEALTHCHECK --interval=30s --timeout=5s --retries=3 \\
  CMD curl -f http://localhost:8000/health || exit 1

# Non-root user for security
RUN useradd -m appuser && chown -R appuser /app
USER appuser

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


def generate_requirements(architecture: str = "unknown") -> str:
    """Generate requirements.txt for the deployment."""
    base = [
        "fastapi>=0.95.0",
        "uvicorn>=0.22.0",
        "onnxruntime>=1.15.0",
        "numpy>=1.24.0",
        "pydantic>=2.0.0",
        "frameworm-cost>=0.1.0",
    ]
    if architecture in ("ddpm", "cfg_ddpm", "vqvae2"):
        base.append("torch>=2.0.0")
    return "\n".join(base) + "\n"
