pcswmm_tools.py

"""Public tool surface for the PCSWMM Engineering MCP.

This product retains its PCSWMM-engineering identity while supporting controlled deterministic evidence backends.  The public workflow starts from a
structured package exported by the PCSWMM Engineering SDK.  The legacy
uploaded-INP functions remain available internally so the service can perform
an independent SWMM verification, but they are not published as MCP tools.
"""
from __future__ import annotations

from pathlib import Path
from typing import Any, Callable

from sessions import STORE
import tools as core
from pcswmm_bridge import extract_inp_content
from revision_pipeline import revision_summary_text
from deterministic_evidence_adapter import (
    build_mcp_package_from_evidence, selected_figure_payloads, validate_evidence_directory,
)


def connect_active_pcswmm_project(
    package: dict | str,
    submission_type: str = "model_only_review",
    submission_number: str = "1",
    previous_submission_reference: str = "",
    auto_verify: bool = False,
) -> dict:
    """Connect the active PCSWMM project package and optionally verify it independently.

    The package must be exported by the PCSWMM Engineering SDK and should include
    project identity, active scenario context, engineering findings, and the exact
    INP text when independent verification is requested.
    """
    ingested = core.ingest_pcswmm_package(
        package=package,
        session_id="",
        create_model_session=True,
    )
    session_id = ingested["session_id"]
    configured = core.configure_submission(
        session_id=session_id,
        submission_type=submission_type,
        submission_number=submission_number,
        previous_submission_reference=previous_submission_reference,
    )

    verification: dict[str, Any] = {
        "status": "not_run",
        "reason": "auto_verify=False",
    }
    normalized = STORE.get(session_id).data.get("pcswmm_package", {})
    has_inp = bool(extract_inp_content(normalized))
    if auto_verify:
        if not has_inp:
            verification = {
                "status": "not_run",
                "reason": "The PCSWMM package does not contain INP text.",
            }
        else:
            verification = core.run_simulation(session_id)

    return {
        "session_id": session_id,
        "product": "PCSWMM Engineering MCP",
        "active_project": ingested["validation"],
        "submission": configured,
        "independent_verification": verification,
        "revision_review": ingested.get("revision_review", {}),
        "next_step": (
            "Call review_active_pcswmm_model, inspect findings, then generate_pcswmm_swmr."
        ),
    }



def connect_deterministic_calgary_evidence(
    evidence_dir: str,
    submission_type: str = "model_only_review",
    submission_number: str = "1",
    auto_attach_figures: bool = True,
    max_figures: int = 12,
) -> dict:
    """Connect an existing deterministic PCSWMM evidence directory to this MCP.

    This is the backend-neutral bridge for local workflows used by PCSWMM, PySWMM, Claude Desktop,
    ChatGPT Work, and Codex. The engineer runs PCSWMM and the native SDK first;
    this tool then ingests the resulting evidence directory without recalculating
    model values.
    """
    check = validate_evidence_directory(evidence_dir)
    if not check["valid"]:
        raise ValueError("Invalid evidence directory: " + "; ".join(check["missing"]))
    package = build_mcp_package_from_evidence(evidence_dir)
    result = connect_active_pcswmm_project(
        package=package,
        submission_type=submission_type,
        submission_number=submission_number,
        auto_verify=False,
    )
    session_id = result["session_id"]
    session = STORE.get(session_id)
    session.data["deterministic_evidence_dir"] = check["evidence_dir"]
    attached = []
    if auto_attach_figures:
        for payload in selected_figure_payloads(evidence_dir, max_figures=max_figures):
            attached.append(core.attach_figure(
                session_id=session_id,
                image_base64=payload["image_base64"],
                caption=payload["caption"],
                section=payload["section"],
            ))
    result.update({
        "evidence_directory": check,
        "attached_primary_figures": len(attached),
        "next_step": (
            "Call review_active_pcswmm_model, set_pcswmm_report_details, "
            "get_pcswmm_submission_readiness, then generate_pcswmm_swmr."
        ),
    })
    return result


