web_app.py

from __future__ import annotations

import argparse
import csv
import gc
import io
import json
import os
import sys
import tempfile
import threading
import time
import traceback
from pathlib import Path
from typing import Any, Optional

from dotenv import load_dotenv

import gradio as gr


PROJECT_DIR = Path(__file__).resolve().parent
if str(PROJECT_DIR) not in sys.path:
    sys.path.insert(0, str(PROJECT_DIR))

from api.chat_agent import ChatAgent, ChatAgentConfig
from api.llm_agent import AgentConfig, LLMAgent
from utils.components import get_punctuation, load_config, read_sentences_input
from utils.distractor_utils import sorted_distractor_keys
from utils.maze_prompt import DistractorGeneratorPrompt, MazeChatPrompt, MazeLLMPrompt

_CACHE_LOCK = threading.RLock()
_LLM_AGENT_CACHE: dict[str, Any] = {"key": None, "agent": None}
_CHAT_AGENT_CACHE: dict[str, Any] = {"key": None, "agent": None}


def _parse_optional_float(value: str) -> Optional[float]:
    text = str(value or "").strip()
    if not text:
        return None
    return float(text)


def _resolve_lexicon_path(language_code: str, cfg: dict, uploaded_lexicon: Optional[str]) -> Optional[str]:
    if uploaded_lexicon:
        p = Path(uploaded_lexicon)
        return str(p) if p.exists() else None

    cfg_path = str(cfg.get("LEXICON_PATH", "")).strip()
    if cfg_path:
        p = Path(cfg_path)
        if p.exists():
            return str(p)

    default_path = PROJECT_DIR / "data" / "lexicon" / f"lexicon_{language_code}.txt"
    if default_path.exists():
        return str(default_path)
    return None


def _load_input_text(text_input: str, file_input: Optional[str]) -> tuple[str, Optional[str]]:
    if file_input:
        p = Path(file_input)
        if not p.exists():
            raise ValueError(f"Uploaded file does not exist: {file_input}")
        return p.read_text(encoding="utf-8"), str(p)

    text = str(text_input or "").strip()
    if not text:
        raise ValueError("Please type sentences or upload a file.")
    return text, None


def _build_template_dir(language_code: str) -> Path:
    p = PROJECT_DIR / "template" / language_code
    if not p.exists():
        raise ValueError(f"Template folder not found for language code '{language_code}': {p}")
    return p


def _save_temp_json(data: list[dict]) -> str:
    fd, path = tempfile.mkstemp(prefix="llmmaze_out_", suffix=".json")
    os.close(fd)
    Path(path).write_text(json.dumps(data, ensure_ascii=False, indent=2), encoding="utf-8")
    return path


def _save_temp_text(content: str, suffix: str) -> str:
    fd, path = tempfile.mkstemp(prefix="llmmaze_out_", suffix=suffix)
    os.close(fd)
    Path(path).write_text(content, encoding="utf-8")
    return path


def _auto_surprisal_device() -> str:
    try:
        import torch

        if torch.cuda.is_available():
            return "cuda"
    except Exception:
        pass
    return "cpu"


def _condition_words_cell(word_data: dict[str, Any]) -> str:
    condition_words = word_data.get("condition_words")
    if not isinstance(condition_words, dict) or not condition_words:
        return ""
    seen: list[str] = []
    for value in condition_words.values():
        token = str(value)
        if token not in seen:
            seen.append(token)
    if not seen:
        return ""
    if len(seen) == 1:
        return seen[0]
    return "|".join(seen)


def _outputs_to_tsv(outputs: list[dict]) -> str:
    rows: list[dict[str, str]] = []
    max_distractors = 0
    for record in outputs:
        for word_data in record.get("words", []):
            target_cell = _condition_words_cell(word_data)
            if not target_cell:
                target_cell = str(word_data.get("target_word", word_data.get("source", "")))
            distractor_keys = sorted_distractor_keys(word_data)
            max_distractors = max(max_distractors, len(distractor_keys))
            row: dict[str, str] = {"target": target_cell}
            for idx, key in enumerate(distractor_keys, start=1):
                row[f"distractor{idx}"] = str(word_data.get(key, ""))
            rows.append(row)

    output = io.StringIO()
    columns = ["target"] + [f"distractor{i}" for i in range(1, max_distractors + 1)]
    writer = csv.DictWriter(output, fieldnames=columns, delimiter="\t", lineterminator="\n")
    writer.writeheader()
    for row in rows:
        writer.writerow(row)
    return output.getvalue()


