"""Agents: oracle (reference solution), null (does nothing), llm (OpenAI-compatible tool loop)."""
import json
import os
import time
import urllib.request

import threading
import urllib.error

from .ledger import add_usage, check_budget
from .proxy import BudgetExceeded


class ServerDown(Exception):
    """The model server does not answer (remote local model). The run ends cleanly and is not a result."""


class WindowEnd(Exception):
    """The time window of the series is over (hard stop, also in the middle of a request)."""

SYSTEM_PROMPT = """You are an ABAP developer. You implement a plan on an SAP system with the tools.

Rules:
- Read the plan. Implement the contract exactly. Do not change a public signature in the contract.
- You make the craft decisions: table types, keys, access paths, SQL strategy, exception design,
  and syntax that fits the release target.
- If a business rule is missing, or two rules contradict, stop. Do not create objects.
  Write the gap in your report.
- Do not change an object in the out-of-scope list.
- Create a new object in package $TMP with sap_create_object. Then write the source with
  sap_push_source.
- Read an object that you do not know before you use it.
- After each write, check the result. If the same check fails two times, stop and report.
- Do not add behavior that the plan does not ask for.

At the end, write a short report:
1. Decisions: one line for each craft decision, with the reason.
2. Objects: object, type, action, reason.
3. Verification: what you checked and the result.
4. Deviations and open points.
"""


class NullAgent:
    name = "null"

    def run(self, task, proxy):
        proxy.note("final", "No action.")
        return "No action."


class OracleAgent:
    """Writes the reference solution. Validates harness and scoring (expected score ~ max)."""
    name = "oracle"

    def run(self, task, proxy):
        if task.meta.get("expected_outcome") == "stop":  # the right answer is to stop and name the gap
            report = "Stopped. Gap: " + task.meta.get("expected_gap", {}).get("description", "")
            proxy.note("final", report)
            return report
        for o in task.objects("reference"):
            ident = {"objectType": o["type"], "objectName": o["name"]}
            if o.get("functionGroup"):
                ident["functionGroup"] = o["functionGroup"]
            proxy.call("sap_create_object", dict(ident, packageName="$TMP",
                                                 description=o.get("description", o["name"])[:60]))
            if o.get("messages"):   # message class: messages are written with sap_push_message
                proxy.call("sap_push_message", {"objectName": o["name"], "messages": o["messages"]})
            if o.get("source"):
                proxy.call("sap_push_source", dict(ident, source=o["source"]))
            if o.get("testclasses_source"):
                proxy.call("sap_push_source", dict(ident, includeType="testclasses",
                                                   source=o["testclasses_source"]))
        proxy.note("final", "Reference solution written.")
        return "Reference solution written."


