from flask import Flask, request, jsonify, send_from_directory, send_file
import solver
import os
import io
import zipfile
import json
import tempfile
import shutil

app = Flask(__name__, static_folder='static')


def _check_password():
    """Return (ok, error_response) tuple."""
    pwd = os.environ.get("SOLVER_PASSWORD")
    if not pwd:
        return True, None  # no password set → skip check
    token = request.args.get("token", "") or request.headers.get("X-Solver-Token", "")
    if token == pwd:
        return True, None
    return False, (jsonify({"error": "Unauthorized: invalid or missing token"}), 401)


@app.before_request
def enforce_password():
    """Require ?token=PASSWORD or X-Solver-Token header on write API requests."""
    # Public paths: index page, static assets, read-only API
    if request.path == '/' or request.path.startswith('/static/'):
        return None
    # Read-only endpoints don't need auth
    if request.method == 'GET' and request.path.startswith('/api/'):
        return None
    ok, err = _check_password()
    if not ok:
        return err


@app.before_request
def ensure_solver_ready():
    """Fetch solver files on first request if not present."""
    if not os.path.isdir("/app/solver"):
        url = os.environ.get(
            "SOLVER_REPO_URL",
            "https://github.com/Usami-Mira/CC_Integrater/raw/refs/heads/main/",
        )
        solver.fetch_solver_files(url)


@app.route('/')
def index():
    return send_from_directory('static', 'solver.html')


@app.route('/api/config')
def api_config():
    # Ensure solver files are present before reading config
    if not os.path.isdir(solver.SOLVER_SRC):
        url = os.environ.get(
            "SOLVER_REPO_URL",
            "https://github.com/Usami-Mira/CC_Integrater/raw/refs/heads/main/",
        )
        solver.fetch_solver_files(url)

    default_config = {}
    cfg_path = os.path.join(solver.SOLVER_SRC, "config.json")
    if os.path.exists(cfg_path):
        try:
            with open(cfg_path, "r", encoding="utf-8") as f:
                default_config = json.load(f)
        except Exception as e:
            print(f"[CONFIG] Failed to load config.json: {e}", flush=True)
    else:
        print(f"[CONFIG] config.json not found at {cfg_path}", flush=True)

    return jsonify({
        "solver_url": os.environ.get("SOLVER_REPO_URL", "https://github.com/Usami-Mira/CC_Integrater/raw/refs/heads/main/"),
        "default_config": default_config,
    })


@app.route('/api/config', methods=['POST'])
def api_config_update():
    data = request.json
    if data and 'solver_url' in data:
        url = data['solver_url']
        os.environ["SOLVER_REPO_URL"] = url
        ok = solver.fetch_solver_files(url)
        return jsonify({"status": "ok" if ok else "error", "solver_url": url})
    return jsonify({"error": "Missing solver_url"}), 400


@app.route('/api/submit', methods=['POST'])
def api_submit():
    data = request.json
    problem = data.get('problem', '').strip()
    if not problem:
        return jsonify({'error': 'No problem provided'}), 400
    name = data.get('name', '').strip() or None
    config = data.get('config')
    secret_key = data.get('secret_key')

    task_id = solver.submit_task(problem, name, config, secret_key)
    return jsonify({'task_id': task_id})


@app.route('/api/submit-file', methods=['POST'])
def api_submit_file():
    """Upload one or more problem.md files."""
    files = request.files.getlist('files')
    if not files:
        return jsonify({'error': 'No files uploaded'}), 400
    config = request.form.get('config')
    secret_key = request.form.get('secret_key')
    config_override = None
    if config:
        try:
            config_override = json.loads(config)
        except json.JSONDecodeError:
            return jsonify({'error': 'Invalid config JSON'}), 400

    task_ids = []
    for f in files:
        if not f.filename:
            continue
        text = f.read().decode('utf-8', errors='replace').strip()
        if not text:
            continue
        name = os.path.splitext(f.filename)[0]
        task_id = solver.submit_task(text, name, config_override, secret_key)
        task_ids.append(task_id)

    return jsonify({'task_ids': task_ids, 'count': len(task_ids)})


