import json
import os
import subprocess
import sys
import threading
import time
from datetime import datetime, timezone
from pathlib import Path

import gradio as gr
import spaces
import torch


def install_local_wheel(env_name: str, default_path: str, import_name: str) -> None:
    """Install a private package supplied as a wheel in the Space repository."""
    try:
        __import__(import_name)
        return
    except ModuleNotFoundError:
        pass

    wheel_path = Path(os.getenv(env_name, default_path))
    if not wheel_path.is_file():
        raise RuntimeError(
            f"Package '{import_name}' is not installed and its wheel is missing: "
            f"{wheel_path}. Add the file to the Space repository or set {env_name}."
        )
    subprocess.check_call(
        [sys.executable, "-m", "pip", "install", str(wheel_path), "--quiet"]
    )


# These files were also required by the supplied original demo.
install_local_wheel(
    "RXLM_WHEEL",
    "/home/user/app/rxlm-0.3.101-py3-none-any.whl",
    "rxlm",
)
install_local_wheel(
    "RXLM_PRO_WHEEL",
    "/home/user/app/rxlm_pro-0.2.71-py3-none-any.whl",
    "rxlm_pro",
)

from rxlm.training.tokenizer import load_tokenizer_from_hf_hub
from rxlm_pro.models.rxq import RxQwenDense


MODEL_ID = os.getenv("MODEL_ID", "AdamF92/RxQwen-Micro-2B-Chat")
TOKENIZER_ID = os.getenv("TOKENIZER_ID", "AdamF92/RxQwen-Nano-0.8B")
MODEL_REVISION = os.getenv("MODEL_REVISION") or None
HF_TOKEN = os.getenv("HF_TOKEN")
MEMORY_DIR = Path(os.getenv("MEMORY_DIR", "./memory"))
MEMORY_DIR.mkdir(parents=True, exist_ok=True)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
tokenizer = load_tokenizer_from_hf_hub(TOKENIZER_ID, token=HF_TOKEN)
model = RxQwenDense.from_pretrained(
    MODEL_ID,
    token=HF_TOKEN,
    tokenizer=tokenizer,
    revision=MODEL_REVISION,
)
model.to(device)
model.init_model(device=device)
model.set_batch_mode(False)

INITIAL_STM = model.export_stm_state().cpu()
MODEL_LOCK = threading.Lock()
STOP_EVENT = threading.Event()


def safe_name(name: str) -> str:
    cleaned = "".join(c for c in name.strip() if c.isalnum() or c in "-_")
    return cleaned[:64] or "default"


def memory_paths(name: str) -> tuple[Path, Path]:
    stem = safe_name(name)
    return MEMORY_DIR / f"{stem}.pt", MEMORY_DIR / f"{stem}.jsonl"


def load_memory(name: str) -> torch.Tensor:
    state_path, _ = memory_paths(name)
    if not state_path.exists():
        return INITIAL_STM.clone()
    state = torch.load(state_path, map_location="cpu", weights_only=True)
    if not isinstance(state, torch.Tensor):
        raise ValueError("The memory file does not contain an STM tensor.")
    return state


def save_memory(name: str, state: torch.Tensor, event: dict) -> None:
    state_path, journal_path = memory_paths(name)
    temporary_path = state_path.with_suffix(".tmp")
    torch.save(state.cpu(), temporary_path)
    os.replace(temporary_path, state_path)
    event = {"timestamp": datetime.now(timezone.utc).isoformat(), **event}
    with journal_path.open("a", encoding="utf-8") as journal:
        journal.write(json.dumps(event, ensure_ascii=False) + "\n")


def stream_generate(prompt: str, stm_state: torch.Tensor):
    """Yield response text, final STM, token delta, and query token count."""
    tokenized = model.tokenize_query(prompt, max_seq_len=8192, device=device)
    query_tokens = int(tokenized["input_ids"].size(-1))
    pieces: list[str] = []
    yield "", None, 0, query_tokens
    with MODEL_LOCK:
        model.load_stm_state(stm_state)
        with torch.inference_mode(), torch.amp.autocast(
            device_type=device.type,
            dtype=torch.bfloat16,
            enabled=device.type == "cuda",
        ):
            for token_id in model.interact(
                **tokenized,
                thinking_mode="extended",
                max_seq_len=8192,
                temperature=0.4,
                top_p=0.9,
            ):
                if token_id == -2:
                    continue
                token = model.stringify_token(
                    token_id,
                    show_memory_update=False,
                    skip_special_tokens=False,
                )
                if token not in {"[T]", "[A]"}:
                    pieces.append(token)
                yield "".join(pieces), None, 1, query_tokens
        new_state = model.export_stm_state().cpu()
    yield "".join(pieces).strip(), new_state, 0, query_tokens


def iteration_prompt(task: str, iteration: int) -> str:
    # Build the prompt from separate strings so web editors cannot lose a closing
    # triple quote while copying the file.
    return (
        "You are an autonomous agent working on one task.\n\n"
        "TASK:\n"
        + task
        + "\n\nThis is iteration "
        + str(iteration)
        + ". Use the information stored in your persistent memory.\n"
        "Take the single most valuable next step. Evaluate the result and store "
        "important facts, decisions, errors, and the next step in memory. Never "
        "claim to have performed actions for which you have no tools. If the task "
        "is genuinely complete, end with the exact marker [TASK_COMPLETE]. "
        "Otherwise, end with a short description of the next step."
    )


