Conflict_Checker/backend/jobs.py

"""
jobs.py - Lightweight async job registry for conflict checks.

A conflict check takes minutes, so the HTTP request must not block on it. Each
upload becomes a job that runs on a background thread; the client gets a job_id
immediately and can either poll GET /jobs/{id} or just close the page and wait
for the completion email.

State is in-memory (fine for a single-user tool); the report is also persisted
to outputs/<job_id>/ so results survive a restart even though live status does
not. No external queue/DB.

A teed stdout/stderr log is kept in memory and written to outputs/<job_id>/job.log
so failed or suspicious runs can be reviewed after the fact.
"""

import json
import os
import time
import traceback
import uuid
import shutil
import threading
from contextlib import contextmanager
from typing import Dict, Iterator, List, Optional

from backend import config
from backend import llm
from backend.job_log import capture_stdio, read_log_file, stamp_line
from backend.agents.runner import run_agent_pipeline
from backend.pipeline.runner import run_pipeline
from backend.email_sender import send_conflict_report, send_review_required

_jobs: Dict[str, Dict] = {}
_lock = threading.Lock()
_LOG_TAIL = 80
PIPELINE_MODES = {"classic", "agent"}
# States where the job will produce no more log output; polls get the full log.
_TERMINAL_STATES = {"done", "error", "needs_review", "finalization_error"}


def _set(job_id: str, **fields) -> None:
    with _lock:
        _jobs[job_id].update(fields)


def _append_log(job_id: str, raw_line: str, log_path: str) -> None:
    """Stamp, store, and append one captured stdout/stderr line."""
    entry = stamp_line(raw_line)
    with _lock:
        job = _jobs.get(job_id)
        if job is not None:
            job.setdefault("log", []).append(entry)
    try:
        os.makedirs(os.path.dirname(log_path), exist_ok=True)
        with open(log_path, "a", encoding="utf-8") as f:
            f.write(entry + "\n")
    except OSError:
        # Don't fail the job over log I/O; avoid print() here — it would
        # re-enter the stdio tee while a job is capturing.
        pass


def create_job(
    pdf_path: str,
    source_filename: str,
    email: Optional[str] = None,
    project_input: Optional[Dict] = None,
    text_local: bool = False,
    pipeline_mode: str = "classic",
    vision_model: Optional[str] = None,
    text_model: Optional[str] = None,
) -> str:
    """Register a job and kick off its background thread. Returns the job_id."""
    pipeline_mode = pipeline_mode.strip().lower()
    if pipeline_mode not in PIPELINE_MODES:
        raise ValueError(f"Unsupported pipeline mode: {pipeline_mode!r}")
    # Agent mode v1 is OpenRouter-only.
    text_local = bool(text_local and pipeline_mode == "classic")
    job_id = uuid.uuid4().hex[:12]
    with _lock:
        _jobs[job_id] = {
            "job_id": job_id,
            "status": "queued",          # queued -> running -> done | needs_review | error
            "source": source_filename,
            "email": email or None,
            "project_input": project_input or {},
            "text_local": text_local,
            "pipeline_mode": pipeline_mode,
            "vision_model": (vision_model or "").strip() or None,
            "text_model": (text_model or "").strip() or None,
            "stage": None,
            "created_at": time.time(),
            "finished_at": None,
            "report": None,
            "error": None,
            "log": [],
        }
    threading.Thread(
        target=_run,
        args=(job_id, pdf_path, project_input, text_local, pipeline_mode,
              vision_model, text_model),
        daemon=True,
    ).start()
    return job_id


