#!/usr/bin/env python3
"""Agent 循环：LLM 决策 + agent-browser 执行 + 断言 + 报告"""
import os
import threading
import time
import traceback

from browser import AgentBrowser, BrowserError
from config import (TASKS_DIR, DEFAULT_MAX_SNAPSHOT_CHARS, MAX_RETRY_SAME_ERROR,
                    HISTORY_LIMIT)
from db import update_task, save_step_log, load_step_logs
from llm import chat_json, LLMError

SYSTEM_PROMPT = """你是一个专业的网页自动化测试工程师，正在使用浏览器工具执行端到端测试任务。
你的目标：根据用户给的测试目标，实际操作网页完成测试，并通过断言判断测试是否通过。

## 可用的浏览器能力
- click @e1        点击元素（@e1 是快照里的引用）
- fill @e1 "值"    在输入框填入文字（会先清空）
- select @e1 "值"  选择下拉框选项
- press "Enter"    按键（Enter/Tab/Escape/Control+a 等）
- wait "2000"      等待毫秒；wait text:"xx" 等待文本出现；wait url:"/path" 等待 URL 变化
- scroll "down 500" 滚动页面
- screenshot       截图留证（重要步骤请调用）
- assert           断言验证（见下）
- done             测试完成（全部通过）
- fail             测试失败（无法继续或断言不通过）

## 断言类型 assert
{"type":"text","expect":"欢迎","present":true}   页面上应出现"欢迎"文本
{"type":"text","expect":"错误提示","present":false} 页面上不应出现"错误提示"
{"type":"url","expect":"/dashboard","present":true} URL 应包含 /dashboard
{"type":"title","expect":"首页","present":true}    页面标题应包含"首页"
{"type":"element","expect":"登录","present":true}  页面上应存在文本为"登录"的元素
断言执行后会返回 PASS 或 FAIL，根据结果决定下一步。

## 输出格式（严格 JSON，不要输出其他内容）
{
  "reason": "简要中文说明当前判断和下一步计划",
  "action": "click|fill|select|press|wait|scroll|screenshot|assert|done|fail",
  "target": "@e1 或 text:按钮文字",
  "value": "填入的值/按键名/等待参数",
  "assert": {断言对象，action=assert 时必填},
  "summary": "测试结论（action=done/fail 时必填，说明验证了什么）"
}

## 工作原则
1. 先理解页面，再逐步操作；每一步只做一件事。
2. 找不到目标元素时，观察快照选择最接近的 ref，或改用 text: 语义定位。
3. 操作后必须用 assert 验证结果，不要盲目继续。
4. 测试失败(断言 FAIL 且无补救)时用 fail 结束，并说明原因。
5. 全部验证通过后用 done 结束，summary 里总结测试覆盖的内容。
6. 快照可能被截断，必要时用 wait 等待页面加载完成再操作。
"""


def _fmt_snapshot(snap, max_chars=DEFAULT_MAX_SNAPSHOT_CHARS):
    """把快照数据转成紧凑文本给 LLM"""
    lines = []
    refs = (snap or {}).get('refs') or {}
    if isinstance(refs, dict):
        for ref, info in refs.items():
            if isinstance(info, dict):
                name = info.get('name') or info.get('text') or info.get('label') or ''
                role = info.get('role') or ''
                lines.append(f'{ref} [{role}] {name}')
            else:
                lines.append(f'{ref} {info}')
    text = '\n'.join(lines)
    if len(text) > max_chars:
        text = text[:max_chars] + '\n...(快照过长已截断)'
    return text


def _fmt_history(steps):
    """格式化历史步骤给 LLM"""
    out = []
    for s in steps[-HISTORY_LIMIT:]:
        act = s.get('action', '')
        tgt = s.get('target', '')
        val = s.get('value', '')
        res = s.get('result', '')
        detail = s.get('detail', '')
        line = f"[步骤{s.get('n')}] {act} {tgt} {val if val else ''} -> {res}"
        if detail:
            line += f' ({detail[:120]})'
        out.append(line)
    return '\n'.join(out)