class LlmAgent:
    """OpenAI-compatible chat completions with tool calls (Ollama, MLX server, vLLM, ...).

    Local models are slow: no time limit by default (tool-call budget limits the run).
    One request at a time; tool calls run in sequence.
    """

    def __init__(self, model, base_url=None, api_key=None, max_turns=80, temperature=0.2,
                 max_seconds=None, max_tokens=None, chat_template_kwargs=None, loop_guard=None,
                 deadline=None, watch=False, empty_retries=2, retry_temperature=None, stream_guard=None):
        self.model = model
        self.name = f"llm:{model}"
        self.base_url = (base_url or os.environ.get("LLM_BASE_URL", "http://127.0.0.1:11434/v1")).rstrip("/")
        self.api_key = api_key or os.environ.get("LLM_API_KEY", "none")
        self.max_turns = max_turns
        self.temperature = temperature
        self.max_seconds = max_seconds
        self.request_timeout = 3600  # local models are slow; a hung request still ends
        # Runaway reasoning: 35 of 1965 DeepSeek turns produced 393k output tokens and no content (46 % of the
        # run cost, 2026-10-03). Normal turns: p95 17.5k, max 82k. A cut turn is retried (see run).
        self.max_tokens = max_tokens or (32000 if ":cloud" in (model or "") else None)
        self.empty_retries = empty_retries
        # empty_response fix (2026-10-06): a retry after an empty turn can use another temperature (the same sample often
        # runs away again: 7 of 13 empty runs had two or three capped turns in a row); stream_guard = reasoning tokens
        # after which a streamed turn without any content or tool call is cut and counted as an empty turn
        self.retry_temperature = retry_temperature
        self.stream_guard = stream_guard
        self.loop_guard = loop_guard  # end the run after this many identical pushes in a row (None = off)
        self.deadline = deadline   # absolute time (time.time()) of the hard stop, or None
        self.watch = watch         # remote model: ping the server during a request, end the run when it is gone
        self.end_reason = None
        self.messages, self.tools, self.reasoning, self.turn_usage = [], [], [], []  # for the trajectory record
        self.chat_template_kwargs = chat_template_kwargs  # local server only, e.g. {"enable_thinking": False}

    def _chat(self, messages, tools, temperature=None):
        body = {"model": self.model, "messages": messages, "tools": tools,
                "temperature": temperature if temperature is not None else self.temperature, "parallel_tool_calls": False}
        if self.max_tokens:
            body["max_tokens"] = self.max_tokens
        if self.chat_template_kwargs:
            body["chat_template_kwargs"] = self.chat_template_kwargs
        req = urllib.request.Request(f"{self.base_url}/chat/completions", json.dumps(body).encode(),
                                     {"Content-Type": "application/json",
                                      "Authorization": f"Bearer {self.api_key}"})
        last = None
        if self.stream_guard and not (self.watch or self.deadline):
            try:
                return self._chat_stream(dict(body), messages)
            except Exception as e:  # noqa: BLE001  streaming not usable (server, parse): the normal request below
                last = e
        for attempt in range(4):  # model server errors (HTTP 5xx, timeouts): retry with backoff
            try:
                if self.watch or self.deadline:
                    data = self._post_watched(req)
                else:
                    with urllib.request.urlopen(req, timeout=self.request_timeout) as r:
                        data = json.loads(r.read().decode())
                return data["choices"][0]["message"], data.get("usage", {})
            except (ServerDown, WindowEnd):
                raise
            except Exception as e:  # noqa: BLE001
                last = e
                code = getattr(e, "code", None)
                if code is not None and code < 500 and code != 429:
                    raise
                if self.deadline and time.time() >= self.deadline:
                    raise WindowEnd()
                if self.watch and not self._ping():
                    raise ServerDown(str(e)[:200])
                time.sleep(10 * (attempt + 1))
        raise RuntimeError(f"model request failed after retries: {last}")

    def _chat_stream(self, body, messages):
        """Streamed request. A turn that has produced only reasoning for `stream_guard` tokens (about 3.2 characters per
        token) and no content and no tool call is cut and returned as an empty turn (the run loop retries it)."""
        body["stream"] = True
        body["stream_options"] = {"include_usage": True}
        req = urllib.request.Request(f"{self.base_url}/chat/completions", json.dumps(body).encode(),
                                     {"Content-Type": "application/json", "Authorization": f"Bearer {self.api_key}"})
        content, reasoning, calls, usage, cut = "", 0, {}, {}, False
        with urllib.request.urlopen(req, timeout=self.request_timeout) as r:
            for raw in r:
                line = raw.decode("utf-8", "replace").strip()
                if not line.startswith("data:"):
                    continue
                data = line[5:].strip()
                if data == "[DONE]":
                    break
                chunk = json.loads(data)
                if chunk.get("usage"):
                    usage = chunk["usage"]
                for ch in chunk.get("choices") or []:
                    d = ch.get("delta") or {}
                    content += d.get("content") or ""
                    reasoning += len(d.get("reasoning") or d.get("reasoning_content") or d.get("thinking") or "")
                    for tc in d.get("tool_calls") or []:
                        c = calls.setdefault(tc.get("index", 0), {"id": tc.get("id") or "", "type": "function",
                                                                  "function": {"name": "", "arguments": ""}})
                        c["id"] = c["id"] or tc.get("id") or ""
                        f = tc.get("function") or {}
                        c["function"]["name"] += f.get("name") or ""
                        a = f.get("arguments")
                        c["function"]["arguments"] += a if isinstance(a, str) else json.dumps(a) if a else ""
                if not content.strip() and not calls and reasoning / 3.2 >= self.stream_guard:
                    cut = True
                    break
        if cut:
            est = {"prompt_tokens": int(len(json.dumps(messages)) / 3.5), "completion_tokens": int(reasoning / 3.2),
                   "estimated": True, "cut_by_stream_guard": True}
            return {"role": "assistant", "content": ""}, est
        msg = {"role": "assistant", "content": content or None}
        if calls:
            msg["tool_calls"] = [calls[k] for k in sorted(calls)]
        return msg, usage

    def _ping(self, timeout=10):
        try:
            urllib.request.urlopen(urllib.request.Request(f"{self.base_url}/models"), timeout=timeout).read()
            return True
        except Exception:  # noqa: BLE001
            return False

    def _post_watched(self, req):
        """The request runs in a thread; this thread watches the deadline and (remote model) the server: a hard stop
        or a dead server ends the wait at once, also in the middle of a long generation."""
        box = {}

        def work():
            try:
                with urllib.request.urlopen(req, timeout=self.request_timeout) as r:
                    box["data"] = json.loads(r.read().decode())
            except BaseException as e:  # noqa: BLE001
                box["err"] = e
        th = threading.Thread(target=work, daemon=True)
        th.start()
        last_ping, fails = time.time(), 0
        while th.is_alive():
            th.join(5)
            if self.deadline and time.time() >= self.deadline:
                raise WindowEnd()
            if self.watch and th.is_alive() and time.time() - last_ping >= 30:
                last_ping = time.time()
                fails = 0 if self._ping() else fails + 1
                if fails >= 3:
                    raise ServerDown("no answer to 3 pings in a row during a request")
        if "err" in box:
            raise box["err"]
        return box["data"]

    def run(self, task, proxy):
        tools = [{"type": "function", "function": {"name": t["name"],
                                                   "description": t.get("description", ""),
                                                   "parameters": t.get("inputSchema", {})}}
                 for t in proxy.schemas()]
        messages = [{"role": "system", "content": SYSTEM_PROMPT},
                    {"role": "user", "content": task.spec}]
        self.messages, self.tools, self.reasoning, self.turn_usage = messages, tools, [], []
        if ":cloud" in self.model:
            check_budget()
        final = ""
        empty = 0
        self.end_reason = "max_turns"
        start = time.time()
        for _ in range(self.max_turns):
            if self.deadline and time.time() >= self.deadline:
                final = "Stopped: time window over."
                self.end_reason = "window_end"
                break
            if self.max_seconds and time.time() - start > self.max_seconds:
                final = f"Stopped: time budget exceeded ({self.max_seconds} s)."
                self.end_reason = "time_budget"
                break
            try:
                msg, usage = self._chat(messages, tools,
                                        self.retry_temperature if (empty and self.retry_temperature) else None)
                add_usage(self.model, usage, kind="run", ref=proxy.prefix)
                self.turn_usage.append(usage)
            except WindowEnd:
                final = "Stopped: time window over."
                self.end_reason = "window_end"
                break
            except ServerDown as e:
                final = f"Stopped: model server not reachable ({e})."
                self.end_reason = "server_down"
                break
            except RuntimeError as e:
                final = f"Stopped: {e}"
                self.end_reason = "model_error"
                break
            proxy.note("assistant", {"content": msg.get("content"),
                                     "tool_calls": msg.get("tool_calls"), "usage": usage})
            messages.append({k: v for k, v in msg.items() if k in ("role", "content", "tool_calls")})
            think = msg.get("reasoning") or msg.get("reasoning_content") or msg.get("thinking")
            if think:  # kept apart from the messages: not training data
                self.reasoning.append({"message_index": len(messages) - 1, "reasoning": think})
            calls = msg.get("tool_calls") or []
            if not calls and not (msg.get("content") or "").strip() and empty < self.empty_retries:
                empty += 1  # empty turn (output limit reached in reasoning): ask the same question again
                messages.pop()
                continue
            if not calls:
                final = msg.get("content") or ("Stopped: empty model response." if empty else "")
                self.end_reason = "report" if (msg.get("content") or "").strip() else "empty_response"
                break
            try:
                for c in calls:
                    fn = c["function"]
                    args = fn.get("arguments") or "{}"
                    args = json.loads(args) if isinstance(args, str) else args
                    err, text = proxy.call(fn["name"], args)
                    messages.append({"role": "tool", "tool_call_id": c.get("id", ""),
                                     "content": ("ERROR: " if err else "") + text[:12000]})
                    streak = max(proxy.same_push_streak, proxy.same_call_streak)
                    if self.loop_guard and streak >= self.loop_guard:
                        final = (f"Stopped: loop (the same source was pushed {proxy.same_push_streak} times in a row)."
                                 if proxy.same_push_streak >= self.loop_guard else
                                 f"Stopped: loop (the same call with the same result {proxy.same_call_streak} times in a row).")
                        self.end_reason = "loop"
                        break
                if self.end_reason == "loop":
                    break
            except BudgetExceeded as e:
                final = f"Stopped: budget exceeded ({e})."
                self.end_reason = "tool_budget"
                break
        proxy.note("final", final)
        return final