def _discover_language_codes() -> list[str]:
    template_dir = PROJECT_DIR / "template"
    if not template_dir.exists():
        return []
    return sorted([p.name for p in template_dir.iterdir() if p.is_dir()])


def _token_joiner_for_language(language_code: str) -> str:
    return "" if language_code in {"zh", "ja", "ko"} else " "


def _release_model_memory() -> None:
    gc.collect()
    try:
        import torch

        if torch.cuda.is_available():
            torch.cuda.empty_cache()
    except Exception:
        pass


def _build_llm_agent(
    *,
    model_id: str,
    lex_path: str,
    puncts: set[str],
    num_distractors: int,
    num_workers: int,
    apply_surprisal: bool,
    min_abs: Optional[float],
    min_delta: float,
    absolute_threshold_only: bool,
    surprisal_device: str,
    token_joiner: str,
    template_dir: Path,
) -> LLMAgent:
    prompt = MazeLLMPrompt(
        path_to_user_template=template_dir / "base.txt",
        path_to_extension_template=template_dir / "extension.txt",
        path_to_system_template=template_dir / "system.txt",
    )
    agent_cfg = AgentConfig(
        model_id=model_id,
        lexicon_path=lex_path,
        punctuations=puncts,
        num_distractors=int(num_distractors),
        num_workers=int(max(1, num_workers)),
        apply_surprisal_threshold=apply_surprisal,
        min_abs=min_abs,
        min_delta=min_delta,
        absolute_threshold_only=bool(absolute_threshold_only),
        surprisal_device=surprisal_device,
        token_joiner=token_joiner,
    )
    return LLMAgent(prompt=prompt, config=agent_cfg)


def _build_chat_agent(
    *,
    model_id: str,
    language_code: str,
    puncts: set[str],
    num_distractors: int,
    token_joiner: str,
    lexicon_mode: bool,
    lex_path: Optional[str],
    apply_surprisal: bool,
    min_abs: Optional[float],
    min_delta: float,
    absolute_threshold_only: bool,
    surprisal_device: str,
    template_dir: Path,
) -> ChatAgent:
    gen_prompt = DistractorGeneratorPrompt(
        path_to_user_template=template_dir / "chat_distractor_gen_base.txt",
        path_to_extension_template=template_dir / "chat_distractor_gen_extension.txt",
        path_to_system_template=template_dir / "system.txt",
    )
    selector_prompt = MazeChatPrompt(
        path_to_user_template=template_dir / "base.txt",
        path_to_extension_template=template_dir / "extension.txt",
        path_to_system_template=template_dir / "system.txt",
    )
    chat_cfg = ChatAgentConfig(
        lexicon_path=lex_path,
        language_code=language_code,
        punctuations=puncts,
        min_candidates=10,
        max_candidates=20,
        num_distractors=int(num_distractors),
        token_joiner=token_joiner,
        lexicon_mode=bool(lexicon_mode),
        apply_surprisal_threshold=apply_surprisal,
        min_abs=min_abs,
        min_delta=min_delta,
        absolute_threshold_only=bool(absolute_threshold_only),
        surprisal_device=surprisal_device,
        model_id=model_id,
    )
    return ChatAgent(
        cfg=chat_cfg,
        selector_prompt=selector_prompt,
        generator_prompt=gen_prompt,
        lexicon_mode=bool(lexicon_mode),
    )