@spaces.GPU
def run_agent(
    task: str,
    memory_name: str,
    max_iterations: int,
    continuous: bool,
):
    state_input_tokens = 0
    llm_input_tokens = 0
    output_tokens = 0

    def result(current_history, current_status):
        return (
            current_history,
            current_status,
            state_input_tokens,
            llm_input_tokens,
            output_tokens,
        )

    if not task or not task.strip():
        yield result([], "Enter a task.")
        return

    STOP_EVENT.clear()
    name = safe_name(memory_name)
    history = []
    try:
        state = load_memory(name)
    except Exception as exc:
        yield result(history, f"Could not load memory: {exc}")
        return

    iteration = 1
    while continuous or iteration <= int(max_iterations):
        if STOP_EVENT.is_set():
            yield result(
                history,
                f"Stopped after {iteration - 1} iterations. Memory was saved.",
            )
            return

        limit_label = "continuous" if continuous else str(int(max_iterations))
        yield result(history, f"Iteration {iteration}/{limit_label}...")
        try:
            response = ""
            final_state = None
            counted_query = False
            for partial_response, possible_state, token_delta, query_tokens in stream_generate(
                iteration_prompt(task.strip(), iteration), state
            ):
                if not counted_query:
                    # State models read only the new query. A conventional LLM must
                    # read the accumulated context again on every turn.
                    previous_state_inputs = state_input_tokens
                    state_input_tokens += query_tokens
                    llm_input_tokens += (
                        previous_state_inputs + output_tokens + query_tokens
                    )
                    counted_query = True
                output_tokens += token_delta
                response = partial_response
                if possible_state is not None:
                    final_state = possible_state
                live_history = history + [
                    {"role": "user", "content": f"Iteration {iteration}"},
                    {"role": "assistant", "content": response},
                ]
                yield result(
                    live_history,
                    f"Iteration {iteration}/{limit_label}: reasoning live...",
                )

            if final_state is None:
                raise RuntimeError("The model did not return an updated STM state.")
            state = final_state
            save_memory(
                name,
                state,
                {"iteration": iteration, "task": task.strip(), "response": response},
            )
        except Exception as exc:
            yield result(history, f"Error in iteration {iteration}: {exc}")
            return

        history = history + [
            {"role": "user", "content": f"Iteration {iteration}"},
            {"role": "assistant", "content": response},
        ]
        if "[TASK_COMPLETE]" in response:
            yield result(
                history,
                f"Task completed in iteration {iteration}. Memory: {name}",
            )
            return
        yield result(history, f"Iteration {iteration} saved; continuing...")
        iteration += 1
        time.sleep(0.1)

    yield result(
        history,
        f"Reached the limit of {max_iterations} iterations. Memory was saved.",
    )


def stop_agent():
    STOP_EVENT.set()
    return "Stop requested..."


def reset_memory(memory_name: str):
    state_path, journal_path = memory_paths(memory_name)
    for path in (state_path, journal_path):
        path.unlink(missing_ok=True)
    return [], f"Deleted memory: {safe_name(memory_name)}", 0, 0, 0


with gr.Blocks(title="Persistent Agent Loop") as demo:
    gr.Markdown(
        "# Live Persistent Agent Loop\n"
        "Watch the model reason token by token while it repeatedly works on one task."
    )
    task = gr.Textbox(
        label="Task",
        lines=5,
        value=(
            "Create a practical five-step plan for launching a small AI-powered "
            "study assistant. Refine the plan on every turn, identify risks, and "
            "finish with a concise implementation checklist."
        ),
    )
    with gr.Row():
        memory_name = gr.Textbox(label="Memory name", value="default")
        max_iterations = gr.Slider(1, 50, value=5, step=1, label="Turns")
        continuous = gr.Checkbox(
            value=False,
            label="Continuous loop (runs until Stop or TASK_COMPLETE)",
        )
    with gr.Row():
        start = gr.Button("Run", variant="primary")
        stop = gr.Button("Stop")
        reset = gr.Button("Clear memory")
    status = gr.Textbox(label="Status", interactive=False)
    with gr.Row():
        state_token_metric = gr.Number(
            value=0,
            label="State Model Input Tokens",
            interactive=False,
        )
        llm_token_metric = gr.Number(
            value=0,
            label="Conventional LLM Input Tokens",
            interactive=False,
        )
        output_token_metric = gr.Number(
            value=0,
            label="Output Tokens",
            interactive=False,
        )
    # The Gradio version used by this Space does not support the `type` parameter.
    log = gr.Chatbot(label="Progress", height=500)

    start.click(
        run_agent,
        [task, memory_name, max_iterations, continuous],
        [
            log,
            status,
            state_token_metric,
            llm_token_metric,
            output_token_metric,
        ],
    )
    stop.click(stop_agent, outputs=status, queue=False)
    reset.click(
        reset_memory,
        memory_name,
        [
            log,
            status,
            state_token_metric,
            llm_token_metric,
            output_token_metric,
        ],
    )


if __name__ == "__main__":
    demo.queue(default_concurrency_limit=1).launch()