"""Local bundled model tier (V-1): Gemma 4 E4B via llama.cpp, on by default.

The model runs as a local ``llama-server`` exposing an OpenAI-compatible endpoint on
127.0.0.1. This module locates and, if needed, launches that server, points an
OpenAI-compatible client at it, and reads a document's full prose for the two things
the deterministic engine cannot do alone: person names that sit in no frame, and an
unmarked-privilege judgment. Every network call this module makes is to 127.0.0.1 and
nowhere else (V-1: localhost only).

Graceful degrade (V-1): if the bundled binary or the model file is absent, or the
server will not start or answer, the tier records a note and returns; the deterministic
engine result still stands and the run never crashes.

The endpoint client is the v1 R-12 shape (OpenAI chat-completions); v2 points it at the
bundled local server and turns the tier on by default. No weights ship in the package;
the model file is resolved from a cache (auto-downloaded on first run, V-8) or env.
"""
from __future__ import annotations

import json
import os
import re
import subprocess
import time
import urllib.request
from pathlib import Path
from typing import Callable, Optional
from urllib.parse import urlparse

from .models import Finding, Report, Severity

# Distinct from the web UI's user-facing port (8799) so the bundled launcher can run
# the web face on 8799 and the local model server here without a port collision.
DEFAULT_PORT = 8791
_PKG_ROOT = Path(__file__).resolve().parent

INSTRUCTION = (
    "You are a privacy pre-flight scanner for legal documents. Read the user's text and "
    "return ONLY a JSON object, no prose and no code fences, with exactly these keys:\n"
    '{"persons": [full personal names of real people mentioned, as written; do NOT '
    "include organizations, companies, courts, agencies, places, statutes, months, or "
    'weekdays], "privilege": {"label": one of "privileged","work_product",'
    '"not_privileged","review", "confidence": a number from 0 to 1, '
    '"reason": "at most 15 words"}}\n'
    "privileged = attorney-client communication conveying or seeking legal advice. "
    "work_product = the lawyer's own analysis, strategy, or mental impressions. "
    "not_privileged = public record, filed document, news, statute, or general "
    "information. review = genuinely unsure. Return strictly valid JSON."
)
_LABEL_SEVERITY = {
    "privileged": Severity.HIGH,
    "work_product": Severity.HIGH,
    "review": Severity.MEDIUM,
}

# --- page limits (V-9): the model pass is the slow part on a CPU laptop ---
PAGE_WORDS = 500          # one "page" of dense single-spaced text
SOFT_PAGES = 30           # below this runs silently; at/above warns and shows progress
DEFAULT_HARD_PAGES = 150  # near here the model pass is skipped unless overridden


def estimate_pages(text: str) -> int:
    import math
    return max(1, math.ceil(len(text.split()) / PAGE_WORDS))


def hard_page_cap() -> int:
    v = os.environ.get("PRIVSWEEP_MAX_PAGES")
    if v:
        try:
            return int(v)
        except ValueError:
            pass
    return DEFAULT_HARD_PAGES


def page_ok_for_model(text: str, report: Report, label: str) -> bool:
    """Gate the model pass by size (V-9). Notes a soft warning; skips past the hard cap."""
    pages = estimate_pages(text)
    cap = hard_page_cap()
    if pages > cap:
        report.note(f"{label}: ~{pages} pages exceeds the {cap}-page model cap; model pass "
                    f"skipped (set PRIVSWEEP_MAX_PAGES or --max-pages to override).")
        return False
    if pages >= SOFT_PAGES:
        report.note(f"{label}: large document (~{pages} pages); the local model pass may "
                    f"take a few minutes.")
    return True


# --- resolution of the bundled runtime + model file --------------------------

def _runtime_dir() -> Optional[Path]:
    env = os.environ.get("PRIVSWEEP_RUNTIME_DIR")
    if env:
        p = Path(env)
        return p if p.exists() else None
    # Shipped layout: runtime/ sits beside source/; source/privsweep/ is this package.
    guess = _PKG_ROOT.parent.parent / "runtime"
    return guess if guess.exists() else None


def llama_server_path() -> Optional[Path]:
    env = os.environ.get("PRIVSWEEP_LLAMA_SERVER")
    if env:
        p = Path(env)
        return p if p.exists() else None
    rt = _runtime_dir()
    if not rt:
        return None
    for sub in (rt, rt / "bin", rt / "llamacpp", rt / "llamacpp" / "bin"):
        for name in ("llama-server.exe", "llama-server"):
            cand = sub / name
            if cand.exists():
                return cand
    return None


