"""Muse CLI backend and persistent session map."""

import glob
import hashlib
import json
import os
import shutil
import subprocess
import threading
import time

from tai_pkg.clirun import Cli
from tai_pkg.compat import fcntl
from tai_pkg.consts import Consts
from tai_pkg.errors import BackendError
from tai_pkg.release import Release
from tai_pkg.text import Text

__all__ = ["Muse", "_MUSE_THREAD_LOCKS", "_MUSE_THREAD_GUARD"]


class Muse:
    """Muse."""

    @staticmethod
    def muse_binary():
        return os.environ.get("MUSE_BIN") or shutil.which("muse")

    @staticmethod
    def muse_output_schema():
        call = {"type": "object", "additionalProperties": False, "properties": {"name": {"type": "string"}, "arguments": {"type": "string"}}, "required": ["name", "arguments"]}
        return {"type": "object", "additionalProperties": False, "properties": {"content": {"type": "string"}, "tool_calls": {"type": "array", "items": call}}, "required": ["content", "tool_calls"]}

    @staticmethod
    def muse_prompt(messages, tools):
        parts = []
        for message in messages or []:
            role = (message.get("role") or "user").upper()
            body = message.get("content") or ""
            if isinstance(body, list):
                body = " ".join(str(part.get("text") or "") for part in body if isinstance(part, dict))
            if message.get("tool_call_id"):
                parts.append("%s %s:\n%s" % (role, message["tool_call_id"], body))
            else:
                parts.append("%s:\n%s" % (role, body))
            for call in message.get("tool_calls") or []:
                func = call.get("function") or {}
                parts.append("TOOL CALL %s:\n%s %s" % (call.get("id") or "", func.get("name") or "", func.get("arguments") or ""))
        if tools:
            parts.append("TOOLS:\n%s" % json.dumps(tools))
        parts.append("Answer with content text plus tool_calls for every tool to invoke now. Empty tool_calls ends the turn.")
        return "\n\n".join(part for part in parts if part.strip())

    @staticmethod
    def muse_argv(config, schema_path, prompt_path, session_id=None):
        argv = [Muse.muse_binary(), "exec", "--json", "--max-model-steps", str(Consts.MUSE_MAX_STEPS), "--disable-web-tools", "--disable-reminders", "--allow-workspace-switch", "--output-schema", schema_path, "--prompt-file", prompt_path]
        if session_id is None:
            argv.append("--no-session-log")
        else:
            argv.extend(["--session-id", session_id])
        if config.explicit_model:
            argv.extend(["--model", config.explicit_model])
        return argv

    @staticmethod
    def muse_write_file(path, text):
        with open(path, "w", encoding="utf-8") as handle:
            handle.write(text)
            handle.flush()
            os.fsync(handle.fileno())

    @staticmethod
    def muse_message_hash(message):
        blob = json.dumps(message, sort_keys=True, separators=(",", ":"), default=str)
        return hashlib.sha256(blob.encode("utf-8")).hexdigest()

    @staticmethod
    def muse_map_path(home):
        return os.path.join(home, Consts.MUSE_SESSIONS_FILE)

    @staticmethod
    def muse_load_map(home):
        path = Muse.muse_map_path(home)
        try:
            with open(path, "r", encoding="utf-8") as handle:
                data = json.load(handle)
        except FileNotFoundError:
            return {}
        except (OSError, ValueError):
            Release.quarantine(path)
            return {}
        if not isinstance(data, dict):
            Release.quarantine(path)
            return {}
        now = time.time()
        pruned = False
        for key in list(data):
            entry = data[key]
            if not isinstance(entry, dict) or not isinstance(entry.get("id"), str) or not isinstance(entry.get("hashes"), list):
                del data[key]
                pruned = True
            elif key.startswith("worker/"):
                try:
                    stamp = Text._fromisoformat(entry.get("updated") or "").timestamp()
                except ValueError:
                    stamp = 0
                if now - stamp > Consts.MUSE_WORKER_TTL:
                    del data[key]
                    pruned = True
        if pruned:
            Muse.muse_save_map(home, data)
        return data

    @staticmethod
    def muse_save_map(home, data):
        Muse.muse_write_file(Muse.muse_map_path(home), json.dumps(data, sort_keys=True) + "\n")

    @staticmethod
    def muse_acquire(home, name):
        with _MUSE_THREAD_GUARD:
            thread_lock = _MUSE_THREAD_LOCKS.setdefault(name, threading.Lock())
        thread_lock.acquire()
        handle = None
        if fcntl is not None:
            try:
                folder = os.path.join(home, Consts.MUSE_LOCKS_DIR)
                os.makedirs(folder, exist_ok=True)
                handle = open(os.path.join(folder, hashlib.sha256(name.encode("utf-8")).hexdigest() + ".lock"), "w")
                fcntl.flock(handle.fileno(), fcntl.LOCK_EX)
            except OSError:
                if handle is not None:
                    try:
                        handle.close()
                    except OSError:
                        pass
                    handle = None

        def release():
            try:
                if handle is not None:
                    if fcntl is not None:
                        try:
                            fcntl.flock(handle.fileno(), fcntl.LOCK_UN)
                        except OSError:
                            pass
                    try:
                        handle.close()
                    except OSError:
                        pass
            finally:
                thread_lock.release()

        return release

    @staticmethod
    def muse_transact(home, func):
        release = Muse.muse_acquire(home, "map")
        try:
            data = Muse.muse_load_map(home)
            result = func(data)
            Muse.muse_save_map(home, data)
            return result
        finally:
            release()

    @staticmethod
    def muse_session_root():
        base = os.environ.get("XDG_DATA_HOME") or os.path.join(os.path.expanduser("~"), ".local", "share")
        return os.path.join(base, "muse", "sessions")

    @staticmethod
    def muse_session_alive(session_id):
        try:
            pattern = os.path.join(Muse.muse_session_root(), "*", "*", "*", session_id, "session.jsonl")
            return bool(glob.glob(pattern))
        except (OSError, ValueError):
            return True

    @staticmethod
    def muse_session_split(messages, hashes):
        if len(hashes) > len(messages):
            return None
        for pos, digest in enumerate(hashes):
            if Muse.muse_message_hash(messages[pos]) != digest:
                return None
        return list(messages[len(hashes):])

    @staticmethod
    def muse_valid_answer(text):
        try:
            data = json.loads(text)
        except ValueError:
            return None
        if not isinstance(data, dict):
            return None
        if not isinstance(data.get("content"), str) or not isinstance(data.get("tool_calls"), list):
            return None
        for call in data["tool_calls"]:
            if not isinstance(call, dict) or not isinstance(call.get("name"), str) or not isinstance(call.get("arguments"), str):
                return None
        return data

    @staticmethod
    def muse_gate(func):
        def gated(ok):
            func(ok)

        gated.armed = threading.Event()
        return gated

    @staticmethod
    def muse_reap(proc, stderr_done, on_reap):
        def target():
            ok = False
            try:
                try:
                    for _line in proc.stdout:
                        pass
                except (OSError, ValueError):
                    pass
                try:
                    proc.stdout.close()
                except (OSError, ValueError):
                    pass
                try:
                    ok = proc.wait(timeout=Consts.MUSE_REAP_GRACE) == 0
                except subprocess.TimeoutExpired:
                    try:
                        proc.kill()
                    except OSError:
                        pass
                    try:
                        proc.wait(timeout=10)
                    except (OSError, subprocess.SubprocessError):
                        pass
                try:
                    stderr_done.wait(timeout=10)
                except (OSError, RuntimeError):
                    pass
            finally:
                if on_reap is not None:
                    try:
                        on_reap(ok)
                    except Exception:
                        pass

        thread = threading.Thread(target=target, daemon=True)
        thread.start()
        return thread

    @staticmethod
    def muse_run(argv, cwd, stream_sink, timeout, on_reap=None):
        try:
            proc = subprocess.Popen(argv, stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True, cwd=cwd, bufsize=1)
        except OSError as exc:
            raise BackendError("muse exec failed: %s" % Text.short_error(exc))
        gate = getattr(on_reap, "armed", None)
        if gate is not None:
            gate.set()
        stderr_parts = []
        stderr_done = threading.Event()

        def collect_stderr():
            try:
                for chunk in proc.stderr:
                    stderr_parts.append(chunk)
            except (OSError, ValueError):
                pass
            finally:
                stderr_done.set()

        threading.Thread(target=collect_stderr, daemon=True).start()
        finished = threading.Event()
        expired = threading.Event()

        def watchdog():
            if not finished.wait(timeout):
                expired.set()
                try:
                    proc.kill()
                except OSError:
                    pass

        threading.Thread(target=watchdog, daemon=True).start()
        answer = None
        terminal = None
        buffer = []
        try:
            for line in proc.stdout:
                text = line.strip()
                if not text.startswith("{"):
                    continue
                try:
                    event = json.loads(text)
                except ValueError:
                    continue
                kind = event.get("payload_type") or ""
                payload = event.get("payload") or {}
                if kind == "run.output.delta" and payload.get("text"):
                    piece = payload["text"]
                    if stream_sink is not None:
                        stream_sink(piece)
                    buffer.append(piece)
                    if answer is None:
                        answer = Muse.muse_valid_answer("".join(buffer))
                        if answer is not None:
                            break
                elif kind.startswith("run.terminal."):
                    terminal = payload
                    break
        finally:
            finished.set()
            Muse.muse_reap(proc, stderr_done, on_reap)
        if terminal is not None:
            if terminal.get("terminal") != "completed":
                raise BackendError("muse %s: %s" % (terminal.get("terminal") or "failed", (terminal.get("reason") or "unknown")[:300]))
            parsed = Muse.muse_valid_answer(terminal.get("text") or "")
            if parsed is None:
                raise BackendError("muse returned invalid structured answer")
            return parsed
        if answer is not None:
            return answer
        if expired.is_set():
            raise BackendError("muse timed out after %ds" % timeout)
        try:
            code = proc.returncode
        except (OSError, ValueError):
            code = "?"
        raise BackendError("muse exit %s: %s" % (code, "".join(stderr_parts).strip()[-300:] or "no terminal event"))

    @staticmethod
    def muse_health():
        binary = Muse.muse_binary()
        if not binary:
            return "missing: muse CLI not found"
        return Cli._cli_version_only_health(binary, "muse")


_MUSE_THREAD_LOCKS = {}


_MUSE_THREAD_GUARD = threading.Lock()