app.py

import os
import json
import html
from typing import Any, Dict, List, Optional, Tuple

import requests
from dotenv import load_dotenv
from fastapi import FastAPI
from fastapi.responses import HTMLResponse, JSONResponse
from pydantic import BaseModel

from openai import OpenAI

# ===============================
# ENV / CONFIG (PROD-like)
# ===============================
load_dotenv()

DEBUG_STARTUP_LOGS = os.getenv("DEBUG_STARTUP_LOGS", "0").strip().lower() in ("1", "true", "yes")

OPENAI_API_KEY = os.getenv("OPENAI_API_KEY", "").strip()
if not OPENAI_API_KEY:
    raise RuntimeError("OPENAI_API_KEY is missing. Put it into .env")

QDRANT_URL = os.getenv("QDRANT_URL", "http://127.0.0.1:6333").strip().rstrip("/")
QDRANT_COLLECTION = os.getenv("QDRANT_COLLECTION", "pms_equipment").strip()
QDRANT_API_KEY = os.getenv("QDRANT_API_KEY", "").strip()

EMBED_MODEL = os.getenv("EMBED_MODEL", "text-embedding-3-small").strip()
VECTOR_SIZE = int(os.getenv("VECTOR_SIZE", "1536").strip())
TOP_K = int(os.getenv("TOP_K", "5").strip())

# ===============================
# Evidence gate (PROD)
# ===============================
SCORE_THRESHOLD = float(os.getenv("SCORE_THRESHOLD", "0.55"))
MIN_STRONG_HITS = int(os.getenv("MIN_STRONG_HITS", "1"))

# ===============================
# Payload / token hygiene
# ===============================
MAX_QUERY_CHARS = int(os.getenv("MAX_QUERY_CHARS", "800").strip())
MIN_QUERY_CHARS = int(os.getenv("MIN_QUERY_CHARS", "3").strip())
MAX_EVIDENCE_CHARS = int(os.getenv("MAX_EVIDENCE_CHARS", "12000").strip())
RETURN_RAW_HITS = os.getenv("RETURN_RAW_HITS", "1").strip().lower() in ("1", "true", "yes")

# ===============================
# LLM
# ===============================
LLM_MODEL = os.getenv("LLM_MODEL", "gpt-4o-mini").strip()  # JSON-only audit answer

if DEBUG_STARTUP_LOGS:
    print("QDRANT_URL =", QDRANT_URL)
    print("QDRANT_COLLECTION =", QDRANT_COLLECTION)
    print("QDRANT_API_KEY =", "SET" if QDRANT_API_KEY else "MISSING")
    print("EMBED_MODEL =", EMBED_MODEL)
    print("VECTOR_SIZE =", VECTOR_SIZE)
    print("TOP_K =", TOP_K)
    print("LLM_MODEL =", LLM_MODEL)

# ===============================
# CLIENTS
# ===============================
oai = OpenAI(api_key=OPENAI_API_KEY)

# ===============================
# APP
# ===============================
app = FastAPI(title="PMS Copilot — RAG MVP")


# ============================================================
# SCHEMAS
# ============================================================
class AskRequest(BaseModel):
    q: str


# ============================================================
# HELPERS
# ============================================================
def embed(text: str) -> List[float]:
    """OpenAI embeddings -> vector[VECTOR_SIZE]."""
    resp = oai.embeddings.create(model=EMBED_MODEL, input=text)
    vec = resp.data[0].embedding
    if len(vec) != VECTOR_SIZE:
        raise RuntimeError(
            f"Embedding dim mismatch: got {len(vec)} but VECTOR_SIZE={VECTOR_SIZE}. "
            f"Check EMBED_MODEL / VECTOR_SIZE in .env"
        )
    return vec


def qdrant_search_rest(query_vec: List[float], limit: int) -> List[Dict[str, Any]]:
    """

    Qdrant REST search (robust, avoids qdrant_client version/SyncApis issues).

    Returns list of points: [{"id":..., "score":..., "payload": {...}}, ...]

    """
    url = f"{QDRANT_URL}/collections/{QDRANT_COLLECTION}/points/search"
    payload = {
        "vector": query_vec,
        "limit": limit,
        "with_payload": True,
        "with_vectors": False,
    }

    headers: Dict[str, str] = {}
    # Qdrant Cloud/self-host can require an API key. For Qdrant Cloud, "api-key" is commonly used.
    if QDRANT_API_KEY:
        headers["api-key"] = QDRANT_API_KEY

    r = requests.post(url, json=payload, headers=headers, timeout=30)
    r.raise_for_status()
    data = r.json()
    return data.get("result", [])