@app.route('/api/submit-bundle', methods=['POST'])
def api_submit_bundle():
    """
    Upload a complete project bundle (zip) containing solver files + problems.

    POST multipart/form-data:
      - project: zip file containing solver source (run.py, spawn.py, outline.md, config.json)
      - problems: one or more .md problem files (or a zip of problem files)

    Response:
      { "task_ids": [...], "count": N }

    Each problem becomes a separate task using the uploaded solver code.
    """
    project_file = request.files.get('project')
    problem_files = request.files.getlist('problems')
    problem_texts = request.form.getlist('problem')  # raw text problems

    if not project_file and not problem_files and not problem_texts:
        return jsonify({'error': 'No project or problems provided'}), 400

    # Extract solver files from project zip
    solver_dir = None
    if project_file and project_file.filename:
        solver_dir = tempfile.mkdtemp(prefix='solver_')
        try:
            with zipfile.ZipFile(project_file.stream, 'r') as zf:
                zf.extractall(solver_dir)
            # Verify required files exist
            required = ['run.py', 'spawn.py', 'outline.md']
            found = []
            for name in required:
                # Check root and first-level subdirs
                if os.path.exists(os.path.join(solver_dir, name)):
                    found.append(name)
                else:
                    for d in os.listdir(solver_dir):
                        subdir = os.path.join(solver_dir, d)
                        if os.path.isdir(subdir) and os.path.exists(os.path.join(subdir, name)):
                            found.append(name)
                            break
            missing = [r for r in required if r not in found]
            if missing:
                shutil.rmtree(solver_dir)
                return jsonify({'error': f'Missing required files in project zip: {", ".join(missing)}'}), 400

            # Replace the active solver source with uploaded files
            # Find the directory containing the solver files
            src_root = solver_dir
            if os.path.exists(os.path.join(solver_dir, 'run.py')):
                src_root = solver_dir
            else:
                for d in os.listdir(solver_dir):
                    subdir = os.path.join(solver_dir, d)
                    if os.path.isdir(subdir) and os.path.exists(os.path.join(subdir, 'run.py')):
                        src_root = subdir
                        break

            # Copy to /app/solver/ (replacing default)
            if os.path.exists(solver.SOLVER_SRC):
                shutil.rmtree(solver.SOLVER_SRC)
            shutil.copytree(src_root, solver.SOLVER_SRC)
            print(f"[BUNDLE] Replaced solver source with uploaded project", flush=True)
        except Exception as e:
            if solver_dir:
                shutil.rmtree(solver_dir)
            return jsonify({'error': f'Invalid project zip: {e}'}), 400
        finally:
            # Clean up temp dir after copy
            if solver_dir and os.path.exists(solver_dir):
                shutil.rmtree(solver_dir)

    # Submit problems
    task_ids = []

    # From file uploads
    for f in problem_files:
        if not f.filename:
            continue
        text = f.read().decode('utf-8', errors='replace').strip()
        if not text:
            continue
        name = os.path.splitext(f.filename)[0]
        task_id = solver.submit_task(text, name)
        task_ids.append(task_id)

    # From form text fields
    for i, text in enumerate(problem_texts):
        text = text.strip()
        if not text:
            continue
        task_id = solver.submit_task(text, f"problem_{i+1}")
        task_ids.append(task_id)

    if not task_ids:
        return jsonify({'error': 'No valid problems found'}), 400

    return jsonify({'task_ids': task_ids, 'count': len(task_ids)})


@app.route('/api/status/<task_id>')
def api_status(task_id):
    secret_key = request.args.get('key')
    task = solver.get_task(task_id, secret_key)
    if not task:
        return jsonify({'error': 'Task not found'}), 404
    return jsonify(task)


@app.route('/api/result/<task_id>')
def api_result(task_id):
    secret_key = request.args.get('key')
    content = solver.get_task_result(task_id, secret_key)
    if content is None:
        return jsonify({'error': 'Result not available yet'}), 404
    return jsonify({'content': content})


@app.route('/api/logs/<task_id>')
def api_logs(task_id):
    lines = request.args.get('lines', '100')
    secret_key = request.args.get('key')
    return jsonify({'logs': solver.get_task_logs(task_id, int(lines), secret_key)})


@app.route('/api/workspace/<task_id>/files')
def api_workspace_files(task_id):
    """List all .md files in a task's workspace."""
    secret_key = request.args.get('key')
    files = solver.list_workspace_files(task_id, secret_key)
    if not files:
        task = solver.get_task(task_id, secret_key)
        if not task:
            return jsonify({'error': 'Task not found'}), 404
    return jsonify({'files': files})


