"""Provider adapter for verified delivery; guards stay independent of the model vendor."""

from __future__ import annotations

import contextvars
import json
import os
import urllib.parse
import urllib.request
import threading
import time
import hashlib
from concurrent.futures import ThreadPoolExecutor
from datetime import datetime, timezone


# -- time budget of one user answer ---------------------------------------------
# One answer may call the model several times. A single deadline bounds the
# whole answer, so the client never waits longer than the tunnel/host allows.
_deadline = contextvars.ContextVar("answer_deadline", default=None)
# Identity of the answer a call belongs to (its original deadline); a reserved
# window shortens the deadline but keeps the owner for fair admission.
_owner = contextvars.ContextVar("answer_owner", default=None)
# Shortest time a call needs to be useful; the gate never sends it with less.
_call_floor = contextvars.ContextVar("call_floor_seconds", default=0.0)


class DeadlineExceeded(RuntimeError):
    """The answer's time budget is spent; callers fall back to checked evidence."""


def answer_budget_seconds() -> float:
    return float(os.getenv("ANSWER_TIME_BUDGET_SECONDS", "55"))


def start_deadline(seconds: float | None = None):
    """Start a budget unless an outer caller already set one; returns a reset token."""
    if _deadline.get() is not None:
        return None
    value = time.monotonic() + (answer_budget_seconds() if seconds is None else seconds)
    return _deadline.set(value), _owner.set(value)


def reset_deadline(token) -> None:
    if token is not None:
        _deadline.reset(token[0])
        _owner.reset(token[1])


class call_floor:
    """Declare the shortest time the enclosed model calls need to finish.

    A call queued until less time is left is not sent: the endpoint would keep
    generating it after this answer gave up, slowing every other answer.
    """

    def __init__(self, seconds: float):
        self.seconds = seconds
        self.token = None

    def __enter__(self):
        self.token = _call_floor.set(self.seconds)
        return self

    def __exit__(self, *exc):
        _call_floor.reset(self.token)
        return False


def remaining_seconds() -> float | None:
    value = _deadline.get()
    return None if value is None else value - time.monotonic()


def has_budget(seconds: float) -> bool:
    left = remaining_seconds()
    return left is None or left >= seconds


class reserved_budget:
    """Temporarily end the budget earlier, keeping `seconds` for a later step.

    Evidence expansion runs inside this window so the final audit is never
    starved by it. Nested windows only shorten the budget; outside an answer
    budget (tests, offline builds) this is a no-op.
    """

    def __init__(self, seconds: float):
        self.seconds = seconds
        self.token = None

    def __enter__(self):
        value = _deadline.get()
        if value is not None:
            self.token = _deadline.set(value - self.seconds)
        return self

    def __exit__(self, *exc):
        if self.token is not None:
            _deadline.reset(self.token)
        return False