def pick_text_from_payload(payload: Dict[str, Any]) -> Optional[str]:
    """Extract readable text from payload (support common field names)."""
    for k in ("text", "chunk", "content", "page_content", "body", "passage", "PROCEDURE"):
        v = payload.get(k)
        if isinstance(v, str) and v.strip():
            return v.strip()

    if payload:
        keys_pref = ["GROUPS", "FREQUENCY TYPE", "MAINTENANCE HEAD", "RESPONSIBILITY", "PROCEDURE"]
        parts = []
        for k in keys_pref:
            if k in payload and payload[k] not in (None, ""):
                parts.append(f"{k}: {payload[k]}")
        if parts:
            return " | ".join(parts)

    return None


def build_evidence_blocks(hits: List[Dict[str, Any]]) -> Tuple[str, List[Dict[str, Any]]]:
    """

    Build evidence list for LLM:

    - evidence_text: lines like [1] ...

    - sources: minimal metadata for UI

    """
    evidence_lines: List[str] = []
    sources: List[Dict[str, Any]] = []

    for i, h in enumerate(hits, start=1):
        payload = h.get("payload") or {}
        text = pick_text_from_payload(payload) or ""
        text = text.replace("\r", " ").replace("\n", " ").strip()
        if not text:
            text = json.dumps(payload, ensure_ascii=False)

        evidence_lines.append(f"[{i}] {text}")

        sources.append(
            {
                "n": i,
                "id": h.get("id"),
                "score": h.get("score"),
                "GROUPS": payload.get("GROUPS"),
                "FREQUENCY TYPE": payload.get("FREQUENCY TYPE"),
                "MAINTENANCE HEAD": payload.get("MAINTENANCE HEAD"),
                "RESPONSIBILITY": payload.get("RESPONSIBILITY"),
            }
        )

    evidence_text = "\n".join(evidence_lines)
    if len(evidence_text) > MAX_EVIDENCE_CHARS:
        evidence_text = evidence_text[:MAX_EVIDENCE_CHARS] + "\n...[TRUNCATED]"

    return evidence_text, sources


def _extract_first_json_object(s: str) -> str:
    """

    Best-effort recovery if LLM outputs extra text.

    Returns substring from first '{' to last '}'.

    """
    if not s:
        return s
    start = s.find("{")
    end = s.rfind("}")
    if start == -1 or end == -1 or end <= start:
        return s
    return s[start : end + 1]


def run_llm_audit_json(query: str, evidence_text: str) -> Dict[str, Any]:
    """

    LLM audit-style answer.

    STRICT JSON ONLY (enforced by system contract + JSON parse).

    """
    system_prompt = """

You are a maritime audit assistant.



RULES (MANDATORY):

- Output MUST be valid JSON

- NO markdown

- NO explanations

- NO text outside JSON

- Use ONLY the provided evidence

- If information is missing, use "Not found in provided records"



JSON SCHEMA (exact):

{

  "summary": string,

  "findings": [

    {

      "topic": string,

      "requirement": string,

      "observation": string,

      "risk": string,

      "evidence_refs": [number]

    }

  ],

  "conclusion": string

}

""".strip()

    user_prompt = f"""

AUDIT QUESTION:

{query}



EVIDENCE:

{evidence_text}

""".strip()

    resp = oai.responses.create(
        model=LLM_MODEL,
        input=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt},
        ],
        temperature=0,
    )

    raw = resp.output_text or ""
    candidate = _extract_first_json_object(raw)

    try:
        return json.loads(candidate)
    except json.JSONDecodeError as e:
        raise RuntimeError(f"LLM returned invalid JSON: {e}\n\nRAW OUTPUT:\n{raw}")


# ============================================================
# API: HEALTH
# ============================================================
@app.get("/health")
def health():
    return {"status": "ok"}


