#!/usr/bin/env python3
"""
forgepilot.py — minimal CLI for the ForgePilot API.

Setup:
    pip install requests
    export FORGEPILOT_API_KEY=fp_live_...   (generate one at forgepilot.dev/settings/api-keys)

Run a tool:
    python forgepilot.py fmea --context "Bolted flange, 5 kN axial, 200 bar seal" --file drawing.pdf
    python forgepilot.py tolerance-stackup --file assembly.step
    python forgepilot.py ask --context "What FoS should I use for a static aluminum bracket?"

Follow up on the same session (real context, not a cold prompt):
    python forgepilot.py chat "What about the tolerance on that mounting hole?"

Sessions are saved to forgepilot_session.json in the current directory by
default (--session to use a different path) — each run appends to history so
a follow-up actually knows what already happened, same as the web workspace.
"""

from __future__ import annotations

import argparse
import json
import os
import sys
from pathlib import Path

try:
    import requests
except ImportError:
    print("Missing dependency: pip install requests", file=sys.stderr)
    sys.exit(1)

BACKEND = "https://forgepilot-production-9575.up.railway.app"
DEFAULT_SESSION = "forgepilot_session.json"

# First line must exactly match a REPORT_MODE_MARKERS entry server-side to get
# full report-quality output (governing equations, verify-yourself section,
# committed verdicts) instead of a conversational answer.
TOOL_MARKERS = {
    "fmea": "FMEA GENERATOR AGENT",
    "test-plan": "TEST PLAN GENERATOR AGENT",
    "compliance-check": "STANDARDS COMPLIANCE CHECKER AGENT",
    "tolerance-stackup": "TOLERANCE STACK-UP AGENT",
    "drawing-review": "DRAWING REVIEWER AGENT",
    "ecn": "ENGINEERING CHANGE NOTICE GENERATOR",
    "beam": "ENGINEERING CALCULATION - Beam Deflection Analysis",
    "stress": "ENGINEERING CALCULATION - Stress Analysis",
    "pressure-vessel": "ENGINEERING CALCULATION - Pressure Vessel Analysis",
    "bolted-joint": "ENGINEERING CALCULATION - Bolted Joint Analysis",
    "bearing-life": "ENGINEERING CALCULATION - Bearing Life Analysis",
    "buckling": "ENGINEERING CALCULATION - Column Buckling Analysis",
    "fatigue": "ENGINEERING CALCULATION - Fatigue Analysis",
    "thermal": "ENGINEERING CALCULATION - Thermal Analysis",
}


def get_api_key(cli_value: str | None) -> str:
    key = cli_value or os.environ.get("FORGEPILOT_API_KEY")
    if not key:
        print(
            "No API key. Set FORGEPILOT_API_KEY or pass --api-key.\n"
            "Generate one at https://forgepilot.dev/settings/api-keys",
            file=sys.stderr,
        )
        sys.exit(1)
    return key


def load_session(path: str) -> list[dict]:
    p = Path(path)
    if not p.exists():
        return []
    try:
        return json.loads(p.read_text(encoding="utf-8"))
    except (json.JSONDecodeError, OSError):
        return []


def save_session(path: str, history: list[dict]) -> None:
    Path(path).write_text(json.dumps(history, indent=2), encoding="utf-8")


def call_analyze(api_key: str, prompt: str, history: list[dict], files: list[str]) -> str:
    open_files = [("files", (Path(f).name, open(f, "rb"))) for f in files]
    try:
        with requests.post(
            f"{BACKEND}/analyze",
            headers={"Authorization": f"Bearer {api_key}"},
            data={"prompt": prompt, "history": json.dumps(history)},
            files=open_files or None,
            stream=True,
            timeout=300,
        ) as r:
            if r.status_code == 401:
                print("\nAuthentication failed — check FORGEPILOT_API_KEY.", file=sys.stderr)
                sys.exit(1)
            if r.status_code == 429:
                body = r.json()
                print(
                    f"\nDaily limit reached ({body.get('count')}/{body.get('limit')}). "
                    "Resets at UTC midnight.",
                    file=sys.stderr,
                )
                sys.exit(1)
            r.raise_for_status()
            full_text = ""
            for chunk in r.iter_content(chunk_size=None, decode_unicode=True):
                if chunk:
                    print(chunk, end="", flush=True)
                    full_text += chunk
            print()
            return full_text
    finally:
        for _, (_, fh) in open_files:
            fh.close()


def cmd_run(args: argparse.Namespace) -> None:
    api_key = get_api_key(args.api_key)
    history = load_session(args.session)

    marker = TOOL_MARKERS.get(args.tool)
    if marker:
        prompt = f"{marker}\n\n{args.context or ''}".strip()
    else:
        prompt = args.context or ""
    if not prompt:
        print("Provide --context describing what to analyze.", file=sys.stderr)
        sys.exit(1)

    result = call_analyze(api_key, prompt, history, args.file or [])

    file_note = f" (files: {', '.join(Path(f).name for f in args.file)})" if args.file else ""
    history.append({"role": "user", "content": f"[Ran {args.tool}{file_note}]"})
    history.append({"role": "assistant", "content": result})
    save_session(args.session, history)


def cmd_chat(args: argparse.Namespace) -> None:
    api_key = get_api_key(args.api_key)
    history = load_session(args.session)
    result = call_analyze(api_key, args.message, history, args.file or [])
    history.append({"role": "user", "content": args.message})
    history.append({"role": "assistant", "content": result})
    save_session(args.session, history)


def main() -> None:
    parser = argparse.ArgumentParser(description="Run ForgePilot engineering analyses from the command line.")
    parser.add_argument("--api-key", help="Overrides FORGEPILOT_API_KEY.")
    parser.add_argument("--session", default=DEFAULT_SESSION, help=f"Session file (default: {DEFAULT_SESSION}).")
    sub = parser.add_subparsers(dest="command", required=True)

    for tool in list(TOOL_MARKERS) + ["ask"]:
        p = sub.add_parser(tool, help=f"Run {tool}.")
        p.add_argument("--context", help="What to analyze (required unless attaching a file that speaks for itself).")
        p.add_argument("--file", action="append", help="File to attach (repeatable).")
        p.set_defaults(func=cmd_run, tool=tool)

    p_chat = sub.add_parser("chat", help="Follow up on the current session.")
    p_chat.add_argument("message")
    p_chat.add_argument("--file", action="append", help="File to attach (repeatable).")
    p_chat.set_defaults(func=cmd_chat)

    args = parser.parse_args()
    args.func(args)


if __name__ == "__main__":
    main()