@app.route('/api/workspace/<task_id>/<filename>')
def api_workspace_file(task_id, filename):
    """Read a specific file from a task's workspace."""
    secret_key = request.args.get('key')
    content = solver.get_workspace_file(task_id, filename, secret_key)
    if content is None:
        return jsonify({'error': 'File not found'}), 404
    return jsonify({'content': content})


@app.route('/api/tasks')
def api_tasks():
    secret_key = request.args.get('key')
    return jsonify({'tasks': solver.list_tasks(secret_key)})


@app.route('/api/download/<task_id>')
def api_download(task_id):
    """Download a single task's workspace as zip (only .md and .json files)."""
    secret_key = request.args.get('key')
    task = solver.get_task(task_id, secret_key)
    if not task:
        return jsonify({'error': 'Task not found'}), 404

    workspace = task.get("workspace", "")
    if not os.path.isdir(workspace):
        return jsonify({'error': 'Workspace not found'}), 404

    buf = io.BytesIO()
    with zipfile.ZipFile(buf, 'w', zipfile.ZIP_DEFLATED) as zf:
        for root, dirs, files in os.walk(workspace):
            for fname in files:
                if fname.startswith('.'):
                    continue
                if not (fname.endswith('.md') or fname.endswith('.json')):
                    continue
                fpath = os.path.join(root, fname)
                arcname = os.path.relpath(fpath, workspace)
                zf.write(fpath, arcname)
    buf.seek(0)
    return send_file(buf, mimetype='application/zip', as_attachment=True,
                     download_name=f"result_{task_id}.zip")


@app.route('/api/download-all')
def api_download_all():
    """Download all completed tasks as a single zip (only .md and .json files)."""
    secret_key = request.args.get('key')
    all_tasks = solver.list_tasks(secret_key)
    buf = io.BytesIO()
    with zipfile.ZipFile(buf, 'w', zipfile.ZIP_DEFLATED) as zf:
        for t in all_tasks:
            if t['status'] != 'complete':
                continue
            ws = solver.get_task(t['task_id'], secret_key)
            if not ws:
                continue
            workspace = ws.get("workspace", "")
            name = t.get("name", t['task_id'])
            for root, dirs, files in os.walk(workspace):
                for fname in files:
                    if fname.startswith('.'):
                        continue
                    if not (fname.endswith('.md') or fname.endswith('.json')):
                        continue
                    fpath = os.path.join(root, fname)
                    arcname = f"{name}/{os.path.relpath(fpath, workspace)}"
                    zf.write(fpath, arcname)
    buf.seek(0)
    return send_file(buf, mimetype='application/zip', as_attachment=True,
                     download_name="all_results.zip")


@app.route('/api/task/<task_id>/config', methods=['PUT'])
def api_update_task_config(task_id):
    """Update config.json for a pending task. Only allowed before the task starts running."""
    data = request.json
    if not data:
        return jsonify({'error': 'Missing config data'}), 400
    secret_key = request.args.get('key')
    task = solver.get_task(task_id, secret_key)
    if not task:
        return jsonify({'error': 'Task not found'}), 404
    config, err = solver.update_task_config(task_id, data)
    if err:
        return jsonify({'error': err}), 400
    return jsonify({'config': config})


@app.route('/api/task/<task_id>', methods=['DELETE'])
def api_delete_task(task_id):
    """Delete a task (remove from memory and delete workspace)."""
    secret_key = request.args.get('key')
    ok = solver.delete_task(task_id, secret_key)
    if not ok:
        return jsonify({'error': 'Task not found'}), 404
    return jsonify({'status': 'deleted', 'task_id': task_id})


@app.route('/api/tasks/cleanup', methods=['POST'])
def api_cleanup_tasks():
    """Delete all completed and errored tasks. Returns count of deleted tasks."""
    count = solver.cleanup_completed_tasks()
    return jsonify({'status': 'ok', 'deleted': count})


@app.route('/api/tasks', methods=['DELETE'])
def api_cleanup_all_tasks():
    """Delete all tasks (alias for cleanup)."""
    count = solver.cleanup_completed_tasks()
    return jsonify({'status': 'ok', 'deleted': count})