def _run(
    job_id: str,
    pdf_path: str,
    project_input: Optional[Dict] = None,
    text_local: bool = False,
    pipeline_mode: str = "classic",
    vision_model: Optional[str] = None,
    text_model: Optional[str] = None,
) -> None:
    out_dir = os.path.join(config.OUTPUT_DIR, job_id)
    log_path = os.path.join(out_dir, "job.log")
    try:
        _set(job_id, status="running")
        os.makedirs(out_dir, exist_ok=True)
        # Truncate any leftover log if job_id somehow collided (shouldn't).
        with open(log_path, "w", encoding="utf-8"):
            pass
        header = (f"=== Job {job_id} | {pipeline_mode} | {_jobs[job_id].get('source')} | "
                  f"vision={vision_model or 'default'} text={text_model or 'default'} | "
                  f"started {time.strftime('%Y-%m-%d %H:%M:%S %Z', time.gmtime())} UTC ===")
        _append_log(job_id, header, log_path)

        def on_line(raw: str) -> None:
            _append_log(job_id, raw, log_path)

        with capture_stdio(on_line):
            _run_pipeline(job_id, pdf_path, out_dir, project_input, text_local,
                          pipeline_mode, vision_model, text_model)
    except Exception as e:
        # Land the failure AND its traceback in the job log so failed runs can
        # be diagnosed from the log alone (the stdio tee is already torn down).
        try:
            _append_log(job_id, f"[Jobs] Job {job_id} failed: {e}", log_path)
            for ln in traceback.format_exc().rstrip().splitlines():
                _append_log(job_id, ln, log_path)
            # Cost-so-far for the failed run (counters are reset per job).
            cost = llm.get_cost()
            _append_log(job_id,
                        f"[Jobs] Estimated LLM cost before failure: "
                        f"${cost['usd']:.4f} over {cost['calls']} live calls "
                        f"({cost.get('cached', 0)} cached)", log_path)
        except Exception:
            pass
        print(f"[Jobs] Job {job_id} failed: {e}")
        _set(job_id, status="error", error=str(e), finished_at=time.time())
        _notify_error(job_id)
    finally:
        try:
            os.remove(pdf_path)
        except OSError:
            pass


def _run_pipeline(job_id: str, pdf_path: str, out_dir: str,
                  project_input: Optional[Dict], text_local: bool,
                  pipeline_mode: str, vision_model: Optional[str],
                  text_model: Optional[str]) -> None:
    """The body of a job run; executes inside the job's tee'd log capture."""
    # Persist minimal job metadata so the disk fallback in get_job can
    # recover the recipient email / pipeline mode after a server restart
    # (plain json.dump, matching the _dump style used elsewhere).
    with open(os.path.join(out_dir, "job.json"), "w", encoding="utf-8") as f:
        json.dump({
            "job_id": job_id,
            "email": _jobs[job_id].get("email"),
            "pipeline_mode": pipeline_mode,
            "source": _jobs[job_id].get("source"),
            "vision_model": vision_model,
            "text_model": text_model,
        }, f, indent=2)
    # Keep a copy of the source PDF so its sheets can be viewed later.
    shutil.copy2(pdf_path, os.path.join(out_dir, "source.pdf"))
    # Raw per-call LLM request/response dumps land in <out_dir>/llm_raw/.
    if config.LLM_RAW_DUMP:
        llm.set_raw_dump_dir(os.path.join(out_dir, "llm_raw"))
    runner = run_agent_pipeline if pipeline_mode == "agent" else run_pipeline
    runner_kwargs = {
        "out_dir": out_dir,
        "on_stage": lambda name: _set(job_id, stage=name),
        "project_input": project_input,
        "source_name": _jobs[job_id].get("source"),
    }
    if pipeline_mode == "classic":
        # run_pipeline takes the picks as params and clears them in finally.
        runner_kwargs["text_local"] = text_local
        runner_kwargs["vision_model"] = vision_model
        runner_kwargs["text_model"] = text_model
    else:
        runner_kwargs["require_review"] = config.AGENT_REQUIRE_REVIEW
        # The agent runner has no override params; set them module-level.
        if vision_model or text_model:
            print(f"[Jobs] Model overrides for this run: "
                  f"vision={vision_model or '(default)'} text={text_model or '(default)'}")
            llm.set_model_overrides(vision_model, text_model)
    try:
        report = runner(pdf_path, **runner_kwargs)
    finally:
        llm.set_raw_dump_dir(None)
        if pipeline_mode == "agent":
            llm.set_model_overrides(None, None)
    report.setdefault("summary", {})["pipeline_mode"] = pipeline_mode
    _log_cost_summary(report.get("summary", {}))
    if report["summary"].get("agent_status") == "needs_review":
        # Human-review gate: hold the job, don't email the unreviewed report.
        _set(job_id, status="needs_review", report=report,
             finished_at=time.time(), stage=None)
        email = _jobs[job_id].get("email")
        if email:
            review_url = f"{config.APP_BASE_URL.rstrip('/')}/?job={job_id}"
            send_review_required(email, report, review_url)
    else:
        _set(job_id, status="done", report=report, finished_at=time.time(), stage=None)
        _notify(job_id, report, out_dir)