class _ModelGate:
    """Fair, deadline-ordered admission to the shared model endpoint.

    The school endpoint serves two requests at a time and keeps generating
    after a client disconnects. Sending more requests than it serves only
    slows every answer and wastes generations nobody reads. The gate keeps at
    most `slots` requests in flight. While several answers share the endpoint
    each holds at most its fair share of slots (one of two), so a second user
    is not starved until his budget ends; within that rule the nearest
    deadline goes first. A request is never sent when its answer budget would
    end before it could finish (MIN_START or the caller's call_floor).
    """

    MIN_START = 3.0

    def __init__(self):
        self.condition = threading.Condition()
        self.active = 0
        self.by_owner = {}
        self.waiting = []
        self.sequence = 0

    @staticmethod
    def slots() -> int:
        return int(os.getenv("ANSWER_LLM_SLOTS", "2"))

    def _next(self, slots):
        """The waiting entry to admit now, or None."""
        if self.active >= slots or not self.waiting:
            return None
        owners = {owner for owner, count in self.by_owner.items() if count} | {e[2] for e in self.waiting}
        share = slots if len(owners) <= 1 else max(1, slots // len(owners))
        for entry in sorted(self.waiting):
            if self.by_owner.get(entry[2], 0) < share:
                return entry
        return None

    def acquire(self, deadline, owner=None, floor=0.0):
        slots = self.slots()
        if slots <= 0:
            return None
        need = max(self.MIN_START, floor)
        with self.condition:
            self.sequence += 1
            owner = owner if owner is not None else ("call", self.sequence)
            entry = (float("inf") if deadline is None else deadline, self.sequence, owner)
            self.waiting.append(entry)
            try:
                while self._next(slots) != entry:
                    wait = None
                    if deadline is not None:
                        wait = deadline - time.monotonic() - need
                        if wait <= 0:
                            raise DeadlineExceeded("Answer time budget exhausted while queued")
                    self.condition.wait(timeout=None if wait is None else min(wait, 1.0))
                if deadline is not None and deadline - time.monotonic() < need:
                    raise DeadlineExceeded("Answer time budget exhausted while queued")
            except BaseException:
                self.waiting.remove(entry)
                self.condition.notify_all()
                raise
            self.waiting.remove(entry)
            self.active += 1
            self.by_owner[owner] = self.by_owner.get(owner, 0) + 1
            self.condition.notify_all()
            return owner

    def release(self, owner):
        with self.condition:
            self.active -= 1
            self.by_owner[owner] -= 1
            if not self.by_owner[owner]:
                del self.by_owner[owner]
            self.condition.notify_all()


_gate = _ModelGate()


def model_busy() -> bool:
    """Whether other model calls are running or queued right now.

    Called between the steps of one answer (it then has no call in flight),
    so a positive value means another answer shares the endpoint.
    """
    with _gate.condition:
        mine = _owner.get()
        return any(owner != mine for owner, count in _gate.by_owner.items() if count) or any(
            entry[2] != mine for entry in _gate.waiting)


def parallel_map(function, items, *, workers: int = 4) -> list:
    """Run independent model calls concurrently, preserving order and the deadline.

    The school endpoint serves parallel requests in about the time of one.
    Each worker runs in a copy of the caller's context, so the deadline applies.
    Exceptions propagate exactly as in a sequential loop.
    """
    items = list(items)
    if len(items) <= 1:
        return [function(item) for item in items]
    contexts = [contextvars.copy_context() for _ in items]
    with ThreadPoolExecutor(max_workers=min(workers, len(items))) as pool:
        futures = [pool.submit(ctx.run, function, item) for ctx, item in zip(contexts, items)]
        return [future.result() for future in futures]


class NoRedirect(urllib.request.HTTPRedirectHandler):
    def redirect_request(self, req, fp, code, msg, headers, newurl):
        return None


def model_configuration() -> dict:
    provider = os.getenv("ANSWER_PROVIDER", "openai_compatible").strip()
    model = os.getenv("ANSWER_MODEL", "").strip()
    if provider == "anthropic":
        key = (os.getenv("ANSWER_API_KEY") or os.getenv("ANTHROPIC_API_KEY", "")).strip()
        endpoint = os.getenv("ANSWER_LLM_URL", "https://api.anthropic.com/v1/messages").strip()
    elif provider == "openai_compatible":
        legacy = os.getenv("OPENWEBUI_API_KEY", "").strip()
        key = os.getenv("ANSWER_API_KEY", "").strip() or legacy
        endpoint = os.getenv("ANSWER_LLM_URL", "https://ui.tukekemt.xyz/api/chat/completions" if legacy else "").strip()
        model = model or ("model120-fast" if legacy else "")
    else:
        raise RuntimeError("Unsupported answer provider")
    parsed = urllib.parse.urlsplit(endpoint)
    if (not key or not model or any(c in key for c in "\r\n") or parsed.scheme != "https"
            or not parsed.hostname or parsed.username is not None or parsed.password is not None
            or parsed.fragment or parsed.query or any(c.isspace() for c in endpoint)):
        raise RuntimeError("Answer model is not configured")
    return {'provider':provider, 'model':model, 'endpoint':endpoint, 'key':key}


_status_lock = threading.Lock()
_status = {}
HEALTH_MAX_AGE_SECONDS = 900


def configuration_key(config: dict) -> str:
    return hashlib.sha256(json.dumps(config,sort_keys=True).encode()).hexdigest()


def generation_status() -> dict:
    try:
        config = model_configuration()
    except (RuntimeError,ValueError):
        return {'ready':False,'status':'unavailable','reason':'model_not_configured','configured':False}
    public={'configured':True,'provider':config['provider'],'model':config['model']}
    with _status_lock:
        status = dict(_status)
    if status.get('configuration') != configuration_key(config) or time.time()-status.get('checked_at',0)>HEALTH_MAX_AGE_SECONDS:
        return {**public,'ready':False,'status':'unknown','reason':'model_availability_not_verified'}
    return {**public,'ready':status['success'],'status':'ready' if status['success'] else 'unavailable',
            'reason':None if status['success'] else 'model_probe_failed',
            'checked_at_utc':datetime.fromtimestamp(status['checked_at'],timezone.utc).isoformat()}


def complete_with_configured_model(messages: list[dict], *, max_tokens: int | None = None, output_schema: dict | None = None) -> dict:
    config = model_configuration()
    try:
        options={'max_tokens':max_tokens}
        # Compatible endpoints do not necessarily implement constrained decoding.
        # Enable only for a provider/model whose capability was configured.
        if output_schema is not None and os.getenv('ANSWER_STRUCTURED_OUTPUT','prompt')=='json_schema':
            options['output_schema']=strict_output_schema(output_schema)
        result = _complete_request(messages,config,**options)
    except DeadlineExceeded:
        raise  # The answer budget ended; this says nothing about model health.
    except Exception:
        left = remaining_seconds()
        if left is not None and left <= 1:
            raise DeadlineExceeded('Answer time budget exhausted') from None
        with _status_lock:
            _status.update(configuration=configuration_key(config), checked_at=time.time(),success=False)
        raise
    with _status_lock:
        _status.update(configuration=configuration_key(config),checked_at=time.time(),success=True)
    return result


def probe_generation() -> dict:
    try:
        complete_with_configured_model([{'role':'user','content':'Reply only with OK.'}],max_tokens=32)
    except Exception:
        pass # The caller gets only sanitized configuration/availability, never raw exceptions.
    return generation_status()


def _complete_request(messages: list[dict], config: dict, *, max_tokens: int | None = None, output_schema: dict | None = None) -> dict:
    if max_tokens is not None and (type(max_tokens) is not int or not 1 <= max_tokens <= 16384):
        raise ValueError("Invalid model output budget")
    provider,model,key,endpoint = (config[k] for k in ('provider','model','key','endpoint'))
    if provider == "anthropic":
        system = "\n\n".join(row["content"] for row in messages if row["role"] == "system")
        body = {
            "model": model,
            "messages": [row for row in messages if row["role"] != "system"],
            "max_tokens": max_tokens or int(os.getenv("ANSWER_MAX_TOKENS", "4096")),
            "temperature": 0,
        }
        if system:
            body["system"] = system
        if output_schema is not None:
            body['output_config']={'format':{'type':'json_schema','schema':output_schema}}
        headers = {"x-api-key": key, "anthropic-version": "2023-06-01"}
    elif provider == "openai_compatible":
        body = {"model": model, "messages": messages, "temperature": 0, "stream": False}
        if max_tokens is not None:
            body["max_tokens"] = max_tokens
        if output_schema is not None:
            body['response_format']={'type':'json_schema','json_schema':{'name':'question_plan','strict':True,'schema':output_schema}}
        headers = {"Authorization": "Bearer " + key}
    else:
        raise RuntimeError("Unsupported answer provider")
    request = urllib.request.Request(
        endpoint,
        data=json.dumps(body, ensure_ascii=False).encode(),
        headers={"Content-Type": "application/json", **headers},
    )
    if remaining_seconds() is not None and remaining_seconds() < 2:
        raise DeadlineExceeded("Answer time budget exhausted")
    admitted = _gate.acquire(_deadline.get(), _owner.get(), _call_floor.get())
    try:
        # The budget is measured again: waiting in the gate consumed part of it.
        timeout = float(os.getenv("ANSWER_LLM_TIMEOUT_SECONDS", "45"))
        left = remaining_seconds()
        if left is not None:
            if left < 2:
                raise DeadlineExceeded("Answer time budget exhausted")
            timeout = min(timeout, left)
        with urllib.request.build_opener(NoRedirect()).open(request, timeout=timeout) as response:
            data = response.read(2 * 1024 * 1024 + 1)
    finally:
        if admitted is not None:
            _gate.release(admitted)
    if len(data) > 2 * 1024 * 1024:
        raise RuntimeError("Answer model output is too large")
    value = json.loads(data)
    if provider == "anthropic":
        if value.get("stop_reason") != "end_turn":
            raise RuntimeError("Incomplete answer model output")
        text = "".join(
            block["text"] for block in value.get("content", []) if block.get("type") == "text"
        )
        if not text:
            raise RuntimeError("Missing answer model output")
        return {"choices": [{"message": {"content": text}}]}
    if not isinstance(value,dict) or not value.get('choices'):
        raise RuntimeError('Invalid answer model response')
    choice=value['choices'][0]
    if choice.get('finish_reason') not in (None,'stop') or not choice.get('message',{}).get('content'):
        raise RuntimeError('Incomplete answer model output')
    return value


def strict_output_schema(value):
    """Provider schema subset; validate the result against full Pydantic constraints."""
    if isinstance(value,list):return [strict_output_schema(v) for v in value]
    if not isinstance(value,dict):return value
    result={k:strict_output_schema(v) for k,v in value.items() if k not in {'default','minimum','maximum','minLength','maxLength','minItems','maxItems','format'}}
    if result.get('type')=='object':
        result['additionalProperties']=False
        result['required']=list(result.get('properties',{}))
    return result