def model_gguf_path() -> Optional[Path]:
    env = os.environ.get("PRIVSWEEP_MODEL_GGUF")
    if env:
        p = Path(env)
        return p if p.exists() else None
    dirs: list[Path] = []
    cache = os.environ.get("PRIVSWEEP_MODEL_DIR")
    if cache:
        dirs.append(Path(cache))
    rt = _runtime_dir()
    if rt:
        dirs.extend([rt.parent / "model cache", rt.parent / "models"])
    for d in dirs:
        if not d.exists():
            continue
        ggufs = sorted(d.glob("*.gguf"))
        # Prefer a text model over a vision projector if both are present.
        text_only = [g for g in ggufs if "mmproj" not in g.name.lower()]
        ggufs = text_only or ggufs
        if ggufs:
            return ggufs[0]
    # Fall back to the first-run download cache (V-8).
    try:
        from .model_download import cached_model_path
        cached = cached_model_path()
        if cached.exists():
            return cached
    except Exception:  # noqa: BLE001
        pass
    return None


def model_endpoint() -> str:
    env = os.environ.get("PRIVSWEEP_MODEL_ENDPOINT")
    if env:
        return env
    port = os.environ.get("PRIVSWEEP_MODEL_PORT", str(DEFAULT_PORT))
    return f"http://127.0.0.1:{port}/v1/chat/completions"


def _base_of(endpoint: str) -> str:
    return endpoint.split("/v1", 1)[0]


def _health_ok(endpoint: str, timeout: float = 3.0) -> bool:
    try:
        with urllib.request.urlopen(_base_of(endpoint) + "/health",
                                    timeout=timeout) as resp:
            return resp.status == 200
    except Exception:  # noqa: BLE001 - any failure means "not ready"
        return False


# --- server lifecycle --------------------------------------------------------

def ensure_server(launch: bool = True):
    """Resolve a reachable local endpoint.

    Returns ``(endpoint|None, note|None, proc|None)``. ``proc`` is set only when this
    call started the server (so the caller can stop it). Never launches over an
    explicitly configured endpoint.
    """
    if os.environ.get("PRIVSWEEP_MODEL_ENDPOINT"):
        ep = model_endpoint()
        if _health_ok(ep):
            return ep, None, None
        return None, f"configured model endpoint not reachable ({ep}).", None

    ep = model_endpoint()
    if _health_ok(ep):
        return ep, None, None  # a local server is already running.

    if not launch:
        return None, "local model not running.", None

    binary = llama_server_path()
    gguf = model_gguf_path()
    if not binary or not gguf:
        return None, ("local model not found; ran deterministic checks only "
                      "(no names-in-prose or unmarked-privilege pass)."), None

    proc = _launch(binary, gguf, urlparse(ep).port or DEFAULT_PORT)
    deadline = time.time() + float(os.environ.get("PRIVSWEEP_MODEL_START_TIMEOUT", "90"))
    while time.time() < deadline:
        if proc.poll() is not None:
            return None, "local model server exited during startup.", None
        if _health_ok(ep, timeout=2.0):
            return ep, None, proc
        time.sleep(1.0)
    stop_server(proc)
    return None, "local model server did not become ready in time.", None


def _launch(binary: Path, gguf: Path, port: int) -> "subprocess.Popen":
    threads = os.environ.get("PRIVSWEEP_MODEL_THREADS") or str(min(8, os.cpu_count() or 4))
    ctx = os.environ.get("PRIVSWEEP_MODEL_CTX", "4096")
    cmd = [str(binary), "-m", str(gguf), "-c", str(ctx), "-t", str(threads),
           "--jinja", "--host", "127.0.0.1", "--port", str(port)]
    env = dict(os.environ)
    env["LD_LIBRARY_PATH"] = str(binary.parent) + os.pathsep + env.get("LD_LIBRARY_PATH", "")
    return subprocess.Popen(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
                            env=env)


def stop_server(proc) -> None:
    if proc is None:
        return
    try:
        proc.terminate()
        proc.wait(timeout=10)
    except Exception:  # noqa: BLE001
        try:
            proc.kill()
        except Exception:  # noqa: BLE001
            pass


# --- the model pass ----------------------------------------------------------