def _log_cost_summary(summary: Dict, label: str = "this run") -> None:
    """
    End-of-log estimated LLM cost block, printed inside the job's stdio tee so
    it lands at the tail of job.log. Both runners populate the same summary
    fields (cost_usd / llm_calls / cached_calls / cost_by_stage / models_used).
    Hybrid note: local text calls carry no usage accounting, so the dollar
    figure covers OpenRouter calls only (local call counts still appear).
    """
    if summary.get("cost_usd") is None and not summary.get("llm_calls"):
        return
    print(f"\n=== Estimated LLM cost ({label}) ===")
    print(f"  Total: ${summary.get('cost_usd', 0.0):.4f} across "
          f"{summary.get('llm_calls', 0)} live calls "
          f"({summary.get('cached_calls', 0)} cached at $0)")
    for name, bucket in (summary.get("cost_by_stage") or {}).items():
        print(f"    {name}: ${bucket.get('usd', 0.0):.4f} "
              f"({bucket.get('calls', 0)} live, {bucket.get('cached', 0)} cached)")
    mu = summary.get("models_used") or {}
    if mu.get("vision"):
        print(f"  Vision models: {', '.join(mu['vision'])}")
    if mu.get("text_local"):
        print(f"  Text models (local): {', '.join(mu['text_local'])}")
    if mu.get("text_cloud"):
        print(f"  Text models (cloud): {', '.join(mu['text_cloud'])}")
    if mu.get("fallback_count"):
        print(f"  Local->cloud fallbacks: {mu['fallback_count']}")
    if mu.get("text_local"):
        print("  Note: local calls have no cost accounting; "
              "the dollar total covers OpenRouter usage only.")


@contextmanager
def capture_job_output(job_id: str, out_dir: str) -> Iterator[None]:
    """
    Re-open the stdio tee + raw LLM dump dir for post-run work that still
    belongs to this job (review finalization): lines append to job.log and
    the in-memory log, raw dumps resume under <out_dir>/llm_raw/. The tee is
    process-global — same overlapping-job caveat as the main run.
    """
    log_path = os.path.join(out_dir, "job.log")
    if config.LLM_RAW_DUMP:
        llm.set_raw_dump_dir(os.path.join(out_dir, "llm_raw"))
    try:
        with capture_stdio(lambda raw: _append_log(job_id, raw, log_path)):
            yield
    finally:
        llm.set_raw_dump_dir(None)


def _notify(job_id: str, report: Dict, out_dir: str) -> None:
    email = _jobs[job_id].get("email")
    if not email:
        return
    results_url = f"{config.APP_BASE_URL.rstrip('/')}/?job={job_id}"
    attachments = [
        os.path.join(out_dir, "report.md"),
        os.path.join(out_dir, "conflicts.json"),
        os.path.join(out_dir, "validated_issues.json"),
        os.path.join(out_dir, "rfis.json"),
    ]
    send_conflict_report(email, report, results_url=results_url, attachments=attachments)


def _notify_error(job_id: str) -> None:
    job = _jobs[job_id]
    email = job.get("email")
    if not email:
        return
    try:
        from backend.email_sender import _smtp_ready, _send
        from email.message import EmailMessage
        if not _smtp_ready():
            print(f"[Email] SMTP not configured - skipping error notice to {email}")
            return
        msg = EmailMessage()
        msg["Subject"] = f"Conflict Checker - {job.get('source','')} - run FAILED"
        msg["From"] = config.SMTP_FROM or config.SMTP_USER
        msg["To"] = email
        msg.set_content(
            "Your conflict check did not complete.\n\n"
            f"Drawing set: {job.get('source','')}\n"
            f"Error: {job.get('error','unknown')}\n\n"
            f"Review the run log at: {config.APP_BASE_URL.rstrip('/')}/?job={job_id}\n\n"
            "Generated by Conflict Checker"
        )
        _send(msg)
    except Exception as e:
        print(f"[Email] Failed to send error notice: {e}")