def generate_calgary_swmr_from_evidence(
    evidence_dir: str,
    project_name: str = "",
    client: str = "",
    consultant: str = "",
    prepared_by: str = "",
    outline_plan_no: str = "",
    submission_type: str = "model_only_review",
    submission_number: str = "1",
    include_model_appendix: bool = True,
    max_figures: int = 12,
) -> dict:
    """One-command City of Calgary SWMR draft from a prior deterministic evidence run.

    Reuses the existing MCP review and report engines. It does not introduce a
    parallel SWMR generator and does not rerun PCSWMM.
    """
    connected = connect_deterministic_calgary_evidence(
        evidence_dir=evidence_dir,
        submission_type=submission_type,
        submission_number=submission_number,
        auto_attach_figures=True,
        max_figures=max_figures,
    )
    session_id = connected["session_id"]
    session = STORE.get(session_id)
    package = session.data.get("pcswmm_package", {}) or {}
    project = package.get("project", {}) or {}
    resolved_project = project_name or str(project.get("project_name") or "PCSWMM Project")
    details = {
        "client": client or project.get("client") or "Not provided",
        "consultant": consultant or project.get("consultant") or "Not provided",
        "prepared_by": prepared_by or project.get("prepared_by") or "Not provided",
        "outline_plan_no": outline_plan_no or project.get("outline_plan_number") or "Not provided",
        "municipality": project.get("municipality") or "City of Calgary",
    }
    core.set_report_details(session_id, details)
    review = review_active_pcswmm_model(session_id)
    readiness = core.get_submission_readiness(session_id)
    report = generate_pcswmm_swmr(
        session_id=session_id,
        project_name=resolved_project,
        client=str(details["client"]),
        consultant=str(details["consultant"]),
        prepared_by=str(details["prepared_by"]),
        outline_plan_no=str(details["outline_plan_no"]),
        include_model_appendix=include_model_appendix,
    )
    return {
        "session_id": session_id,
        "project_name": resolved_project,
        "review": review,
        "submission_readiness": readiness,
        "report": report,
        "evidence_directory": connected["evidence_directory"],
        "attached_primary_figures": connected["attached_primary_figures"],
        "note": "Draft generated from deterministic PCSWMM evidence for professional engineering review.",
    }


def connect_pyswmm_calgary_evidence(
    evidence_dir: str,
    submission_type: str = "model_only_review",
    submission_number: str = "1",
    auto_attach_figures: bool = True,
    max_figures: int = 12,
) -> dict:
    """Connect deterministic evidence produced by the PySWMM Calgary backend.

    Reuses the existing PCSWMM-engineering MCP review, readiness, figure, response
    matrix, and SWMR tools. Hydraulic values are not recalculated during connection.
    """
    result = connect_deterministic_calgary_evidence(
        evidence_dir=evidence_dir,
        submission_type=submission_type,
        submission_number=submission_number,
        auto_attach_figures=auto_attach_figures,
        max_figures=max_figures,
    )
    session = STORE.get(result["session_id"])
    package = session.data.get("pcswmm_package", {}) or {}
    backend = (package.get("source") or {}).get("backend")
    result["backend"] = backend or "PySWMM / EPA SWMM"
    result["workflow"] = "PySWMM deterministic evidence -> existing PCSWMM-engineering MCP"
    return result


def generate_calgary_swmr_from_pyswmm_evidence(
    evidence_dir: str,
    project_name: str = "",
    client: str = "",
    consultant: str = "",
    prepared_by: str = "",
    outline_plan_no: str = "",
    submission_type: str = "model_only_review",
    submission_number: str = "1",
    include_model_appendix: bool = True,
    max_figures: int = 12,
) -> dict:
    """Generate a City of Calgary SWMR draft from PySWMM deterministic evidence."""
    result = generate_calgary_swmr_from_evidence(
        evidence_dir=evidence_dir, project_name=project_name, client=client,
        consultant=consultant, prepared_by=prepared_by, outline_plan_no=outline_plan_no,
        submission_type=submission_type, submission_number=submission_number,
        include_model_appendix=include_model_appendix, max_figures=max_figures,
    )
    result["backend"] = "PySWMM / EPA SWMM"
    result["note"] = (
        "Draft generated by the existing PCSWMM-engineering MCP from deterministic "
        "PySWMM/EPA SWMM evidence for professional engineering review."
    )
    return result

