ui_app.py

from __future__ import annotations

import subprocess
import sys
from pathlib import Path
from typing import Any

import yaml
from flask import Flask, jsonify, render_template, request


ROOT = Path(__file__).resolve().parent
CONFIG_DIR = ROOT / "configs"

app = Flask(__name__)


@app.get("/")
def index() -> str:
    return render_template("index.html")


@app.get("/api/configs")
def list_configs() -> Any:
    CONFIG_DIR.mkdir(parents=True, exist_ok=True)
    configs = sorted(path.name for path in CONFIG_DIR.iterdir() if path.suffix in {".yaml", ".yml"})
    return jsonify({"ok": True, "configs": configs})


@app.get("/api/config/<name>")
def load_config(name: str) -> Any:
    path = _safe_config_path(name)
    if not path.exists():
        return jsonify({"ok": False, "error": "Config not found"}), 404
    yaml_text = path.read_text(encoding="utf-8")
    try:
        config = _load_yaml_config(yaml_text)
    except yaml.YAMLError as exc:
        return jsonify({"ok": False, "error": f"YAML error: {exc}"}), 400
    return jsonify({"ok": True, "yaml": yaml_text, "config": config})


@app.post("/api/parse")
def parse_config() -> Any:
    payload = request.get_json(force=True)
    yaml_text = payload.get("yaml") or ""
    try:
        config = _load_yaml_config(yaml_text)
    except yaml.YAMLError as exc:
        return jsonify({"ok": False, "error": f"YAML error: {exc}"}), 400
    return jsonify({"ok": True, "config": config})


@app.post("/api/save")
def save_config() -> Any:
    payload = request.get_json(force=True)
    name = payload.get("name") or "ui_generated.yaml"
    yaml_text = payload.get("yaml") or ""
    path = _safe_config_path(name)
    try:
        _load_yaml_config(yaml_text)
    except yaml.YAMLError as exc:
        return jsonify({"ok": False, "error": f"YAML error: {exc}"}), 400
    path.write_text(yaml_text, encoding="utf-8")
    return jsonify({"ok": True, "path": str(path.relative_to(ROOT))})


@app.post("/api/generate")
def generate_config() -> Any:
    payload = request.get_json(force=True)
    config = build_config(payload)
    return jsonify({"ok": True, "yaml": yaml.safe_dump(config, allow_unicode=True, sort_keys=False)})


@app.post("/api/run")
def run_config() -> Any:
    payload = request.get_json(force=True)
    name = payload.get("name") or "ui_generated.yaml"
    yaml_text = payload.get("yaml") or ""
    path = _safe_config_path(name)
    try:
        _load_yaml_config(yaml_text)
    except yaml.YAMLError as exc:
        return jsonify({"ok": False, "error": f"YAML error: {exc}"}), 400
    path.write_text(yaml_text, encoding="utf-8")
    process = subprocess.run(
        [sys.executable, "main.py", "--config", str(path.relative_to(ROOT))],
        cwd=ROOT,
        text=True,
        capture_output=True,
        timeout=300,
    )
    return jsonify(
        {
            "ok": process.returncode == 0,
            "returncode": process.returncode,
            "stdout": process.stdout,
            "stderr": process.stderr,
        }
    )


def build_config(payload: dict[str, Any]) -> dict[str, Any]:
    task_name = payload.get("task_name") or "ui_task"
    request_type = payload.get("request_type") or "static"
    pagination_type = payload.get("pagination_type") or "none"
    storage_type = payload.get("storage_type") or "jsonl"
    output_path = payload.get("output_path") or f"data/output/{task_name}.{_storage_ext(storage_type)}"
    fields = _build_fields(payload.get("fields") or "")
    api_fields = _build_api_fields(payload.get("api_fields") or "")
    media_enabled = bool(payload.get("media_enabled"))
    parse_attachments = bool(payload.get("parse_attachments"))
    dynamic_enabled = request_type == "dynamic" or bool(payload.get("browser_enabled"))

    config: dict[str, Any] = {
        "task": {
            "name": task_name,
            "purpose": payload.get("purpose") or "信息采集",
            "description": payload.get("description") or "Generated by local UI",
        },
        "target": {
            "site_name": payload.get("site_name") or task_name,
            "base_url": payload.get("base_url") or "",
            "entry_urls": _lines(payload.get("entry_urls") or ""),
            "allowed_domains": _lines(payload.get("allowed_domains") or ""),
        },
        "request": {
            "method": payload.get("method") or "GET",
            "type": request_type,
            "headers": {
                "User-Agent": payload.get("user_agent")
                or "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36",
            },
            "cookies": {},
            "params": {},
            "payload": {},
            "encoding": payload.get("encoding") or "auto",
            "timeout": {"connect": 10, "read": 30},
        },
        "browser": {
            "enabled": dynamic_enabled,
            "engine": payload.get("browser_engine") or "chromium",
            "headless": bool(payload.get("headless", True)),
            "wait_until": payload.get("wait_until") or "networkidle",
            "wait_selector": payload.get("wait_selector") or "",
            "actions": _browser_actions(payload),
        },
        "pagination": _pagination(payload, pagination_type),
        "list": {
            "item_selector": payload.get("item_selector") or "",
            "fields": fields,
        },
        "detail": {
            "enabled": bool(payload.get("detail_enabled")),
            "url_field": payload.get("detail_url_field") or "detail_url",
            "keep_on_detail_failed": True,
            "fields": _build_fields(payload.get("detail_fields") or ""),
        },
        "api": {
            "enabled": request_type == "api",
            "data_path": payload.get("api_data_path") or "$",
            "fields": api_fields,
        },
        "media": {
            "enabled": media_enabled,
            "output_dir": payload.get("media_output_dir") or f"data/media/{task_name}",
            "skip_existing": True,
            "parse_attachments": parse_attachments,
            "parse_config": {"max_chars": int(payload.get("parse_max_chars") or 200000)},
            "resources": _media_resources(payload),
        },
        "cleaning": {
            "strip": True,
            "normalize_space": True,
            "html_unescape": True,
            "remove_tracking_params": True,
        },
        "dedupe": {
            "enabled": bool(payload.get("dedupe_enabled", True)),
            "type": payload.get("dedupe_type") or "persistent",
            "key_fields": _csv(payload.get("dedupe_fields") or ""),
            "path": f"data/state/{task_name}_dedupe.txt",
        },
        "resume": {
            "enabled": bool(payload.get("resume_enabled", True)),
            "retry_failed_first": bool(payload.get("retry_failed_first", True)),
            "completed_path": f"data/state/{task_name}_completed_urls.txt",
            "failed_path": f"data/state/{task_name}_failed.jsonl",
        },
        "rate_limit": {
            "delay_min": float(payload.get("delay_min") or 0),
            "delay_max": float(payload.get("delay_max") or 0),
            "concurrency": 1,
        },
        "retry": {
            "times": int(payload.get("retry_times") or 2),
            "backoff": float(payload.get("retry_backoff") or 1),
            "retry_status": [429, 500, 502, 503, 504],
        },
        "storage": {
            "type": storage_type,
            "path": output_path,
            "table": payload.get("table") or "items",
        },
        "logging": {
            "enabled": bool(payload.get("logging_enabled", True)),
            "level": "INFO",
            "request_log": f"data/logs/{task_name}_requests.jsonl",
            "error_log": f"data/logs/{task_name}_errors.jsonl",
            "summary_log": f"data/logs/{task_name}_summary.jsonl",
        },
    }
    if payload.get("authorization"):
        config["request"]["headers"]["Authorization"] = payload["authorization"]
    if payload.get("referer"):
        config["request"]["headers"]["Referer"] = payload["referer"]
    return config