def _log_from_disk(job_id: str) -> List[str]:
    return read_log_file(os.path.join(config.OUTPUT_DIR, job_id, "job.log"))


def get_job_log(job_id: str) -> Optional[List[str]]:
    """Full job log lines, from memory or disk. None if job unknown."""
    with _lock:
        job = _jobs.get(job_id)
        if job is not None:
            return list(job.get("log") or [])

    log = _log_from_disk(job_id)
    # Job exists on disk if we have a log or a report artifact.
    report_path = os.path.join(config.OUTPUT_DIR, job_id, "conflicts.json")
    if log or os.path.isfile(report_path):
        return log
    return None


def get_job(job_id: str) -> Optional[Dict]:
    """Public job view. Includes the full report only when done.

    Falls back to the on-disk artifacts (conflicts.json / job.log) when the
    job isn't in the in-memory registry (e.g. after a server restart).
    """
    with _lock:
        job = _jobs.get(job_id)
        if job:
            out = dict(job)
            log = list(job.get("log") or [])
            out["log_tail"] = log[-_LOG_TAIL:]
            # Full log on terminal states so the UI can show it without a
            # second fetch; keep polls light while running.
            if out.get("status") in _TERMINAL_STATES:
                out["log"] = log
            else:
                out.pop("log", None)
            return out

    # Try loading from disk
    report_path = os.path.join(config.OUTPUT_DIR, job_id, "conflicts.json")
    log = _log_from_disk(job_id)
    if not os.path.isfile(report_path) and not log:
        return None
    try:
        report = None
        if os.path.isfile(report_path):
            with open(report_path, encoding="utf-8") as f:
                report = json.load(f)
        summary = (report or {}).get("summary", {})
        if report is None:
            # Crashed before writing a report; the log is the only artifact.
            status = "error"
        else:
            # Recover the job's real state: a job that stopped at the review gate
            # must come back as needs_review (not done) or it can never finalize.
            status = "needs_review" if summary.get("agent_status") == "needs_review" else "done"
        # job.json (written at job start) carries the recipient email and
        # pipeline mode so the final notification still fires after a restart.
        # Missing/corrupt job.json degrades to the previous derivations.
        meta: Dict = {}
        meta_path = os.path.join(config.OUTPUT_DIR, job_id, "job.json")
        try:
            with open(meta_path, encoding="utf-8") as f:
                loaded = json.load(f)
            if isinstance(loaded, dict):
                meta = loaded
        except (OSError, json.JSONDecodeError):
            pass
        source_pdf = os.path.join(config.OUTPUT_DIR, job_id, "source.pdf")
        job = {
            "job_id":       job_id,
            "status":       status,
            "source":       meta.get("source") or (report or {}).get("source", os.path.basename(report_path)),
            "email":        meta.get("email"),
            "project_input": (report or {}).get("project_input", {}),
            "text_local":   summary.get("text_backend") == "local",
            "pipeline_mode": meta.get("pipeline_mode") or summary.get("pipeline_mode", "classic"),
            "vision_model": meta.get("vision_model"),
            "text_model":   meta.get("text_model"),
            "stage":        None,
            "created_at":   os.path.getmtime(source_pdf) if os.path.isfile(source_pdf) else None,
            "finished_at":  os.path.getmtime(report_path) if os.path.isfile(report_path) else None,
            "report":       report,
            "error":        None if report is not None else "Report missing; see job log",
            "log":          log,
            "log_tail":     log[-_LOG_TAIL:],
        }
        # Hydrate the in-memory registry so _set(...) transitions (reviewing,
        # finalizing, done) work for restart-recovered jobs.
        with _lock:
            return dict(_jobs.setdefault(job_id, job))
    except Exception as e:
        print(f"[Jobs] Failed to load job {job_id} from disk: {e}")
        return None