def validate_active_pcswmm_project(session_id: str) -> dict:
    """Validate the connected PCSWMM package and report available capabilities."""
    result = core.validate_pcswmm_package(session_id)
    session = STORE.get(session_id)
    result["independent_results_available"] = bool(session.data.get("results"))
    result["active_project_connected"] = True
    return result


def run_independent_pcswmm_verification(session_id: str) -> dict:
    """Independently rerun the active PCSWMM model through the isolated SWMM worker."""
    session = STORE.get(session_id)
    if not session.data.get("pcswmm_package"):
        raise ValueError("This session is not connected to a PCSWMM engineering package.")
    if not session.data.get("inp_path"):
        raise ValueError(
            "The PCSWMM package did not include INP content; independent execution is unavailable."
        )
    return core.run_simulation(session_id)


def optimize_pcswmm_design(
    session_id: str,
    objective: dict | str,
    variables: dict | str | None = None,
    candidates: list | str | None = None,
    constraints: dict | str | None = None,
    outfall_link_id: str = "",
    max_evaluations: int = 24,
) -> dict:
    """Screen candidate designs against FIXED City of Calgary design-standard
    constraints and rank the feasible ones by a VARIABLE design objective.

    Requires run_independent_pcswmm_verification to have been called first on
    this session (optimization re-simulates the model, so it needs the same
    INP text and base run that independent verification requires).

    objective: {"metric": "maximum_link_velocity" | "maximum_modelled_depth_ratio"
                | "maximum_node_flooding" | "peak_link_flow"
                | "maximum_storage_volume" | "maximum_storage_depth"
                | "peak_subcatchment_runoff" | "total_conduit_volume_m3",
                "direction": "minimize" | "maximize"}
    variables: grid-search mode -- {"conduit_diameter_overrides":
                {"C1": [0.3, 0.375, 0.45]}, "storage_depth_overrides":
                {"ST1": [1.0, 1.5]}}; every combination is evaluated, capped
                by max_evaluations. Supported categories:
                conduit_diameter_overrides, conduit_roughness_overrides,
                storage_depth_overrides.
    candidates: shortlist mode -- a list of fully-formed override dicts to
                evaluate directly instead of a grid. Supply either variables
                or candidates, not conceptually both.
    constraints: overrides merged onto the fixed COC defaults (max_depth_ratio,
                min_velocity_mps, max_velocity_mps, max_node_flooding_cms,
                max_continuity_error_pct, max_allowable_outfall_flow_cms).
                These defaults are commonly-used stormwater design thresholds,
                NOT a transcription of a specific City of Calgary SWMDM or
                Industry Bulletin edition -- confirm every value against
                current criteria before treating a PASS as compliance.
    outfall_link_id: the link ID whose peak flow is checked against
                max_allowable_outfall_flow_cms, if that constraint is set.
    """
    if not isinstance(objective, dict) and not isinstance(objective, str):
        raise ValueError('objective must be an object like {"metric": ..., "direction": ...}')
    return core.optimize_design(
        session_id=session_id,
        objective=objective,
        variables=variables,
        candidates=candidates,
        constraints=constraints,
        outfall_link_id=outfall_link_id,
        max_evaluations=max_evaluations,
    )