class TaskRunner(threading.Thread):
    def __init__(self, task_id, url, goal, max_steps, timeout):
        super().__init__(daemon=True, name=f'task-{task_id}')
        self.task_id = task_id
        self.url = url
        self.goal = goal
        self.max_steps = max_steps
        self.timeout = timeout
        self.stop_flag = threading.Event()
        self.task_dir = os.path.join(TASKS_DIR, task_id)
        os.makedirs(self.task_dir, exist_ok=True)

    def stop(self):
        self.stop_flag.set()

    def _log(self, msg):
        ts = time.strftime('%H:%M:%S')
        print(f'[{ts}][{self.task_id}] {msg}', flush=True)

    def _screenshot(self, browser, name):
        try:
            path = os.path.join(self.task_dir, name)
            browser.screenshot(path)
            return os.path.basename(path)
        except Exception:
            return None

    def run(self):
        update_task(self.task_id, status='running', started_at=time.time(),
                    result='running')
        steps = []
        browser = None
        try:
            browser = AgentBrowser(namespace=f'task-{self.task_id}')
            deadline = time.time() + self.timeout
            self._log(f'打开页面: {self.url}')
            browser.open(self.url, timeout=60)
            browser.wait('--load', 'networkidle', timeout=45)
            self._log('页面已打开')

            step_n = 0
            while step_n < self.max_steps:
                if self.stop_flag.is_set():
                    self._finish(browser, 'stopped', '任务被手动停止', steps)
                    return
                if time.time() > deadline:
                    self._finish(browser, 'error', '任务超时', steps)
                    return

                step_n += 1
                self._log(f'--- 步骤 {step_n}/{self.max_steps} ---')

                # 1. 快照
                try:
                    snap = browser.snapshot(interactive=True, compact=True, timeout=45)
                    snap_text = _fmt_snapshot(snap)
                except BrowserError as e:
                    snap_text = f'(快照失败: {e})'

                cur_url = browser.url()
                cur_title = browser.title()

                # 2. LLM 决策
                user_msg = (
                    f'## 测试目标\n{self.goal}\n\n'
                    f'## 当前页面\nURL: {cur_url}\n标题: {cur_title}\n\n'
                    f'## 页面可交互元素\n{snap_text}\n\n'
                    f'## 已执行步骤\n' + (_fmt_history(steps) if steps else '(尚无)') +
                    f'\n\n请输出下一步动作的 JSON。'
                )
                try:
                    decision = chat_json(
                        [{'role': 'system', 'content': SYSTEM_PROMPT},
                         {'role': 'user', 'content': user_msg}],
                        temperature=0.2, max_tokens=500)
                except LLMError as e:
                    self._log(f'LLM 错误: {e}')
                    self._finish(browser, 'error', f'LLM 决策失败: {e}', steps)
                    return

                action = decision.get('action', '')
                if action not in ('click', 'fill', 'select', 'press', 'wait',
                                  'scroll', 'screenshot', 'assert', 'done', 'fail'):
                    self._log(f'非法动作: {action}')
                    steps.append(self._record(step_n, decision, 'error',
                                              f'非法动作: {action}'))
                    continue

                # 3. 执行动作（带自愈重试）
                result, detail, extra = self._execute(browser, decision)
                self._log(f'动作 {action} -> {result} {detail}')

                if action in ('done', 'fail'):
                    self._finish(browser,
                                 'pass' if action == 'done' else 'fail',
                                 decision.get('summary', detail), steps)
                    return

                step_rec = self._record(step_n, decision, result, detail, extra)
                steps.append(step_rec)

                # 关键步骤截图（截图后统一写日志，避免重复）
                if result == 'ok' and action in ('click', 'fill', 'assert'):
                    shot = self._screenshot(browser, f'step{step_n:02d}.png')
                    if shot:
                        step_rec['screenshot'] = shot
                save_step_log(self.task_id, step_rec)

            self._finish(browser, 'error', f'超过最大步数({self.max_steps})未完成', steps)
        except BrowserError as e:
            self._log(f'浏览器错误: {e}')
            self._finish(browser, 'error', f'浏览器错误: {e}', steps)
        except Exception as e:
            self._log(f'未知异常: {traceback.format_exc()}')
            self._finish(browser, 'error', f'异常: {e}', steps)

    def _record(self, n, decision, result, detail='', extra=None):
        rec = {
            'n': n,
            'action': decision.get('action'),
            'reason': decision.get('reason', ''),
            'target': decision.get('target', ''),
            'value': decision.get('value', ''),
            'assert': decision.get('assert'),
            'result': result,
            'detail': detail,
            'ts': time.strftime('%Y-%m-%d %H:%M:%S'),
        }
        if extra:
            rec.update(extra)
        return rec

    def _execute(self, browser, decision):
        """执行单个动作，带自愈重试。返回 (result, detail, extra)"""
        action = decision.get('action')
        target = decision.get('target', '')
        value = decision.get('value', '')
        assert_obj = decision.get('assert') or {}

        if action == 'done':
            return 'ok', decision.get('summary', '完成'), {}
        if action == 'fail':
            return 'ok', decision.get('summary', '失败'), {}
        if action == 'screenshot':
            shot = self._screenshot(browser, f'shot_{int(time.time())}.png')
            return ('ok', f'已截图 {shot}', {'screenshot': shot}) if shot else \
                   ('error', '截图失败', {})
        if action == 'assert':
            return self._do_assert(browser, assert_obj)

        # 带重试的动作
        last_err = None
        for attempt in range(MAX_RETRY_SAME_ERROR + 1):
            if attempt > 0:
                self._log(f'重试 {attempt}: {action} {target}')
            try:
                if action == 'click':
                    browser.click(self._resolve(target))
                    return 'ok', '点击成功', {}
                elif action == 'fill':
                    browser.fill(self._resolve(target), value)
                    return 'ok', '填入成功', {}
                elif action == 'select':
                    browser.select(self._resolve(target), value)
                    return 'ok', '选择成功', {}
                elif action == 'press':
                    browser.press(value)
                    return 'ok', f'按键 {value}', {}
                elif action == 'wait':
                    if value.startswith('text:'):
                        browser.wait('--text', value[5:])
                    elif value.startswith('url:'):
                        browser.wait('--url', value[4:])
                    elif value.isdigit():
                        browser.wait(value)
                    else:
                        browser.wait(value)
                    return 'ok', f'等待完成', {}
                elif action == 'scroll':
                    parts = value.split()
                    browser.scroll if hasattr(browser, 'scroll') else None
                    args = parts if parts else ['down', '500']
                    self._run_scroll(browser, args)
                    return 'ok', f'滚动 {value}', {}
            except BrowserError as e:
                last_err = str(e)
            # 失败后重新快照让 LLM 换个策略
            try:
                snap = browser.snapshot(interactive=True, compact=True, timeout=40)
                snap_text = _fmt_snapshot(snap, max_chars=5000)
            except Exception:
                snap_text = '(快照失败)'
            retry_dec = chat_json(
                [{'role': 'system', 'content': SYSTEM_PROMPT},
                 {'role': 'user', 'content': (
                     f'## 测试目标\n{self.goal}\n\n'
                     f'## 刚才执行失败\n动作: {action} 目标: {target} 值: {value}\n'
                     f'错误: {last_err}\n\n'
                     f'## 当前页面元素\n{snap_text}\n\n'
                     f'请换一种方式完成相同意图，输出下一步动作 JSON。'
                     f'如果确认无法完成，输出 {{"action":"fail","reason":"...","summary":"..."}}'
                 )}],
                temperature=0.2, max_tokens=400)
            new_action = retry_dec.get('action', '')
            if new_action == 'fail':
                return 'error', f'自愈放弃: {retry_dec.get("summary", last_err)}', {}
            if new_action in ('click', 'fill', 'select', 'press', 'wait', 'scroll'):
                # 更新决策重试
                decision = retry_dec
                target = decision.get('target', target)
                value = decision.get('value', value)
                action = new_action
                continue
            last_err = f'自愈输出非法动作: {new_action}'
        return 'error', f'执行失败: {last_err}', {}

    def _run_scroll(self, browser, args):
        import subprocess
        from config import AGENT_BROWSER, NODE_BIN_DIR, XDG_RUNTIME_DIR
        import os
        env = os.environ.copy()
        env['PATH'] = f'{NODE_BIN_DIR}:{env.get("PATH", "")}'
        env['XDG_RUNTIME_DIR'] = XDG_RUNTIME_DIR
        subprocess.run([AGENT_BROWSER, '--namespace', f'task-{self.task_id}',
                        'scroll'] + args, capture_output=True, timeout=20, env=env,
                       check=True)

    def _resolve(self, target):
        """把目标转成 agent-browser 定位：@ref 原样，text:xxx 转 find 语法"""
        if not target:
            raise BrowserError('缺少目标元素')
        if target.startswith('@'):
            return target
        if target.startswith('text:'):
            return f'find text "{target[5:]}"'
        if target.startswith('role:'):
            parts = target[5:].split(':', 1)
            if len(parts) == 2:
                return f'find role {parts[0]} --name "{parts[1]}"'
            return f'find role {parts[0]}'
        return target

    def _do_assert(self, browser, assert_obj):
        atype = assert_obj.get('type', '')
        expect = assert_obj.get('expect', '')
        present = bool(assert_obj.get('present', True))
        try:
            if atype == 'text':
                found = self._page_has_text(browser, expect)
            elif atype == 'element':
                found = self._element_exists(browser, expect)
            elif atype == 'url':
                cur = browser.url()
                found = expect in cur
                if not found:
                    return ('fail', f'URL 应为包含"{expect}"，实际: {cur}', {})
                return 'ok', f'URL 包含 "{expect}" (PASS)', {}
            elif atype == 'title':
                cur = browser.title()
                found = expect in cur
                if not found:
                    return ('fail', f'标题应包含"{expect}"，实际: {cur}', {})
                return 'ok', f'标题包含 "{expect}" (PASS)', {}
            else:
                return 'error', f'未知断言类型: {atype}', {}
        except BrowserError as e:
            return 'error', f'断言执行出错: {e}', {}

        if found == present:
            tag = 'PASS' if present else 'PASS(确认不存在)'
            return 'ok', f'断言通过: {atype}="{expect}" present={present} ({tag})', {}
        tag = 'FAIL' if present else 'FAIL(不应出现却出现)'
        return 'fail', f'断言失败: {atype}="{expect}" present={present} ({tag})', {}

    def _page_has_text(self, browser, text):
        r = browser.eval_js(
            f'document.body && document.body.innerText.includes({json_dumps(text)})')
        if isinstance(r, dict):
            d = r.get('data', {})
            val = d.get('value') if isinstance(d, dict) else d
            return bool(val)
        return bool(r)

    def _element_exists(self, browser, text):
        r = browser.eval_js(
            f'!![...document.querySelectorAll("button,a,input,textarea,[role=button]")]'
            f'.find(el => (el.innerText||el.value||"").trim().includes({json_dumps(text)}))')
        if isinstance(r, dict):
            d = r.get('data', {})
            val = d.get('value') if isinstance(d, dict) else d
            return bool(val)
        return bool(r)

    def _finish(self, browser, result, summary, steps):
        if browser:
            browser.close()
        update_task(self.task_id, status='finished', result=result,
                    finished_at=time.time(), steps=len(steps), summary=summary)
        report = self._build_report(result, summary, steps)
        rpath = os.path.join(self.task_dir, 'report.html')
        with open(rpath, 'w', encoding='utf-8') as f:
            f.write(report)
        update_task(self.task_id, report_path=f'/api/tasks/{self.task_id}/report')
        self._log(f'完成: result={result} summary={summary}')

    def _build_report(self, result, summary, steps):
        """生成 HTML 报告"""
        title_map = {'pass': '✅ 测试通过', 'fail': '❌ 测试失败',
                     'error': '⚠️ 测试错误', 'stopped': '⏹️ 已停止'}
        color_map = {'pass': '#16a34a', 'fail': '#dc2626',
                     'error': '#d97706', 'stopped': '#6b7280'}
        rows = []
        for s in steps:
            cls = {'ok': 'ok', 'fail': 'bad', 'error': 'bad'}.get(s.get('result'), '')
            shot = ''
            if s.get('screenshot'):
                shot = (f'<div class="shot"><img src="/api/tasks/{self.task_id}/'
                        f'screenshot/{s["screenshot"]}" loading="lazy"></div>')
            rows.append(f'''<tr class="{cls}">
<td>{s.get('n')}</td>
<td>{s.get('action')}</td>
<td>{s.get('target','')} {s.get('value','')}</td>
<td class="reason">{s.get('reason','')}</td>
<td>{s.get('result')} {s.get('detail','')}</td>
<td>{shot}</td></tr>''')
        steps_html = '\n'.join(rows) if rows else '<tr><td colspan="6">无步骤记录</td></tr>'
        dur = ''
        t = get_task_safe(self.task_id)
        if t and t.get('started_at') and t.get('finished_at'):
            dur = f'{t["finished_at"] - t["started_at"]:.1f}s'
        return f'''<!DOCTYPE html>
<html lang="zh-CN"><head><meta charset="utf-8">
<title>{title_map.get(result, result)} - webtest-agent</title>
<style>
body{{font-family:-apple-system,'Segoe UI',sans-serif;margin:0;background:#f3f4f6;color:#111}}
.wrap{{max-width:1100px;margin:0 auto;padding:24px}}
h1{{font-size:22px}}
.badge{{display:inline-block;padding:6px 16px;border-radius:999px;color:#fff;background:{color_map.get(result, '#6b7280')};font-size:14px}}
.meta{{color:#666;font-size:13px;margin:8px 0 20px}}
.summary{{background:#fff;border-radius:10px;padding:16px;margin-bottom:20px;border:1px solid #e5e7eb}}
table{{width:100%;border-collapse:collapse;background:#fff;border-radius:10px;overflow:hidden;font-size:13px}}
th,td{{padding:10px 12px;border-bottom:1px solid #f0f0f0;text-align:left;vertical-align:top}}
th{{background:#f9fafb;font-weight:600}}
tr.bad td{{background:#fef2f2}}
tr.ok td{{background:#f0fdf4}}
.reason{{color:#555;max-width:260px}}
.shot img{{max-width:260px;border-radius:6px;border:1px solid #e5e7eb;display:block}}
</style></head><body><div class="wrap">
<h1>网页测试报告 <span class="badge">{title_map.get(result, result)}</span></h1>
<div class="meta">任务: {self.task_id} | 目标: {self.goal} | 步骤数: {len(steps)} | 耗时: {dur}</div>
<div class="summary"><b>结论：</b>{summary}</div>
<table><thead><tr><th>#</th><th>动作</th><th>目标</th><th>原因</th><th>结果</th><th>截图</th></tr></thead>
<tbody>{steps_html}</tbody></table>
</div></body></html>'''


def json_dumps(s):
    import json
    return json.dumps(str(s))


def get_task_safe(tid):
    from db import get_task
    try:
        return get_task(tid)
    except Exception:
        return None