# ============================================================
# UI (HTML)
# ============================================================
@app.get("/", response_class=HTMLResponse)
def home():
    qdrant_url_html = html.escape(QDRANT_URL)
    coll_html = html.escape(QDRANT_COLLECTION)
    embed_html = html.escape(EMBED_MODEL)
    llm_html = html.escape(LLM_MODEL)

    return f"""

<!doctype html>

<html>

  <head>

    <meta charset="utf-8" />

    <title>PMS Copilot — RAG MVP</title>

    <style>

      body {{

        font-family: Arial, sans-serif;

        max-width: 1200px;

        margin: 34px auto;

        padding: 0 16px;

      }}

      h1 {{ margin: 0 0 14px 0; font-size: 44px; letter-spacing: -0.5px; }}

      .meta {{

        color:#666; font-size: 13px; margin: 8px 0 18px 0;

      }}

      .row {{ display:flex; gap:10px; margin: 14px 0; align-items: stretch; }}

      input {{

        flex:1; padding: 14px; font-size: 16px;

        border: 1px solid #bbb; border-radius: 6px;

      }}

      button {{

        padding: 14px 18px; font-size: 16px; cursor: pointer;

        border: 2px solid #222; background: #eee; border-radius: 6px;

        min-width: 88px;

      }}

      .panel {{

        background: #f6f6f6;

        border-radius: 12px;

        padding: 16px;

        margin-top: 14px;

        border: 1px solid #e2e2e2;

      }}

      .error {{

        background: #fdecec;

        border: 1px solid #f3b6b6;

      }}

      .title {{ font-size: 18px; font-weight: 700; margin: 0 0 10px 0; }}

      .sub {{ color:#333; margin: 0 0 10px 0; }}

      .kv {{ margin: 0; color:#111; }}

      .kv b {{ display:inline-block; min-width: 140px; }}

      .findings {{

        margin-top: 14px;

        display: grid;

        grid-template-columns: 1fr;

        gap: 10px;

      }}

      .card {{

        background: #fff;

        border-radius: 10px;

        border: 1px solid #e5e5e5;

        padding: 14px;

      }}

      .card h3 {{

        margin: 0 0 8px 0;

        font-size: 16px;

      }}

      .muted {{ color:#666; font-size: 13px; }}

      .evidence {{

        margin-top: 14px;

      }}

      table {{

        width: 100%;

        border-collapse: collapse;

        background: #fff;

        border-radius: 10px;

        overflow: hidden;

        border: 1px solid #e5e5e5;

      }}

      th, td {{

        padding: 10px;

        border-bottom: 1px solid #eee;

        font-size: 13px;

        vertical-align: top;

      }}

      th {{ text-align: left; background: #fafafa; }}

      .row2 {{

        display:flex; justify-content: space-between; align-items: center;

        gap: 12px; margin-top: 10px;

      }}

      pre {{

        margin: 0;

        white-space: pre-wrap;

        background: #111;

        color: #eee;

        padding: 12px;

        border-radius: 10px;

        overflow: auto;

        font-size: 12px;

      }}

      .right {{

        display:flex; gap: 10px; align-items: center;

      }}

      .checkbox {{

        display:flex; gap: 8px; align-items: center;

        font-size: 13px; color:#333;

      }}

    </style>

  </head>

  <body>

    <h1>PMS Copilot — RAG MVP</h1>

    <div class="meta">

      Qdrant: <b>{qdrant_url_html}</b> · Collection: <b>{coll_html}</b> ·

      Embed: <b>{embed_html}</b> · TopK: <b>{TOP_K}</b> · LLM: <b>{llm_html}</b>

    </div>



    <div class="row">

      <input id="q" placeholder="Введите запрос..." />

      <button onclick="send()">Ask</button>

    </div>



    <div id="result" class="panel" style="display:none;"></div>



    <script>

      function esc(s) {{

        return String(s ?? "").replaceAll("&", "&amp;").replaceAll("<","&lt;").replaceAll(">","&gt;");

      }}



      function renderAudit(audit) {{

        const summary = audit?.summary ?? "";

        const findings = Array.isArray(audit?.findings) ? audit.findings : [];

        const conclusion = audit?.conclusion ?? "";



        let html = '';

        html += `<div class="title">Summary</div>`;

        html += `<div class="sub">${{esc(summary)}}</div>`;



        html += `<div class="title" style="margin-top:14px;">Findings</div>`;

        if (!findings.length) {{

          html += `<div class="muted">No findings returned.</div>`;

        }} else {{

          html += `<div class="findings">`;

          for (const f of findings) {{

            const refs = Array.isArray(f?.evidence_refs) ? f.evidence_refs.join(", ") : "";

            html += `

              <div class="card">

                <h3>${{esc(f?.topic ?? "Finding")}}</h3>

                <p class="kv"><b>Requirement:</b> ${{esc(f?.requirement ?? "")}}</p>

                <p class="kv"><b>Observation:</b> ${{esc(f?.observation ?? "")}}</p>

                <p class="kv"><b>Risk:</b> ${{esc(f?.risk ?? "")}}</p>

                <p class="muted"><b>Evidence refs:</b> ${{esc(refs)}}</p>

              </div>

            `;

          }}

          html += `</div>`;

        }}



        html += `<div class="title" style="margin-top:14px;">Conclusion</div>`;

        html += `<div class="sub">${{esc(conclusion)}}</div>`;

        return html;

      }}



      function renderEvidenceTable(sources) {{

        if (!Array.isArray(sources) || !sources.length) return '';

        let rows = '';

        for (const s of sources) {{

          rows += `

            <tr>

              <td>${{esc(s.n)}}</td>

              <td>${{esc(s.id)}}</td>

              <td>${{esc(s.score)}}</td>

              <td>${{esc(s["GROUPS"])}}</td>

              <td>${{esc(s["FREQUENCY TYPE"])}}</td>

              <td>${{esc(s["RESPONSIBILITY"])}}</td>

            </tr>

          `;

        }}

        return `

          <div class="evidence">

            <div class="title">Evidence</div>

            <table>

              <thead>

                <tr>

                  <th>#</th>

                  <th>ID</th>

                  <th>Score</th>

                  <th>GROUPS</th>

                  <th>FREQUENCY</th>

                  <th>RESPONSIBILITY</th>

                </tr>

              </thead>

              <tbody>${{rows}}</tbody>

            </table>

          </div>

        `;

      }}



      async function send() {{

        const q = document.getElementById('q').value;

        const panel = document.getElementById('result');

        panel.style.display = 'block';

        panel.className = 'panel';

        panel.innerHTML = `<div class="title">Working...</div><div class="muted">Embedding → Qdrant → LLM</div>`;



        try {{

          const r = await fetch('/ask', {{

            method: 'POST',

            headers: {{ 'Content-Type': 'application/json' }},

            body: JSON.stringify({{ q }})

          }});



          const data = await r.json();



          if (!data.ok) {{

            panel.className = 'panel error';

            panel.innerHTML = `

              <div class="title">Error</div>

              <div class="sub">${{esc(data.error ?? "Request failed")}}</div>

              <pre>${{esc(JSON.stringify(data, null, 2))}}</pre>

            `;

            return;

          }}



          const audit = data.audit;

          const sources = data.sources;



          const auditHtml = renderAudit(audit);

          const evidenceHtml = renderEvidenceTable(sources);



          panel.innerHTML = `

            ${{auditHtml}}

            ${{evidenceHtml}}

            <div class="row2">

              <div class="checkbox">

                <input id="rawToggle" type="checkbox" onchange="toggleRaw()" />

                <label for="rawToggle">Show raw JSON</label>

              </div>

              <div class="right muted">TopK: ${{esc(data.debug?.top_k)}}</div>

            </div>

            <div id="rawBlock" style="display:none; margin-top:10px;">

              <pre>${{esc(JSON.stringify(data, null, 2))}}</pre>

            </div>

          `;



        }} catch (e) {{

          panel.className = 'panel error';

          panel.innerHTML = `<div class="title">Error</div><pre>${{esc(String(e))}}</pre>`;

        }}

      }}



      function toggleRaw() {{

        const cb = document.getElementById('rawToggle');

        const block = document.getElementById('rawBlock');

        if (!cb || !block) return;

        block.style.display = cb.checked ? 'block' : 'none';

      }}

    </script>

  </body>

</html>

""".strip()