def get_pcswmm_overland_flow_assessment(
    session_id: str,
    major_link_ids: list | str | None = None,
) -> dict:
    """Extract peak Q/v/d for major-system (overland) flow routes from the
    connected model and screen them against the Alberta Environment / City
    of Calgary depth-velocity criteria (SWMDM 2011, Table 3-20/Figure 3-23),
    returning the Table-11-13-style tabulation and a Figure-11-3-style chart
    ready to attach to the SWMR via attach_pcswmm_figure -- the extraction
    and calculation step consultants have otherwise had to do by hand from
    raw model output.

    major_link_ids: optional list of link IDs to treat as the major/overland
    system, matching your own drawing classification. If omitted, overland
    routes are auto-detected from the INP by cross-section shape
    (TRAPEZOIDAL, RECT_OPEN, TRIANGULAR, IRREGULAR, STREET).

    Requires run_independent_pcswmm_verification to have been called first
    on this session (this reads the same simulated node/link results).

    Returns: detection_method, the depth-velocity curve and its source
    citation, a unit_system flag, the tabulated per-link results (native
    model units), and chart_png_base64 -- a ready-to-attach PNG chart.
    """
    return core.get_overland_flow_assessment(session_id, major_link_ids=major_link_ids)


def get_pcswmm_discharge_volume_summary(
    session_id: str,
    locations: list | str | None = None,
    outfall_node_ids: list | str | None = None,
) -> dict:
    """Table 11-12-style permissible-discharge/runoff-volume summary per
    location/manhole, plus cumulative system volume at one or more outfalls
    (City of Calgary SWMDM 2011, Section 11.1.7.2.8 item iv).

    Requires run_independent_pcswmm_verification to have been called first
    on this session.

    From the model: Invert, Obvert, maximum HGL, and Runoff Volume (a
    trapezoidal integration of that node's simulated inflow, independently
    validated against the SWMM engine's own continuity totals to 9
    significant figures) are computed directly. Area is auto-summed from
    subcatchments draining to that node unless overridden.

    NOT computed, by design -- must be supplied if wanted in the table:
    permissible Discharge Rate (L/s/ha, a City-assigned/design value) and
    Storage Volume (a preliminary on-site storage result, e.g. from the
    wbscc or calgary-storm-retention tools). Omitted fields are left blank
    with a note, never fabricated.

    locations: [{"node_id": "MH17-3", "location_label": "Site 1",
        "manhole_number": "17-3", "area_ha": 1.2 (optional override),
        "discharge_rate_lps_ha": 50 (optional), "storage_volume_m3": 120
        (optional)}, ...]
    outfall_node_ids: outfall node IDs to compute cumulative volume for; if
        omitted, every OUTFALLS node in the model is used.
    """
    return core.get_discharge_volume_summary(
        session_id, locations=locations, outfall_node_ids=outfall_node_ids
    )


def get_active_pcswmm_project(session_id: str) -> dict:
    """Return connected PCSWMM project, package, scenario, and execution identity."""
    session = STORE.get(session_id)
    validation = session.data.get("pcswmm_validation", {})
    package = session.data.get("pcswmm_package", {})
    project = package.get("project") or package.get("project_summary") or {}
    source = package.get("source") or {}
    metadata = (session.data.get("results") or {}).get("metadata", {})
    return {
        "session_id": session_id,
        "project": project,
        "source": source,
        "package_validation": validation,
        "active_scenario": package.get("active_scenario") or project.get("active_scenario"),
        "pcswmm_sdk_version": source.get("sdk_version") or package.get("sdk_version"),
        "independent_execution": {
            "available": bool(session.data.get("results")),
            "run_id": metadata.get("run_id"),
            "model_sha256": metadata.get("model_sha256"),
            "results_usable": metadata.get("results_usable"),
            "engine": "EPA SWMM / pyswmm isolated local worker",
        },
    }