def run_web_generation(
    config_path: str,
    agent_type: str,
    processing_mode: str,
    language_code: str,
    model_id: str,
    num_distractors: int,
    lexicon_mode: bool,
    lexicon_file: Optional[str],
    min_abs_text: str,
    min_delta: float,
    absolute_threshold_only: bool,
    output_format: str,
    text_input: str,
    file_input: Optional[str],
):
    start = time.perf_counter()
    try:
        cfg = load_config(config_path)

        language_code = str(language_code or cfg.get("LANGUAGE_CODE", "zh")).strip()
        model_id = str(model_id or cfg.get("MODEL_ID", "")).strip()
        # Web app input convention: words must be whitespace-tokenized.
        word_separator = " "
        agent_type = str(agent_type or "LLM").strip().lower()
        processing_mode = str(processing_mode or "NATURALISTIC_READING").strip().lower()
        output_format = str(output_format or "TSV").strip().upper()
        min_abs = _parse_optional_float(min_abs_text)
        min_delta = float(min_delta or 0.0)
        surprisal_device = _auto_surprisal_device()
        apply_surprisal = (min_abs is not None) or (min_delta != 0.0) or bool(absolute_threshold_only)
        token_joiner = _token_joiner_for_language(language_code)

        puncts = get_punctuation(language_code)
        template_dir = _build_template_dir(language_code)
        raw_text, uploaded_path = _load_input_text(text_input, file_input)
        input_path = uploaded_path

        if processing_mode == "controlled_experiment":
            if input_path is None:
                # controlled parser expects file-style tabular rows
                fd, tmp_path = tempfile.mkstemp(prefix="llmmaze_controlled_", suffix=".txt")
                os.close(fd)
                Path(tmp_path).write_text(raw_text, encoding="utf-8")
                input_path = tmp_path
            input_data = None
        else:
            # For naturalistic mode, parse line-wise to support both typed text and uploaded files.
            lines = [ln.strip() for ln in raw_text.splitlines() if ln.strip()]
            input_data = read_sentences_input(lines, split_on=word_separator)

        cache_hit = False
        if agent_type == "llm":
            lex_path = _resolve_lexicon_path(language_code, cfg, lexicon_file)
            if not lex_path:
                raise ValueError("Could not find lexicon file for LLM agent.")
            llm_key = (
                model_id,
                str(lex_path),
                int(num_distractors),
                int(max(1, cfg.get("NUM_WORKERS", 1))),
                bool(apply_surprisal),
                min_abs,
                float(min_delta),
                bool(absolute_threshold_only),
                str(surprisal_device),
                token_joiner,
                str(template_dir),
                tuple(sorted(puncts)),
            )
            with _CACHE_LOCK:
                if _LLM_AGENT_CACHE["key"] == llm_key and _LLM_AGENT_CACHE["agent"] is not None:
                    agent = _LLM_AGENT_CACHE["agent"]
                    cache_hit = True
                else:
                    old_agent = _LLM_AGENT_CACHE.get("agent")
                    _LLM_AGENT_CACHE["agent"] = None
                    _LLM_AGENT_CACHE["key"] = None
                    if old_agent is not None:
                        del old_agent
                        _release_model_memory()
                    agent = _build_llm_agent(
                        model_id=model_id,
                        lex_path=lex_path,
                        puncts=puncts,
                        num_distractors=int(num_distractors),
                        num_workers=int(max(1, cfg.get("NUM_WORKERS", 1))),
                        apply_surprisal=apply_surprisal,
                        min_abs=min_abs,
                        min_delta=float(min_delta),
                        absolute_threshold_only=bool(absolute_threshold_only),
                        surprisal_device=surprisal_device,
                        token_joiner=token_joiner,
                        template_dir=template_dir,
                    )
                    _LLM_AGENT_CACHE["agent"] = agent
                    _LLM_AGENT_CACHE["key"] = llm_key
            if processing_mode == "controlled_experiment":
                if not input_path:
                    raise ValueError("Controlled mode requires tabular input (upload file or tabular text).")
                controlled_items = agent.read_controlled_input(input_path, split_on=word_separator)
                outputs = agent.run_controlled_experiment(controlled_items)
            else:
                outputs = agent.run(input_data)
        else:
            lex_path = _resolve_lexicon_path(language_code, cfg, lexicon_file) if lexicon_mode else None
            chat_key = (
                model_id,
                language_code,
                bool(lexicon_mode),
                str(lex_path or ""),
                int(num_distractors),
                bool(apply_surprisal),
                min_abs,
                float(min_delta),
                bool(absolute_threshold_only),
                str(surprisal_device),
                token_joiner,
                str(template_dir),
                tuple(sorted(puncts)),
            )
            with _CACHE_LOCK:
                if _CHAT_AGENT_CACHE["key"] == chat_key and _CHAT_AGENT_CACHE["agent"] is not None:
                    agent = _CHAT_AGENT_CACHE["agent"]
                    cache_hit = True
                else:
                    old_agent = _CHAT_AGENT_CACHE.get("agent")
                    _CHAT_AGENT_CACHE["agent"] = None
                    _CHAT_AGENT_CACHE["key"] = None
                    if old_agent is not None:
                        del old_agent
                        _release_model_memory()
                    agent = _build_chat_agent(
                        model_id=model_id,
                        language_code=language_code,
                        puncts=puncts,
                        num_distractors=int(num_distractors),
                        token_joiner=token_joiner,
                        lexicon_mode=bool(lexicon_mode),
                        lex_path=lex_path,
                        apply_surprisal=apply_surprisal,
                        min_abs=min_abs,
                        min_delta=float(min_delta),
                        absolute_threshold_only=bool(absolute_threshold_only),
                        surprisal_device=surprisal_device,
                        template_dir=template_dir,
                    )
                    _CHAT_AGENT_CACHE["agent"] = agent
                    _CHAT_AGENT_CACHE["key"] = chat_key
            if processing_mode == "controlled_experiment":
                raise ValueError("Controlled mode is currently implemented for LLMAgent only.")
            outputs = agent.run(input_data)

        if output_format == "TSV":
            out_text = _outputs_to_tsv(outputs)
            out_file = _save_temp_text(out_text, ".tsv")
        else:
            out_text = json.dumps(outputs, ensure_ascii=False, indent=2)
            out_file = _save_temp_json(outputs)
        elapsed = time.perf_counter() - start
        status = (
            f"Done in {elapsed:.2f}s | agent={agent_type.upper()} | mode={processing_mode.upper()} | "
            f"records={len(outputs)} | cache={'HIT' if cache_hit else 'MISS'}"
        )
        return status, out_text, out_file
    except Exception as exc:
        elapsed = time.perf_counter() - start
        status = f"Error after {elapsed:.2f}s: {exc.__class__.__name__}: {exc}"
        err_json = json.dumps(
            {
                "error_type": exc.__class__.__name__,
                "error_message": str(exc),
                "traceback": traceback.format_exc(),
            },
            ensure_ascii=False,
            indent=2,
        )
        return status, err_json, None


