ai_features.py

import os
import json
import re
import logging
from openai import OpenAI

logger = logging.getLogger(__name__)

# the newest OpenAI model is "gpt-5" which was released August 7, 2025.
# do not change this unless explicitly requested by the user
AI_MODEL = "gpt-5"

AI_INTEGRATIONS_OPENAI_API_KEY = os.environ.get("AI_INTEGRATIONS_OPENAI_API_KEY")
AI_INTEGRATIONS_OPENAI_BASE_URL = os.environ.get("AI_INTEGRATIONS_OPENAI_BASE_URL")

client = OpenAI(
    api_key=AI_INTEGRATIONS_OPENAI_API_KEY,
    base_url=AI_INTEGRATIONS_OPENAI_BASE_URL
)


def _extract_json(text: str) -> dict:
    if not text:
        return {}
    text = text.strip()
    if text.startswith("```"):
        text = re.sub(r'^```\w*\n?', '', text)
        text = re.sub(r'\n?```$', '', text)
        text = text.strip()
    try:
        return json.loads(text)
    except json.JSONDecodeError:
        match = re.search(r'\{[\s\S]*\}', text)
        if match:
            try:
                return json.loads(match.group())
            except json.JSONDecodeError:
                pass
        return {}


def ai_review_rpt(rpt_content: str, inp_content: str = "") -> dict:
    system_prompt = """You are a senior stormwater engineer reviewing EPA SWMM simulation results.
Write a professional technical review memo. Include:

1. **Executive Summary** — 2-3 sentence overview
2. **Key Findings** — bullet list of important results
3. **Warnings & Concerns** — nodes near flooding, conduits near capacity, continuity errors
4. **Recommendations** — specific actionable suggestions
5. **Rating** — PASS, MARGINAL, or FAIL with justification

Use specific node/conduit names. Cite numbers. Be concise. Format in markdown."""

    user_msg = f"Review this SWMM simulation report:\n\n{rpt_content[:6000]}"
    if inp_content:
        user_msg += f"\n\nModel input for context:\n{inp_content[:3000]}"

    try:
        response = client.chat.completions.create(
            model=AI_MODEL,
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_msg}
            ],
        )
        memo = response.choices[0].message.content or ""
        if not memo:
            return {"status": "error", "message": "AI returned empty response. Please try again."}
        return {"status": "success", "memo": memo}
    except Exception as e:
        logger.error(f"AI review error: {e}")
        return {"status": "error", "message": str(e)}


def ai_build_model(description: str) -> dict:
    system_prompt = """You are an expert EPA SWMM modeler. Given a plain-English description,
generate a complete, valid SWMM 5.x .INP file.

Rules:
- Include all required sections: [TITLE], [OPTIONS], [EVAPORATION], [RAINGAGES], [SUBCATCHMENTS], [SUBAREAS], [INFILTRATION], [JUNCTIONS], [OUTFALLS], [CONDUITS], [XSECTIONS], [TIMESERIES], [REPORT], [COORDINATES], [Polygons]
- Use realistic parameter values
- Use KINWAVE flow routing
- Include a design storm timeseries
- Include coordinates for visualization
- Use CFS flow units by default
- Name elements: S1, S2 for subcatchments; J1, J2 for junctions; C1, C2 for conduits

Return ONLY the .INP file content. Start with [TITLE]. No markdown fencing or extra text."""

    try:
        response = client.chat.completions.create(
            model=AI_MODEL,
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": f"Create a SWMM model for: {description}"}
            ],
        )
        inp_content = response.choices[0].message.content or ""
        inp_content = inp_content.strip()
        if inp_content.startswith("```"):
            inp_content = re.sub(r'^```\w*\n?', '', inp_content)
            inp_content = re.sub(r'\n?```$', '', inp_content)
            inp_content = inp_content.strip()

        if not ("[TITLE]" in inp_content and "[OPTIONS]" in inp_content):
            return {"status": "error", "message": "Generated model is missing required sections. Please try again with a more specific description."}

        return {"status": "success", "inp_content": inp_content}
    except Exception as e:
        logger.error(f"AI model build error: {e}")
        return {"status": "error", "message": str(e)}


def ai_what_if(question: str, inp_content: str, previous_results: str = "") -> dict:
    system_prompt = """You are an expert SWMM modeling assistant answering "what if" questions.

Your job:
1. Explain what would happen in plain English
2. List specific changes needed
3. Give your engineering recommendation

Respond in this exact JSON format (no markdown fencing):
{"explanation":"...","changes_made":["change1","change2"],"modified_inp":null,"confidence":"high","recommendation":"..."}

Set modified_inp to null. Keep explanation concise (2-3 paragraphs max). Use actual element names from the model."""

    user_msg = f"My SWMM model:\n{inp_content[:4000]}\n\nQuestion: {question}"

    try:
        response = client.chat.completions.create(
            model=AI_MODEL,
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_msg}
            ],
        )
        raw = response.choices[0].message.content or ""
        if not raw:
            return {"status": "error", "message": "AI returned empty response. Please try again."}
        result = _extract_json(raw)
        if not result:
            return {"status": "success", "result": {"explanation": raw, "changes_made": [], "modified_inp": None, "confidence": "medium", "recommendation": ""}}
        return {"status": "success", "result": result}
    except Exception as e:
        logger.error(f"AI what-if error: {e}")
        return {"status": "error", "message": str(e)}


def ai_anomaly_detection(inp_content: str, rpt_content: str = "") -> dict:
    system_prompt = """You are a senior SWMM QA engineer. Analyze the model for anomalies and issues.

Check for: numerical instability, configuration issues, parameter warnings, hydraulic issues, data quality.

Respond in this exact JSON format (no markdown fencing):
{"anomalies":[{"severity":"warning","category":"parameter","element":"J1","description":"issue desc","recommendation":"fix"}],"overall_health":"healthy","summary":"summary paragraph"}

Severity: critical, warning, or info. Category: numerical, configuration, parameter, hydraulic, or data.
Use actual element names. Be specific and concise."""

    user_msg = ""
    if inp_content:
        user_msg += f"SWMM Input File:\n{inp_content[:5000]}\n"
    if rpt_content:
        user_msg += f"\nSWMM Report File:\n{rpt_content[:5000]}"
    if not user_msg:
        return {"status": "error", "message": "No input or report content provided"}

    try:
        response = client.chat.completions.create(
            model=AI_MODEL,
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_msg}
            ],
        )
        raw = response.choices[0].message.content or ""
        if not raw:
            return {"status": "error", "message": "AI returned empty response. Please try again."}
        result = _extract_json(raw)
        if not result:
            return {"status": "success", "result": {"anomalies": [], "overall_health": "unknown", "summary": raw}}
        return {"status": "success", "result": result}
    except Exception as e:
        logger.error(f"AI anomaly detection error: {e}")
        return {"status": "error", "message": str(e)}