def review_active_pcswmm_model(session_id: str) -> dict:
    """Run the consolidated deterministic review of the connected active PCSWMM model."""
    session = STORE.get(session_id)
    if not session.data.get("pcswmm_package"):
        raise ValueError("No active PCSWMM project is connected to this session.")

    package_validation = core.validate_pcswmm_package(session_id)
    package = session.data.get("pcswmm_package", {}) or {}
    revision = session.data.get("revision_review", {}) or {}
    submission_type = str(session.data.get("submission_type", "first_submission"))
    revision_applicable = submission_type in {"revised_submission", "final_submission"} or bool(revision.get("object_impacts"))
    revision_review = (
        {"applicable": True, "summary": revision_summary_text(revision), **revision}
        if revision.get("object_impacts")
        else {
            "applicable": False,
            "status": "not_applicable" if not revision_applicable else "missing",
            "message": (
                "Revision review is not applicable to this first/model-only submission."
                if not revision_applicable else
                "A revised/final submission requires a baseline snapshot and structured revision review."
            ),
        }
    )
    review: dict[str, Any] = {
        "session_id": session_id,
        "package_validation": package_validation,
        "revision_review": revision_review,
        "pcswmm_engineering_review": package.get("engineering_review", {}),
        "pcswmm_result_summary": (package.get("results", {}) or {}),
    }

    if session.data.get("results"):
        review.update({
            "execution_reconciliation": core.get_reconciliation(session_id),
            "calgary_screening": core.calgary_screening(session_id),
            "engineering_findings": core.preliminary_design_review(session_id),
            "node_results": core.get_node_results(session_id, limit=20),
            "link_results": core.get_link_results(session_id, limit=20),
        })
    else:
        review["independent_verification"] = {
            "status": "not_available",
            "reason": "Run run_independent_pcswmm_verification before relying on hydraulic conclusions.",
        }

    reasoning = ((package.get("engineering_review") or {}).get("reasoning_narrative"))
    review["executive_summary"] = (
        revision_summary_text(revision) if revision.get("object_impacts")
        else (reasoning or "PCSWMM-native engineering review completed from the active project package.")
    )
    review["disclaimer"] = (
        "Deterministic engineering screening for professional review; not a sealed design determination."
    )
    return review


def get_pcswmm_node_results(
    session_id: str,
    node_type: str = "",
    sort_by: str = "Depth Ratio",
    limit: int = 20,
) -> dict:
    """Return independently verified node results for the active PCSWMM project."""
    return core.get_node_results(session_id, node_type=node_type, sort_by=sort_by, limit=limit)


def get_pcswmm_link_results(
    session_id: str,
    sort_by: str = "Peak Velocity (m/s)",
    limit: int = 20,
) -> dict:
    """Return independently verified link results for the active PCSWMM project."""
    return core.get_link_results(session_id, sort_by=sort_by, limit=limit)


def get_pcswmm_subcatchment_results(session_id: str, limit: int = 30) -> dict:
    """Return independently verified subcatchment results for the active PCSWMM project."""
    return core.get_subcatchment_results(session_id, limit=limit)


def get_pcswmm_timeseries(
    session_id: str,
    object_type: str,
    object_id: str,
    variable: str,
) -> dict:
    """Return a bounded time series for an active PCSWMM project object."""
    return core.get_timeseries(session_id, object_type, object_id, variable)


def review_pcswmm_revision(session_id: str) -> dict:
    """Return the normalized baseline-to-revised PCSWMM impact assessment."""
    return core.get_design_revision_review(session_id)


def configure_pcswmm_submission(
    session_id: str,
    submission_type: str = "first_submission",
    submission_number: str = "1",
    previous_submission_reference: str = "",
    certified_checklist_attached: bool = False,
    no_objections_letter_attached: bool = False,
    prior_approved_report_attached: bool = False,
    correspondence_appendix_attached: bool = False,
    clean_copy: bool = False,
) -> dict:
    """Configure the Calgary submission route for the active PCSWMM project."""
    return core.configure_submission(
        session_id=session_id,
        submission_type=submission_type,
        submission_number=submission_number,
        previous_submission_reference=previous_submission_reference,
        certified_checklist_attached=certified_checklist_attached,
        no_objections_letter_attached=no_objections_letter_attached,
        prior_approved_report_attached=prior_approved_report_attached,
        correspondence_appendix_attached=correspondence_appendix_attached,
        clean_copy=clean_copy,
    )


def set_pcswmm_city_comments(
    session_id: str,
    comments: list | str,
    report_changes: dict | str = {},
) -> dict:
    """Attach City comments and documented changes to the active PCSWMM project session."""
    return core.set_city_comments(session_id, comments, report_changes)


def build_pcswmm_city_response_matrix(session_id: str) -> dict:
    """Build the evidence-linked City comment response matrix for the PCSWMM project."""
    return core.build_city_comment_response_matrix(session_id)