INTEGRATION_GUIDE = """# CC Solver API 集成指南

## 基础地址
```
https://your-space.hf.space/api/
```

## 认证
所有请求必须携带密码，二选一：

### 方式 1：URL 参数（推荐，最简单）
```
?token=YOUR_PASSWORD
```

### 方式 2：HTTP 请求头
```
X-Solver-Token: YOUR_PASSWORD
```

如果服务端未设置 `SOLVER_PASSWORD` 环境变量，则跳过认证。

## 1. 提交题目（文本）
```
POST /api/submit?token=YOUR_PASSWORD
Content-Type: application/json

{
  "problem": "物理题目内容...",
  "name": "可选的任务名称"
}
```
返回：`{"task_id": "task-1234567890"}`

## 2. 提交题目（文件上传）
```
POST /api/submit-file
Content-Type: multipart/form-data

files: problem1.md, problem2.md, ...
```
返回：`{"task_ids": ["task-111", "task-222"], "count": 2}`

## 3. 提交项目包（自定义 solver + 题目）
```
POST /api/submit-bundle
Content-Type: multipart/form-data

project: solver-project.zip   (包含 run.py, spawn.py, outline.md, config.json)
problems: problem1.md, problem2.md, ...   (可选，题目文件)
problem: "物理题目文本"                   (可选，直接传文本)
```
返回：`{"task_ids": [...], "count": N}`

## 4. 查询任务状态
```
GET /api/status/{task_id}
```
返回：
```json
{
  "task_id": "task-123",
  "name": "Problem 1",
  "status": "running|complete|error",
  "phase": "starting|planning|building|evaluating|complete",
  "problem_preview": "题目预览...",
  "problem_length": 1234,
  "created_at": "2026-05-14T...",
  "started_at": "2026-05-14T...",
  "finished_at": null,
  "files_done": ["plan", "solution"],
  "error": null
}
```

## 5. 获取结果
```
GET /api/result/{task_id}
```
返回：`{"content": "完整的 final_summary.md 内容"}`

## 6. 查看日志
```
GET /api/logs/{task_id}?lines=100
```
返回：`{"logs": ["[2026-05-14T...] Task created", ...]}`

## 7. 获取所有任务列表
```
GET /api/tasks
```
返回：`{"tasks": [...]}`

## 8. 下载单个任务结果
```
GET /api/download/{task_id}
```
返回：ZIP 文件（包含 workspace 中的用户文件，排除 run.py/spawn.py/outline.md）

## 9. 下载所有已完成任务
```
GET /api/download-all
```
返回：ZIP 文件（包含所有已完成任务的 workspace，排除 solver 源文件）

## 10. 更新任务配置（仅 pending 状态）
```
PUT /api/task/{task_id}/config
Content-Type: application/json

{
  "wallClock": 3600,
  "model": "qwen3.6-plus"
}
```
返回：`{"config": {...}}` — 更新后的完整 config
仅对尚未开始执行的任务有效。

## 11. 删除单个任务
```
DELETE /api/task/{task_id}
```
返回：`{"status": "deleted", "task_id": "task-123"}`
从内存和磁盘中完全删除该任务。

## 12. 清理已完成/出错的任务
```
POST /api/tasks/cleanup
```
返回：`{"status": "ok", "deleted": 5}`
批量删除所有 `complete` 和 `error` 状态的任务，释放空间。

## 状态说明
| status | 含义 |
|--------|------|
| pending | 排队等待 |
| running | 正在解题中 |
| complete | 解题完成，可获取结果 |
| error | 出错 |

## files_done 说明
自动扫描 workspace 中生成的 `.md` 文件（排除 problem.md），文件名去掉 `.md` 后缀。例如：
- `["plan"]` — Planner 已完成
- `["plan", "solution"]` — Builder 已完成
- `["plan", "solution", "review"]` — Evaluator 已完成
- `["plan", "solution", "review", "final_summary"]` — 全部完成

## 轮询建议
每 3~5 秒调用一次 `/api/status/{task_id}`，当 `status` 变为 `complete` 后调用 `/api/result/{task_id}` 获取结果。
"""


@app.route('/api/integration-guide')
def api_integration_guide():
    """Return API documentation for external integrations."""
    return jsonify({
        "content": INTEGRATION_GUIDE,
        "format": "markdown"
    })


@app.route('/api/integration-guide.md')
def api_integration_guide_md():
    """Return API documentation as plain markdown."""
    return INTEGRATION_GUIDE, 200, {'Content-Type': 'text/markdown; charset=utf-8'}


if __name__ == '__main__':
    app.run(host='0.0.0.0', port=7860)