# ============================================================
# API
# ============================================================
@app.post("/ask")
def ask(req: AskRequest):
    q = (req.q or "").strip()
    if not q:
        return JSONResponse({"ok": False, "error": "Empty query"}, status_code=400)

    if len(q) < MIN_QUERY_CHARS:
        return JSONResponse(
            {"ok": False, "error": f"Query too short (min {MIN_QUERY_CHARS} chars)"},
            status_code=400,
        )

    if len(q) > MAX_QUERY_CHARS:
        q = q[:MAX_QUERY_CHARS]

    # 1) Embedding
    try:
        query_vec = embed(q)
    except Exception as e:
        return JSONResponse(
            {
                "ok": False,
                "error": "Embedding failed",
                "details": str(e),
                "debug": {
                    "embed_model": EMBED_MODEL,
                    "vector_size": VECTOR_SIZE,
                },
            },
            status_code=500,
        )

    # 2) Qdrant search (REST)
    try:
        raw_points = qdrant_search_rest(query_vec, TOP_K)
    except Exception as e:
        return JSONResponse(
            {
                "ok": False,
                "error": "Qdrant search failed",
                "details": str(e),
                "debug": {
                    "qdrant_url": QDRANT_URL,
                    "collection": QDRANT_COLLECTION,
                    "qdrant_api_key_set": bool(QDRANT_API_KEY),
                },
            },
            status_code=500,
        )

    # Normalize hits for downstream
    hits: List[Dict[str, Any]] = []
    for p in raw_points:
        hits.append(
            {
                "id": p.get("id"),
                "score": p.get("score"),
                "payload": p.get("payload") or {},
            }
        )

    # Evidence gate
    strong_hits = sum(1 for h in hits if (h.get("score") or 0) >= SCORE_THRESHOLD)
    evidence_text, sources = build_evidence_blocks(hits)

    if strong_hits < MIN_STRONG_HITS:
        return {
            "ok": True,
            "query": q,
            "audit": {
                "summary": "Insufficient evidence found in PMS data for a grounded audit answer.",
                "findings": [
                    {
                        "topic": "Evidence gating",
                        "requirement": f"At least {MIN_STRONG_HITS} hits with score >= {SCORE_THRESHOLD}",
                        "observation": f"Only {strong_hits} strong hits were retrieved.",
                        "risk": "Answer may be speculative without sufficient PMS evidence.",
                        "evidence_refs": [],
                    }
                ],
                "conclusion": "Please refine the question or ensure the relevant PMS/manual records exist in the collection.",
            },
            "sources": sources,
            "hits": hits if RETURN_RAW_HITS else [],
            "debug": {
                "qdrant_url": QDRANT_URL,
                "collection": QDRANT_COLLECTION,
                "top_k": TOP_K,
                "embed_model": EMBED_MODEL,
                "vector_size": VECTOR_SIZE,
                "llm_model": LLM_MODEL,
                "strong_hits": strong_hits,
                "score_threshold": SCORE_THRESHOLD,
                "min_strong_hits": MIN_STRONG_HITS,
                "llm_called": False,
            },
        }

    # 3) LLM audit JSON (strict)
    try:
        audit = run_llm_audit_json(q, evidence_text)
    except Exception as e:
        return JSONResponse(
            {
                "ok": False,
                "error": "LLM failed",
                "details": str(e),
                "debug": {"llm_model": LLM_MODEL},
                "sources": sources,
                "hits": hits if RETURN_RAW_HITS else [],
            },
            status_code=500,
        )

    return {
        "ok": True,
        "query": q,
        "audit": audit,       # STRICT JSON (parsed)
        "sources": sources,   # compact evidence table for UI
        "hits": hits if RETURN_RAW_HITS else [],
        "debug": {
            "qdrant_url": QDRANT_URL,
            "collection": QDRANT_COLLECTION,
            "top_k": TOP_K,
            "embed_model": EMBED_MODEL,
            "vector_size": VECTOR_SIZE,
            "llm_model": LLM_MODEL,
            "strong_hits": strong_hits,
            "score_threshold": SCORE_THRESHOLD,
            "min_strong_hits": MIN_STRONG_HITS,
            "llm_called": True,
        },
    }