def build_app() -> gr.Blocks:
    default_config = str(PROJECT_DIR / "config.yaml")
    language_choices = _discover_language_codes() or ["zh"]
    css = """
    #app-title { text-align: center; margin-bottom: 0.2rem; }
    #app-subtitle { text-align: center; margin-top: 0; color: #666; }
    #download-wrap { display: flex; justify-content: center; }
    #download-btn { min-width: 280px; }
    """
    with gr.Blocks(title="Agentic A-Maze Studio", css=css) as app:
        gr.HTML("<h1 id='app-title'>Agentic A-Maze Studio</h1>")
        gr.HTML(
            "<p id='app-subtitle'>Multilingual distractor generation for naturalistic reading and controlled Maze task experiments.</p>"
        )
        gr.Markdown(
            "- **AGENT TYPE**: **LLM** (RECOMMENDED, local Hugging Face model + lexicon) or **CHAT** (OpenAI API).\n"
            "- **LANGUAGE CODE**: Use [wordfreq language codes](https://pypi.org/project/wordfreq/) such as `en`, `zh`, `fa`.\n"
            "- **MODEL_ID**: For LLM use a Hugging Face model ID (e.g. `Qwen/Qwen2.5-3B-Instruct`); for CHAT use an OpenAI model (e.g. `gpt-4o-mini`)."
        )
        config_path = gr.State(default_config)
        with gr.Row():
            agent_type = gr.Dropdown(
                label="AGENT TYPE",
                choices=["LLM", "CHAT"],
                value="LLM",
                info="LLM is recommended for stable lexicon-based multilingual runs.",
            )
            processing_mode = gr.Dropdown(
                label="EXPERIMENT MODE",
                choices=["naturalistic_reading", "controlled_experiment"],
                value="naturalistic_reading",
                info="naturalistic_reading: plain sentence list. controlled_experiment: tabular item/condition/sentence input.",
            )
        with gr.Row():
            language_code = gr.Dropdown(
                label="LANGUAGE CODE",
                choices=language_choices,
                value="zh" if "zh" in language_choices else language_choices[0],
                info="Use codes consistent with wordfreq (e.g., en, zh, fa).",
            )
            model_id = gr.Textbox(
                label="MODEL_ID",
                value="Qwen/Qwen2.5-3B-Instruct",
                info="Find LLM model IDs on Hugging Face model pages; CHAT model IDs from OpenAI docs.",
            )
        with gr.Row():
            num_distractors = gr.Slider(label="Num Distractors", minimum=1, maximum=10, step=1, value=3)
            lexicon_mode = gr.Checkbox(
                label="USE LEXICON MATCHING (CHAT ONLY)",
                value=False,
                info="When ON, CHAT tries lexicon-based candidates before fallback. Recommended for CHAT agent.",
            )
            output_format = gr.Dropdown(
                label="OUTPUT FORMAT",
                choices=["JSON", "TSV"],
                value="TSV",
                info="TSV format: first column target/condition words, following columns distractors.",
            )
        with gr.Row():
            min_abs_text = gr.Textbox(
                label="MIN_ABS",
                value="",
                info="Minimum absolute distractor surprisal. Leave blank to disable.",
            )
            min_delta = gr.Number(
                label="MIN_DELTA",
                value=0.0,
                info="Minimum surprisal gap: distractor surprisal - target surprisal.",
            )
            absolute_threshold_only = gr.Checkbox(
                label="ABSOLUTE THRESHOLD ONLY",
                value=False,
                info="If ON, only MIN_ABS is enforced; MIN_DELTA is ignored.",
            )
        lexicon_file = gr.File(
            label="Optional Lexicon Upload (.txt/.tsv/.csv)",
            type="filepath",
            file_types=[".txt", ".tsv", ".csv"],
        )
        gr.Markdown(
            "If uploaded, this lexicon file is used first; otherwise the app uses config/default lexicon path."
        )
        gr.Markdown("### Input")
        gr.Markdown(
            f"**Input requirement:** words must be separated by whitespace, regardless of language. Please mark numbers, symbols, etc. with **, e.g., **2006**, so that the app will not generate distractors for them."
        )
        text_input = gr.Textbox(
            label="Sentence(s) or tabular controlled text",
            lines=8,
            placeholder="Naturalistic: one tokenized sentence per line.\nControlled: tabular rows (item_id/condition_id/sentence).",
        )
        file_input = gr.File(label="Or upload input file (.txt/.tsv/.csv)", type="filepath")
        run_btn = gr.Button("Run Generation", variant="primary")
        gr.Markdown("### Output")
        status = gr.Textbox(label="Status")
        output_json = gr.Textbox(label="Output Preview", lines=20)
        with gr.Row(elem_id="download-wrap"):
            output_file = gr.DownloadButton(
                label="Download Output File",
                value=None,
                elem_id="download-btn",
                variant="primary",
            )

        run_btn.click(
            fn=run_web_generation,
            inputs=[
                config_path,
                agent_type,
                processing_mode,
                language_code,
                model_id,
                num_distractors,
                lexicon_mode,
                lexicon_file,
                min_abs_text,
                min_delta,
                absolute_threshold_only,
                output_format,
                text_input,
                file_input,
            ],
            outputs=[status, output_json, output_file],
        )
    return app


def main() -> None:
    load_dotenv()
    parser = argparse.ArgumentParser(description="Run the LLM Maze Gradio web app.")
    parser.add_argument("--host", default="0.0.0.0", help="Server host.")
    parser.add_argument("--port", type=int, default=7860, help="Server port.")
    parser.add_argument("--share", action="store_true", help="Create a public Gradio share link.")
    args = parser.parse_args()

    app = build_app()
    app.launch(server_name=args.host, server_port=args.port, share=args.share)


if __name__ == "__main__":
    main()