1#!/usr/bin/env python3
2"""
3Chart Skill — project-based static server + save APIs
4 
5Usage:
6  python3 chart_server.py [serve_dir] [port]
7 
8Defaults:
9  serve_dir = /data/workspace/output/chart-html
10  port = 7860
11 
12Endpoints:
13- POST /save-chart : Save PNG to current project directory as screenshot.png (or filename)
14- GET  /           : Static files from serve_dir
15 
16Notes:
17- Gallery and library APIs are intentionally removed.
18- Each chart should live in: output/chart-html/<project>/index.html
19"""
20 
21import sys
22import os
23import json
24import base64
25import re
26from pathlib import Path
27from urllib.parse import urlsplit
28from http.server import HTTPServer, SimpleHTTPRequestHandler
29 
30WORKSPACE = Path("/data/workspace")
31DEFAULT_SERVE_DIR = WORKSPACE / "output" / "chart-html"
32DEFAULT_SERVE_DIR.mkdir(parents=True, exist_ok=True)
33 
34 
35def _safe_filename(name: str, ext: str) -> str:
36    name = re.sub(r"[^\w\u4e00-\u9fff.\-]", "_", name or "")
37    if not name:
38        name = f"screenshot{ext}"
39    if not name.endswith(ext):
40        name += ext
41    return name
42 
43 
44class ChartHandler(SimpleHTTPRequestHandler):
45    def __init__(self, *args, serve_dir=None, **kwargs):
46        self._serve_dir = Path(serve_dir or DEFAULT_SERVE_DIR)
47        super().__init__(*args, directory=str(self._serve_dir), **kwargs)
48 
49    def log_message(self, format, *args):
50        pass
51 
52    def _normalized_path(self):
53        # Preview proxy may forward as /preview/{id}/<endpoint>
54        p = self.path.split('?', 1)[0]
55        m = re.match(r"^/preview/[^/]+/(.*)$", p)
56        if m:
57            p = '/' + m.group(1)
58        return p
59 
60    def _translate_preview_path(self, raw_path: str) -> str:
61        """Map /preview/<id>/<project>/... to /<project>/... before static lookup."""
62        # keep query string when present
63        split = urlsplit(raw_path)
64        p = split.path
65        m = re.match(r"^/preview/[^/]+/(.*)$", p)
66        if m:
67            mapped = '/' + m.group(1)
68        else:
69            mapped = p
70        if split.query:
71            mapped = f"{mapped}?{split.query}"
72        return mapped
73 
74    def translate_path(self, path):
75        # SimpleHTTPRequestHandler resolves filesystem paths via this method.
76        # We normalize preview-prefixed URLs first so static files map correctly.
77        return super().translate_path(self._translate_preview_path(path))
78 
79    def do_OPTIONS(self):
80        self.send_response(200)
81        self.send_header("Access-Control-Allow-Origin", "*")
82        self.send_header("Access-Control-Allow-Methods", "GET, POST, OPTIONS")
83        self.send_header("Access-Control-Allow-Headers", "Content-Type")
84        self.end_headers()
85 
86    def do_GET(self):
87        return super().do_GET()
88 
89    def do_POST(self):
90        p = self._normalized_path()
91        if p in ("/save-chart", "/save-chart/"):
92            return self._handle_save_chart()
93        self.send_error(404, "Not found")
94 
95    def _read_json_body(self):
96        length = int(self.headers.get("Content-Length", 0))
97        raw = self.rfile.read(length)
98        return json.loads(raw)
99 
100    def _send_json(self, data, status=200):
101        body = json.dumps(data, ensure_ascii=False)
102        self.send_response(status)
103        self.send_header("Content-Type", "application/json; charset=utf-8")
104        self.send_header("Content-Length", str(len(body.encode("utf-8"))))
105        self.send_header("Access-Control-Allow-Origin", "*")
106        self.end_headers()
107        self.wfile.write(body.encode("utf-8"))
108 
109    def _resolve_project_dir(self, project: str):
110        # project like "btc-90d-20260401" or "btc-90d-20260401/subdir"
111        project = (project or "").strip().strip('/')
112        if not project:
113            return None
114        target = (self._serve_dir / project).resolve()
115 
116        # Prevent path traversal
117        try:
118            target.relative_to(self._serve_dir.resolve())
119        except ValueError:
120            return None
121 
122        if not target.exists() or not target.is_dir():
123            return None
124        return target
125 
126    def _handle_save_chart(self):
127        try:
128            data = self._read_json_body()
129            data_url = data.get("dataUrl", "")
130            filename = _safe_filename(data.get("filename", "screenshot"), ".png")
131            project = data.get("project", "")
132 
133            project_dir = self._resolve_project_dir(project)
134            if project_dir is None:
135                return self._send_json({"ok": False, "error": "Invalid or missing project"}, status=400)
136 
137            if "," in data_url:
138                _, b64 = data_url.split(",", 1)
139            else:
140                b64 = data_url
141            png_bytes = base64.b64decode(b64)
142 
143            out_path = project_dir / filename
144            out_path.write_bytes(png_bytes)
145 
146            rel_path = out_path.relative_to(self._serve_dir).as_posix()
147            self._send_json({"ok": True, "filename": filename, "path": str(out_path), "url": f"/{rel_path}"})
148        except Exception as e:
149            self._send_json({"ok": False, "error": str(e)}, status=500)
150 
151 
152def run(serve_dir, port):
153    os.chdir(serve_dir)
154    server = HTTPServer(("127.0.0.1", port), lambda *a, **k: ChartHandler(*a, serve_dir=serve_dir, **k))
155    print(f"[chart-server] serving {serve_dir} on port {port}", flush=True)
156    server.serve_forever()
157 
158 
159if __name__ == "__main__":
160    serve_dir = sys.argv[1] if len(sys.argv) > 1 else str(DEFAULT_SERVE_DIR)
161    port = int(sys.argv[2]) if len(sys.argv) > 2 else 7860
162    Path(serve_dir).mkdir(parents=True, exist_ok=True)
163    run(serve_dir, port)
164