# ******************************************************************************
#
# Pentaho
#
# Copyright (C) 2026 by Pentaho Canada Inc. : http://www.pentaho.com
#
# Use of this software is governed by the Business Source License included
# in the LICENSE.TXT file.
#
# Change Date: 2030-06-15
# ******************************************************************************
# =============================================================================
# agent.py — Maintenance Assessment Agent
# Agent as a Service  |  Pentaho Academy
#
# Single endpoint: POST /assess
# Receives: log_id, asset_id, log_text, history (list of {logged_at, log_text})
# Returns:  priority, fault_type, pattern, assessment, confidence
#
# Run:
#   source agent-venv/bin/activate          # Linux / macOS
#   .\agent-venv\Scripts\Activate.ps1       # Windows PowerShell
#   uvicorn agent:app --host 0.0.0.0 --port 8000
# =============================================================================

import os
import json
import logging
import httpx
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel

# ── Configuration ─────────────────────────────────────────────────────────────
MODEL_URL   = os.getenv("AGENT_MODEL_URL",   "http://localhost:11434/api/generate")
MODEL_NAME  = os.getenv("AGENT_MODEL_NAME",  "llama3.1:8b")
TEMPERATURE = float(os.getenv("AGENT_TEMPERATURE", "0.1"))
TIMEOUT     = int(os.getenv("AGENT_TIMEOUT",       "120"))
MAX_RETRIES = int(os.getenv("AGENT_MAX_RETRIES",   "2"))

logging.basicConfig(level=logging.INFO)
log = logging.getLogger("maintenance-agent")

app = FastAPI(title="Maintenance Assessment Agent", version="1.0.0")


# ── Pydantic models ───────────────────────────────────────────────────────────
class HistoryEntry(BaseModel):
    logged_at: str
    log_text:  str

class AssessRequest(BaseModel):
    log_id:   str
    asset_id: str
    log_text: str
    history:  list[HistoryEntry]

class AssessResponse(BaseModel):
    log_id:     str
    asset_id:   str
    priority:   str
    fault_type: str
    pattern:    str
    assessment: str
    confidence: int


# ── Prompt builder ────────────────────────────────────────────────────────────
def build_prompt(req: AssessRequest) -> str:
    """
    Build the assessment prompt.

    History is presented BEFORE the current entry, oldest first.
    Priority and pattern rules are explicit to reduce model ambiguity.
    """
    if req.history:
        hist_lines = []
        for i, h in enumerate(req.history, 1):
            hist_lines.append(f"  {i}. [{h.logged_at}] {h.log_text}")
        history_block = (
            f"Prior history for {req.asset_id} (oldest first):\n"
            + "\n".join(hist_lines)
        )
    else:
        history_block = (
            f"Prior history for {req.asset_id}: "
            "none (first recorded entry for this asset)"
        )

    return f"""You are a maintenance engineer assistant.
Assess the current log entry for asset {req.asset_id} given its history.
Return ONLY valid JSON, no other text:
{{
  "priority":   "CRITICAL|HIGH|MEDIUM|LOW",
  "fault_type": "short_snake_case_label",
  "pattern":    "RECURRENCE|ESCALATION|NEW_FAULT|NORMAL_VARIATION",
  "assessment": "one or two sentences explaining your reasoning",
  "confidence": 0-100
}}

Priority rules - apply in order, use the first that matches:
- CRITICAL: a hard safety limit has been breached (e.g. temperature above alarm threshold),
  or the same fault is actively worsening within the current shift
- HIGH: current symptoms match the description or fault type of a prior history entry
  (bearing noise, vibration, temperature drift), especially if a repair was done recently
  and symptoms are returning
- MEDIUM: fault is novel with no matching history entry, requires investigation within 48 hours
- LOW: engineer explicitly states the observation is within normal range and no action is needed

Pattern rules:
- RECURRENCE: current symptoms closely match the description or fault type of one or more
  prior history entries for this asset
- ESCALATION: the same fault described in an earlier entry from the CURRENT shift is now
  worsening (use when two entries for the same asset appear on the same day and the second
  is more severe)
- NEW_FAULT: no prior history entry describes anything similar to the current symptom
- NORMAL_VARIATION: engineer explicitly states the observation is normal, no concern raised

{history_block}

Current entry [{req.asset_id}]: {req.log_text}"""


# ── LLM call ──────────────────────────────────────────────────────────────────
async def call_llm(prompt: str) -> str:
    async with httpx.AsyncClient(timeout=TIMEOUT) as client:
        r = await client.post(MODEL_URL, json={
            "model":   MODEL_NAME,
            "prompt":  prompt,
            "stream":  False,
            "format":  "json",
            "options": {
                "temperature": TEMPERATURE,
                "num_predict": 300
            }
        })
        r.raise_for_status()
        return r.json()["response"]


def extract_json(text: str) -> dict:
    s = text.find("{")
    e = text.rfind("}") + 1
    if s < 0 or e <= s:
        raise ValueError(f"No JSON object found in response: {text[:150]!r}")
    return json.loads(text[s:e])


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


@app.post("/assess", response_model=AssessResponse)
async def assess(req: AssessRequest):
    log.info(
        f"Assessing {req.log_id} for {req.asset_id} "
        f"({len(req.history)} history entries)"
    )

    prompt     = build_prompt(req)
    last_error = None

    for attempt in range(MAX_RETRIES + 1):
        try:
            raw  = await call_llm(prompt)
            data = extract_json(raw)

            priority   = str(data.get("priority",   "MEDIUM")).upper()
            fault_type = str(data.get("fault_type", "unknown"))
            pattern    = str(data.get("pattern",    "NEW_FAULT")).upper()
            assessment = str(data.get("assessment", ""))
            confidence = int(data.get("confidence", 50))
            confidence = max(0, min(100, confidence))

            return AssessResponse(
                log_id     = req.log_id,
                asset_id   = req.asset_id,
                priority   = priority,
                fault_type = fault_type,
                pattern    = pattern,
                assessment = assessment,
                confidence = confidence,
            )

        except (ValueError, KeyError, json.JSONDecodeError) as e:
            last_error = e
            log.warning(f"Attempt {attempt + 1} failed for {req.log_id}: {e}")
            prompt = (
                "You must return ONLY a raw JSON object. "
                "No markdown, no explanation, no code fences.\n\n"
                + build_prompt(req)
            )

        except httpx.HTTPError as e:
            log.error(f"LLM backend error for {req.log_id}: {e}")
            raise HTTPException(
                status_code=502,
                detail=f"LLM backend unavailable: {e}"
            )

    log.error(
        f"All {MAX_RETRIES + 1} attempts failed for {req.log_id}: {last_error}"
    )
    raise HTTPException(
        status_code=500,
        detail=f"Failed after {MAX_RETRIES + 1} attempts: {last_error}"
    )