"""Score a candidate. The data and all scoring stay in this (trusted) process.

The candidate runs in a separate confined process (see host.py) and is handed the
stream one byte at a time, each only after it has committed to a prediction.
"""
from __future__ import annotations

import json
import os
import struct
import subprocess
import sys
import tempfile
import threading
import time
from array import array
from pathlib import Path

import numpy as np

from compsearch.runner import (ALPHABET, CandidateError, byte_cost_bits, code_cost_bits,
                               curriculum_parts)
from compsearch.safety import check_source

PROJECT_ROOT = Path(__file__).resolve().parent.parent
HOST = Path(__file__).resolve().parent / "host.py"
FRAME_BYTES = ALPHABET * 8
STARTUP_SLACK_SECONDS = 10
MAX_ERROR_BYTES = 4000


def evaluate(candidate: str | Path, data: str | Path,
             seconds: float = 300.0, mem_mb: int = 1024) -> dict:
    """Score a candidate file. Always returns a dict with a 'status' key."""
    source = Path(candidate).read_text(encoding="utf-8")
    problem = check_source(source)
    if problem:
        return {"status": "rejected", "error": problem}
    stream = Path(data).read_bytes()
    if not stream:
        return {"status": "error", "error": "the data stream is empty"}

    env = dict(os.environ)
    for name in ("OMP_NUM_THREADS", "OPENBLAS_NUM_THREADS", "MKL_NUM_THREADS",
                 "NUMEXPR_NUM_THREADS"):
        env[name] = "1"  # one core per candidate, so parallel runs stay fair
    env["PYTHONHASHSEED"] = "0"

    timed_out = threading.Event()
    with tempfile.TemporaryDirectory() as empty_dir:  # candidate's working folder has nothing in it
        process = subprocess.Popen(
            [sys.executable, str(HOST), str(mem_mb)], cwd=empty_dir, env=env,
            stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.DEVNULL)

        def stop() -> None:
            timed_out.set()
            process.kill()

        watchdog = threading.Timer(seconds + STARTUP_SLACK_SECONDS, stop)
        watchdog.start()
        started = time.monotonic()
        try:
            result = _converse(process, source, stream)
        finally:
            watchdog.cancel()
            process.kill()
            process.wait()
            for pipe in (process.stdin, process.stdout):
                try:
                    pipe.close()
                except OSError:
                    pass
        elapsed = time.monotonic() - started

    if timed_out.is_set() or (result["status"] == "ok" and elapsed > seconds):
        return {"status": "timeout", "error": f"not finished after {seconds:.0f} seconds"}
    if result["status"] != "ok":
        return result
    code_bits = code_cost_bits(source)
    byte_bits = result["byte_bits"]
    data_bpb = float(byte_bits.sum()) / len(stream)
    scored = {
        "status": "ok",
        "bpb": data_bpb + code_bits / len(stream),
        "data_bpb": data_bpb,
        "code_bytes": code_bits // 8,
        "n_bytes": len(stream),
        "seconds": round(elapsed, 2),
    }
    curriculum = load_curriculum(data, len(stream))
    if curriculum:
        parts = curriculum_parts(byte_bits, curriculum)
        scored["parts"] = {name: round(value, 4) for name, value in parts.items()}
        scored["bpb"] = sum(parts.values()) / len(parts) + code_bits / len(stream)
    return scored


def load_curriculum(data: str | Path, stream_length: int) -> dict | None:
    """A stream may come with a .json file beside it that marks its segments."""
    path = Path(data).with_suffix(".json")
    if not path.is_file():
        return None
    curriculum = json.loads(path.read_text(encoding="utf-8"))
    for segment in curriculum["segments"]:
        if not 0 <= segment["start"] < segment["end"] <= stream_length:
            raise ValueError(f"{path}: segment {segment} does not fit the stream")
    return curriculum


def _converse(process: subprocess.Popen, source: str, stream: bytes) -> dict:
    died = {"status": "crashed", "error": "the candidate process ended early"}
    encoded = source.encode("utf-8")
    try:
        process.stdin.write(struct.pack("<I", len(encoded)) + encoded)
        process.stdin.flush()
        byte_bits = array("d")
        for actual in stream:
            tag = process.stdout.read(1)
            if tag == b"E":
                return _read_error(process)
            frame = process.stdout.read(FRAME_BYTES) if tag == b"P" else b""
            if len(frame) != FRAME_BYTES:
                return died
            byte_bits.append(byte_cost_bits(np.frombuffer(frame, dtype=np.float64), actual))
            process.stdin.write(bytes((actual,)))
            process.stdin.flush()
    except CandidateError as error:
        return {"status": "invalid", "error": str(error)}
    except OSError:  # broken pipe: the process was killed or exited
        return died
    return {"status": "ok", "byte_bits": np.frombuffer(byte_bits, dtype=np.float64)}


def _read_error(process: subprocess.Popen) -> dict:
    header = process.stdout.read(4)
    if len(header) != 4:
        return {"status": "crashed", "error": "the candidate process ended early"}
    length = min(struct.unpack("<I", header)[0], MAX_ERROR_BYTES)
    status, _, message = process.stdout.read(length).decode("utf-8", "replace").partition("\n")
    known = {"rejected", "error", "invalid", "memory", "crashed"}
    return {"status": status if status in known else "crashed", "error": message}