def classify(text: str, endpoint: str, timeout: float = 240.0) -> dict:
    """Call the local OpenAI-compatible endpoint (127.0.0.1) and parse its JSON reply."""
    payload = {
        "model": "privsweep-e4b",
        "messages": [
            {"role": "system", "content": INSTRUCTION},
            {"role": "user", "content": text},
        ],
        "temperature": 0,
        "max_tokens": 512,
    }
    req = urllib.request.Request(endpoint, data=json.dumps(payload).encode("utf-8"),
                                 headers={"Content-Type": "application/json"},
                                 method="POST")
    with urllib.request.urlopen(req, timeout=timeout) as resp:  # localhost only
        body = json.load(resp)
    content = body["choices"][0]["message"]["content"]
    return _extract_json(content)


def _extract_json(content: str) -> dict:
    content = re.sub(r"```$", "", re.sub(r"^```(?:json)?", "", content.strip()).strip()).strip()
    span = content[content.find("{"): content.rfind("}") + 1] if "{" in content else ""
    for cand in (content, span):
        try:
            obj = json.loads(cand)
            if isinstance(obj, dict):
                return obj
        except Exception:  # noqa: BLE001
            continue
    return {}


def _overlaps(span, existing) -> bool:
    s, e = span
    return any(not (e <= es[0] or s >= es[1]) for es in existing)


def model_findings(text: str, result: dict, file: str,
                   locator: Callable[[int], str], existing_spans) -> list[Finding]:
    """Turn one model reply into Findings: residual prose names + a privilege judgment."""
    out: list[Finding] = []
    persons = result.get("persons") or []
    if not isinstance(persons, list):
        persons = [str(persons)]
    seen: set[str] = set()
    for raw in persons:
        name = str(raw).strip()
        if not name or name.lower() in seen:
            continue
        seen.add(name.lower())
        occurrences = [(m.start(), m.end()) for m in re.finditer(re.escape(name), text)]
        if occurrences:
            free = [o for o in occurrences if not _overlaps(o, existing_spans)]
            if not free:
                continue  # every occurrence is already anchored by a frame; frames own it
            span = free[0]
            loc = locator(span[0])
        else:
            span, loc = None, "model"
        out.append(Finding(
            detector="model.e4b", type="person_name", severity=Severity.HIGH,
            file=file, location=loc, span=span, preview=name, role="unknown",
            confidence=None, detail={"source": "model"},
        ))
    priv = result.get("privilege") or {}
    label = str(priv.get("label", "")).lower().strip()
    if label in _LABEL_SEVERITY:
        out.append(Finding(
            detector="model.privilege", type="privilege_classification",
            severity=_LABEL_SEVERITY[label], file=file, location="document",
            preview=f"{label}: {priv.get('reason', '')}"[:120],
            confidence=priv.get("confidence"),
            detail={"label": label, "reason": priv.get("reason"), "source": "model"},
        ))
    return out


def augment_with_model(paths, report: Report, endpoint: Optional[str] = None,
                       *, launch: bool = True) -> None:
    """On-by-default model pass. Degrades to a note if the model is unavailable (V-1)."""
    proc = None
    started = False
    if endpoint is None:
        endpoint, note, proc = ensure_server(launch=launch)
        if endpoint is None:
            report.note(note or "local model unavailable; ran deterministic checks only.")
            return
        started = proc is not None
    elif not _health_ok(endpoint):
        report.note(f"model endpoint not reachable ({endpoint}); deterministic results stand.")
        return

    try:
        from .engine import (_effective_kind, _normalize_paths, locator_for_offset,
                             walk_paths)
        from .extract import extract

        existing_by_file: dict[str, list] = {}
        for f in report.findings:
            if f.span:
                existing_by_file.setdefault(f.file, []).append(f.span)

        for fp in walk_paths(_normalize_paths(paths)):
            if _effective_kind(fp, None) in (None, "image", "eml"):
                continue
            try:
                doc = extract(fp)
            except Exception:  # noqa: BLE001 - one bad file must not abort the pass
                continue
            if not page_ok_for_model(doc.full_text, report, fp):
                continue
            try:
                result = classify(doc.full_text, endpoint)
            except Exception as exc:  # noqa: BLE001 - endpoint died mid-run -> warn, stop
                report.note(f"local model call failed ({type(exc).__name__}); "
                            f"deterministic results stand.")
                break
            report.add(*model_findings(
                doc.full_text, result, fp,
                lambda off, d=doc: locator_for_offset(d, off),
                existing_by_file.get(fp, [])))
    finally:
        if started:
            stop_server(proc)