def _pagination(payload: dict[str, Any], pagination_type: str) -> dict[str, Any]:
    enabled = pagination_type != "none"
    return {
        "enabled": enabled,
        "type": pagination_type,
        "page_param": payload.get("page_param") or "page",
        "start_page": int(payload.get("start_page") or 1),
        "end_page": int(payload.get("end_page") or 1),
        "url_template": payload.get("url_template") or "",
        "next_page_selector": payload.get("next_page_selector") or "",
        "max_pages": int(payload.get("max_pages") or 1),
        "cursor_param": payload.get("cursor_param") or "cursor",
        "cursor_start": payload.get("cursor_start") or "",
        "next_cursor_path": payload.get("next_cursor_path") or "$.data.next_cursor",
        "stop_when_empty_path": payload.get("stop_when_empty_path") or "$.data.items",
    }


def _browser_actions(payload: dict[str, Any]) -> list[dict[str, Any]]:
    actions = []
    if payload.get("scroll_enabled"):
        actions.append(
            {
                "type": "scroll",
                "times": int(payload.get("scroll_times") or 3),
                "pause": float(payload.get("scroll_pause") or 1),
            }
        )
    wait_seconds = payload.get("wait_seconds")
    if wait_seconds:
        actions.insert(0, {"type": "wait", "seconds": float(wait_seconds)})
    return actions


def _media_resources(payload: dict[str, Any]) -> list[dict[str, Any]]:
    resources = []
    for media_type in ("image", "audio", "video", "subtitle", "attachment"):
        fields = payload.get(f"{media_type}_fields")
        if fields:
            resources.append(
                {
                    "name": f"{media_type}s",
                    "type": media_type,
                    "url_fields": _csv(fields),
                    "filename_prefix": media_type,
                    "parse": media_type == "attachment" and bool(payload.get("parse_attachments")),
                    "include_text": media_type == "attachment",
                }
            )
    return resources


def _build_fields(text: str) -> dict[str, Any]:
    fields = {}
    for line in _lines(text):
        parts = [part.strip() for part in line.split("|")]
        if len(parts) < 2:
            continue
        name, selector = parts[0], parts[1]
        rule: dict[str, Any] = {"selector": selector}
        if len(parts) >= 3 and parts[2]:
            rule["mode"] = parts[2]
        if len(parts) >= 4 and parts[3].lower() in {"true", "many", "1"}:
            rule["many"] = True
        if len(parts) >= 5 and parts[4]:
            rule["transform"] = parts[4]
        fields[name] = rule
    return fields


def _build_api_fields(text: str) -> dict[str, str]:
    fields = {}
    for line in _lines(text):
        if "|" not in line:
            continue
        name, path = [part.strip() for part in line.split("|", 1)]
        fields[name] = path
    return fields


def _lines(text: str) -> list[str]:
    return [line.strip() for line in text.splitlines() if line.strip()]


def _csv(text: str) -> list[str]:
    return [item.strip() for item in text.split(",") if item.strip()]


def _storage_ext(storage_type: str) -> str:
    return {"csv": "csv", "sqlite": "db"}.get(storage_type, "jsonl")


def _load_yaml_config(yaml_text: str) -> dict[str, Any]:
    data = yaml.safe_load(yaml_text) or {}
    if not isinstance(data, dict):
        raise yaml.YAMLError("Top-level YAML must be a mapping/object")
    return data


def _safe_config_path(name: str) -> Path:
    filename = Path(name).name
    if not filename.endswith((".yaml", ".yml")):
        filename += ".yaml"
    CONFIG_DIR.mkdir(parents=True, exist_ok=True)
    return CONFIG_DIR / filename


if __name__ == "__main__":
    app.run(host="127.0.0.1", port=5000, debug=False)