def get_pcswmm_submission_readiness(session_id: str) -> dict:
    """Apply staged draft, engineer-review, and submission-readiness gates."""
    return core.get_submission_readiness(session_id)


def get_pcswmm_hydraulic_summary(session_id: str) -> dict:
    """Return critical deterministic hydraulic results without rerunning the model."""
    return core.get_hydraulic_summary(session_id)


def set_pcswmm_report_details(session_id: str, details: dict | str) -> dict:
    """Set project, client, consultant, author, and report metadata."""
    return core.set_report_details(session_id, details)


def set_pcswmm_report_configuration(session_id: str, configuration: dict | str) -> dict:
    """Set confirmed report criteria and presentation options."""
    return core.set_report_configuration(session_id, configuration)


def attach_pcswmm_figure(
    session_id: str,
    image_base64: str,
    caption: str,
    section: str = "Engineering Review",
) -> dict:
    """Attach a PCSWMM profile, graph, map, or engineering figure to the report."""
    return core.attach_figure(session_id, image_base64, caption, section)


def _native_pcswmm_swmr(session, project_name: str, client: str, consultant: str,
                           prepared_by: str, outline_plan_no: str,
                           include_model_appendix: bool) -> dict:
    """Generate a structured consultant-demonstration SWMR from PCSWMM-native evidence."""
    from pcswmm_native_report import generate_native_swmr
    return generate_native_swmr(
        session=session, project_name=project_name, client=client, consultant=consultant,
        prepared_by=prepared_by, outline_plan_no=outline_plan_no,
        include_model_appendix=include_model_appendix,
    )

def generate_pcswmm_swmr(
    session_id: str,
    project_name: str,
    client: str = "",
    consultant: str = "",
    prepared_by: str = "",
    outline_plan_no: str = "",
    include_model_appendix: bool = True,
) -> dict:
    """Generate an SWMR from PCSWMM evidence; use the full reconciled engine when optional verification exists."""
    session = STORE.get(session_id)
    if not session.data.get("pcswmm_package"):
        raise ValueError("A connected PCSWMM project is required.")
    if session.data.get("results"):
        return core.generate_report(
            session_id=session_id, project_name=project_name, client=client,
            consultant=consultant, prepared_by=prepared_by,
            outline_plan_no=outline_plan_no,
            include_model_appendix=include_model_appendix,
        )
    return _native_pcswmm_swmr(session, project_name, client, consultant, prepared_by,
                               outline_plan_no, include_model_appendix)

def list_pcswmm_sessions() -> dict:
    """List local PCSWMM Engineering MCP sessions."""
    return core.list_sessions()


def close_pcswmm_session(session_id: str) -> dict:
    """Close a PCSWMM Engineering MCP session and remove temporary working files."""
    return core.close_session(session_id)


PCSWMM_TOOL_REGISTRY: dict[str, Callable[..., dict]] = {
    fn.__name__: fn
    for fn in [
        connect_active_pcswmm_project,
        connect_deterministic_calgary_evidence,
        connect_pyswmm_calgary_evidence,
        generate_calgary_swmr_from_evidence,
        generate_calgary_swmr_from_pyswmm_evidence,
        validate_active_pcswmm_project,
        get_active_pcswmm_project,
        run_independent_pcswmm_verification,
        optimize_pcswmm_design,
        get_pcswmm_overland_flow_assessment,
        get_pcswmm_discharge_volume_summary,
        review_active_pcswmm_model,
        get_pcswmm_node_results,
        get_pcswmm_link_results,
        get_pcswmm_subcatchment_results,
        get_pcswmm_timeseries,
        review_pcswmm_revision,
        configure_pcswmm_submission,
        set_pcswmm_city_comments,
        build_pcswmm_city_response_matrix,
        get_pcswmm_submission_readiness,
        get_pcswmm_hydraulic_summary,
        set_pcswmm_report_details,
        set_pcswmm_report_configuration,
        attach_pcswmm_figure,
        generate_pcswmm_swmr,
        list_pcswmm_sessions,
        close_pcswmm_session,
    ]
}