"""
Civil Engineer Agent — Inference Server

Commercial marketplace agent wrapping the fine-tuned Llama 3.1 8B
civil engineering model. Runs locally via Ollama, exposed as a
standard Brynq agent endpoint.

The model stays on the creator's machine. Users call through Brynq
marketplace — they never get the weights.

Endpoints:
    GET  /health   — Liveness check
    GET  /info     — Agent metadata
    POST /execute  — Run inference
"""

from __future__ import annotations

import json
import logging
import time
import urllib.error
import urllib.request
from typing import Any

from fastapi import FastAPI, HTTPException
from fastapi.responses import JSONResponse

log = logging.getLogger("brynq.agents.civil_engineer")

OLLAMA_URL = "http://localhost:11434"
MODEL_NAME = "civil-engineer"
AGENT_VERSION = "1.0.0"
START_TIME = time.time()

app = FastAPI(title="Civil Engineer Agent", version=AGENT_VERSION)


def _trim_repetition(text: str) -> str:
    """Cut output at the first sign of repetition.

    The fine-tuned model sometimes repeats answer blocks. This finds the
    first complete answer and stops there.
    """
    lines = text.split("\n")
    seen_answers = 0
    result_lines = []

    for line in lines:
        stripped = line.strip().lower()
        if stripped.startswith("**answer") or stripped.startswith("answer:"):
            seen_answers += 1
            if seen_answers > 1:
                break  # Stop at second answer block
        result_lines.append(line)

    return "\n".join(result_lines).strip()


def _ollama_chat(prompt: str, system_prompt: str = "", temperature: float = 0.4) -> dict[str, Any]:
    """Call the civil-engineer model via Ollama."""
    messages = []
    if system_prompt:
        messages.append({"role": "system", "content": system_prompt})
    messages.append({"role": "user", "content": prompt})

    body = json.dumps({
        "model": MODEL_NAME,
        "messages": messages,
        "stream": False,
        "options": {
            "temperature": temperature,
            "num_predict": 2048,
            "repeat_penalty": 1.5,
        },
    }).encode("utf-8")

    req = urllib.request.Request(
        f"{OLLAMA_URL}/api/chat",
        data=body,
        headers={"Content-Type": "application/json"},
        method="POST",
    )

    try:
        with urllib.request.urlopen(req, timeout=120) as resp:
            result = json.loads(resp.read().decode("utf-8"))
            return result
    except urllib.error.URLError as e:
        raise HTTPException(503, f"Ollama unreachable: {e}")
    except Exception as e:
        raise HTTPException(500, f"Inference error: {e}")


@app.get("/health")
async def health() -> JSONResponse:
    """Liveness check."""
    # Verify Ollama + model are accessible
    try:
        req = urllib.request.Request(f"{OLLAMA_URL}/api/tags", method="GET")
        with urllib.request.urlopen(req, timeout=5) as resp:
            data = json.loads(resp.read().decode("utf-8"))
            models = [m.get("name", "") for m in data.get("models", [])]
            model_ready = any(MODEL_NAME in m for m in models)
    except Exception:
        model_ready = False

    return JSONResponse({
        "status": "ok" if model_ready else "degraded",
        "uptime_seconds": round(time.time() - START_TIME, 1),
        "model": MODEL_NAME,
        "model_ready": model_ready,
    })


@app.get("/info")
async def info() -> JSONResponse:
    """Agent metadata."""
    return JSONResponse({
        "name": "Civil Engineer",
        "version": AGENT_VERSION,
        "description": (
            "Structural analysis, geotechnical, hydraulic, and transportation "
            "engineering. Trained on 90K verified problems with code references "
            "(AISC 360, IS 456, AASHTO LRFD, ACI 318, Eurocode)."
        ),
        "capabilities": [
            "structural_analysis",
            "beam_design",
            "column_design",
            "foundation_design",
            "hydraulics",
            "geotechnical",
            "transportation",
            "code_compliance",
        ],
        "protocol_version": 1,
        "model": MODEL_NAME,
        "parameters": "8B",
        "training_examples": 90549,
        "pricing": {"per_call": 0.02, "currency": "USD"},
    })


@app.post("/execute")
async def execute(request: dict[str, Any]) -> JSONResponse:
    """Run civil engineering inference."""
    prompt = request.get("prompt", "").strip()
    if not prompt:
        raise HTTPException(400, "prompt is required")

    context = request.get("context", {})
    config = request.get("config", {})

    # Build full prompt with context from prior chain steps
    parts = []
    for key, val in context.items():
        if val and isinstance(val, str):
            parts.append(f"Context ({key}):\n{val}")
    parts.append(prompt)
    full_prompt = "\n\n".join(parts)

    system_prompt = config.get("system_prompt", (
        "You are a civil engineering expert. Solve problems step-by-step "
        "with calculations and units. Cite governing codes where applicable. "
        "Keep answers concise and well-structured."
    ))
    temperature = config.get("temperature", 0.4)

    t0 = time.time()
    result = _ollama_chat(full_prompt, system_prompt, temperature)
    elapsed_ms = round((time.time() - t0) * 1000)

    raw_output = result.get("message", {}).get("content", "")
    eval_count = result.get("eval_count", 0)
    prompt_count = result.get("prompt_eval_count", 0)

    # Trim repetition: stop at first repeated "Answer:" block
    output = _trim_repetition(raw_output)

    return JSONResponse({
        "output": output,
        "tokens_used": eval_count + prompt_count,
        "elapsed_ms": elapsed_ms,
        "metadata": {
            "model": MODEL_NAME,
            "eval_tokens": eval_count,
            "prompt_tokens": prompt_count,
        },
    })


def main():
    """Run the inference server."""
    import uvicorn
    log.info("Starting Civil Engineer agent on port 9100")
    uvicorn.run(app, host="0.0.0.0", port=9100)


if __name__ == "__main__":
    main()
