#!/usr/bin/env python3
"""Run live FastGPT workflow evaluations against the FastAPI service."""

import argparse
import json
import os
import re
import sys
from dataclasses import dataclass, field
from datetime import datetime, timezone
from pathlib import Path
from uuid import uuid4

import httpx
import yaml


EVALS_DIR = Path(__file__).resolve().parent
DEFAULT_CASES_DIR = EVALS_DIR / "cases"
DEFAULT_OUTPUT_DIR = EVALS_DIR / "results"


@dataclass
class ChatResult:
    stage_codes: list[str]
    errors: list[dict]
    text: str = ""


@dataclass
class TranscriptTurn:
    step: int
    input: str
    expected_stage_code: str
    actual_stage_codes: list[str]
    reply: str
    errors: list[dict]
    ok: bool
    error: str | None = None


@dataclass
class CaseTranscript:
    name: str
    path: str
    client_mode: str
    session_id: str
    passed: bool
    turns: list[TranscriptTurn] = field(default_factory=list)
    warnings: list[str] = field(default_factory=list)


def timestamp() -> str:
    return datetime.now(timezone.utc).isoformat()


def local_run_stamp() -> str:
    return datetime.now().strftime("%Y%m%d-%H%M%S")


def load_cases(cases_dir: Path, name_filters: list[str]) -> list[dict]:
    cases = []
    for path in sorted(cases_dir.rglob("*.yaml")):
        case = yaml.safe_load(path.read_text(encoding="utf-8"))
        if not isinstance(case, dict):
            raise ValueError(f"{path}: case must be a YAML object")

        name = case.get("name")
        client_mode = case.get("clientMode")
        steps = case.get("steps")
        if not isinstance(name, str) or not name:
            raise ValueError(f"{path}: name is required")
        if client_mode not in {"direct", "browser_addon"}:
            raise ValueError(f"{path}: unsupported clientMode={client_mode!r}")
        if not isinstance(steps, list) or not steps:
            raise ValueError(f"{path}: at least one step is required")

        for index, step in enumerate(steps, start=1):
            expected_code = (
                step.get("expect", {}).get("stageCode")
                if isinstance(step, dict)
                else None
            )
            if not isinstance(step, dict) or not isinstance(step.get("input"), str):
                raise ValueError(f"{path}: step {index} input must be a string")
            if not isinstance(expected_code, str) or not expected_code:
                raise ValueError(
                    f"{path}: step {index} expect.stageCode is required"
                )

        if name_filters and not any(value in name for value in name_filters):
            continue

        case["_path"] = str(path.relative_to(EVALS_DIR))
        cases.append(case)

    if not cases:
        raise ValueError(f"No evaluation cases found under {cases_dir}")
    return cases


def parse_sse(response: httpx.Response) -> ChatResult:
    stage_codes: list[str] = []
    errors: list[dict] = []
    text_parts: list[str] = []
    event_name = ""
    data_lines: list[str] = []

    def consume_event() -> None:
        nonlocal event_name, data_lines
        if data_lines:
            payload = json.loads("\n".join(data_lines))
            if event_name == "stage_code":
                stage_codes.append(payload["nextStageCode"])
            elif event_name == "text_delta":
                chunk = payload.get("text")
                if isinstance(chunk, str) and chunk:
                    text_parts.append(chunk)
            elif event_name == "error":
                errors.append(payload)
        event_name = ""
        data_lines = []

    for line in response.iter_lines():
        if not line:
            consume_event()
        elif line.startswith("event:"):
            event_name = line.removeprefix("event:").strip()
        elif line.startswith("data:"):
            data_lines.append(line.removeprefix("data:").strip())

    consume_event()
    return ChatResult(
        stage_codes=stage_codes,
        errors=errors,
        text="".join(text_parts),
    )


def run_chat_step(
    client: httpx.Client,
    *,
    session_id: str,
    client_mode: str,
    text: str,
) -> ChatResult:
    with client.stream(
        "POST",
        "/chat",
        params={"stream": "true"},
        json={
            "sessionId": session_id,
            "timeStamp": timestamp(),
            "text": text,
            "clientMode": client_mode,
            "needFormUpdate": False,
            "useTextChunk": False,
        },
    ) as response:
        response.raise_for_status()
        return parse_sse(response)


def cleanup_session(client: httpx.Client, session_id: str) -> str | None:
    try:
        response = client.request(
            "DELETE",
            "/delete_session",
            json={"sessionId": session_id, "timeStamp": timestamp()},
        )
        response.raise_for_status()
        payload = response.json()
        if payload.get("code") != "200":
            return f"cleanup response={payload}"
    except (httpx.HTTPError, ValueError) as exc:
        return f"cleanup error={exc}"
    return None


def session_id_for(case_name: str) -> str:
    slug = re.sub(r"[^a-zA-Z0-9_-]+", "-", case_name).strip("-")
    return f"eval-{slug[:24]}-{uuid4().hex[:12]}"


def print_turn(turn: TranscriptTurn) -> None:
    status = "OK" if turn.ok else "FAIL"
    print(f"  [{status}] step {turn.step}")
    print(f"    user: {turn.input}")
    if turn.reply:
        print(f"    bot:  {turn.reply}")
    else:
        print("    bot:  <empty>")
    print(
        f"    stage: expected={turn.expected_stage_code} "
        f"actual={turn.actual_stage_codes}"
    )
    if turn.error:
        print(f"    error: {turn.error}")
    if turn.errors:
        print(f"    sse_errors: {turn.errors}")


def transcript_to_markdown(transcript: CaseTranscript) -> str:
    lines = [
        f"# {transcript.name}",
        "",
        f"- path: `{transcript.path}`",
        f"- clientMode: `{transcript.client_mode}`",
        f"- sessionId: `{transcript.session_id}`",
        f"- result: `{'PASS' if transcript.passed else 'FAIL'}`",
        "",
    ]
    for turn in transcript.turns:
        status = "OK" if turn.ok else "FAIL"
        lines.extend(
            [
                f"## Step {turn.step} ({status})",
                "",
                f"**User:** {turn.input}",
                "",
                f"**Bot:** {turn.reply or '<empty>'}",
                "",
                f"- expected stageCode: `{turn.expected_stage_code}`",
                f"- actual stageCodes: `{turn.actual_stage_codes}`",
            ]
        )
        if turn.error:
            lines.append(f"- error: `{turn.error}`")
        if turn.errors:
            lines.append(f"- sse_errors: `{turn.errors}`")
        lines.append("")
    if transcript.warnings:
        lines.append("## Warnings")
        lines.append("")
        for warning in transcript.warnings:
            lines.append(f"- {warning}")
        lines.append("")
    return "\n".join(lines)


def transcript_to_json(transcript: CaseTranscript) -> dict:
    return {
        "name": transcript.name,
        "path": transcript.path,
        "clientMode": transcript.client_mode,
        "sessionId": transcript.session_id,
        "passed": transcript.passed,
        "warnings": transcript.warnings,
        "turns": [
            {
                "step": turn.step,
                "input": turn.input,
                "expectedStageCode": turn.expected_stage_code,
                "actualStageCodes": turn.actual_stage_codes,
                "reply": turn.reply,
                "errors": turn.errors,
                "ok": turn.ok,
                "error": turn.error,
            }
            for turn in transcript.turns
        ],
    }


def safe_case_filename(name: str) -> str:
    slug = re.sub(r"[^a-zA-Z0-9_-]+", "-", name).strip("-")
    return slug or "case"


def save_transcript(output_dir: Path, transcript: CaseTranscript) -> Path:
    output_dir.mkdir(parents=True, exist_ok=True)
    base = safe_case_filename(transcript.name)
    json_path = output_dir / f"{base}.json"
    md_path = output_dir / f"{base}.md"
    json_path.write_text(
        json.dumps(transcript_to_json(transcript), ensure_ascii=False, indent=2),
        encoding="utf-8",
    )
    md_path.write_text(transcript_to_markdown(transcript), encoding="utf-8")
    return md_path


def run_case(client: httpx.Client, case: dict) -> CaseTranscript:
    session_id = session_id_for(case["name"])
    transcript = CaseTranscript(
        name=case["name"],
        path=case["_path"],
        client_mode=case["clientMode"],
        session_id=session_id,
        passed=True,
    )

    print(
        f"\n=== {case['name']} clientMode={case['clientMode']} "
        f"({case['_path']}) session={session_id} ==="
    )

    try:
        for step_number, step in enumerate(case["steps"], start=1):
            expected = step["expect"]["stageCode"]
            user_input = step["input"]
            try:
                result = run_chat_step(
                    client,
                    session_id=session_id,
                    client_mode=case["clientMode"],
                    text=user_input,
                )
            except (httpx.HTTPError, json.JSONDecodeError, KeyError) as exc:
                turn = TranscriptTurn(
                    step=step_number,
                    input=user_input,
                    expected_stage_code=expected,
                    actual_stage_codes=[],
                    reply="",
                    errors=[],
                    ok=False,
                    error=str(exc),
                )
                transcript.turns.append(turn)
                transcript.passed = False
                print_turn(turn)
                break

            ok = result.stage_codes == [expected]
            turn = TranscriptTurn(
                step=step_number,
                input=user_input,
                expected_stage_code=expected,
                actual_stage_codes=result.stage_codes,
                reply=result.text,
                errors=result.errors,
                ok=ok,
                error=None if ok else "stageCode mismatch",
            )
            transcript.turns.append(turn)
            print_turn(turn)
            if not ok:
                transcript.passed = False
                break
    finally:
        cleanup_warning = cleanup_session(client, session_id)
        if cleanup_warning:
            transcript.warnings.append(cleanup_warning)
            print(f"  warning: {cleanup_warning}")

    return transcript


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Run live workflow evaluations and assert stage codes."
    )
    parser.add_argument(
        "--base-url",
        default=os.getenv("EVAL_BASE_URL", "http://127.0.0.1:8000"),
        help="FastAPI base URL (default: EVAL_BASE_URL or http://127.0.0.1:8000)",
    )
    parser.add_argument(
        "--cases-dir",
        type=Path,
        default=DEFAULT_CASES_DIR,
        help="Directory containing YAML cases",
    )
    parser.add_argument(
        "--case",
        action="append",
        default=[],
        help="Run cases whose names contain this value; repeatable",
    )
    parser.add_argument(
        "--output-dir",
        type=Path,
        default=None,
        help=(
            "Directory for conversation transcripts "
            f"(default: {DEFAULT_OUTPUT_DIR}/<timestamp>)"
        ),
    )
    parser.add_argument(
        "--no-save",
        action="store_true",
        help="Print dialogue live but do not write transcript files",
    )
    parser.add_argument("--timeout", type=float, default=90.0)
    parser.add_argument(
        "--insecure",
        action="store_true",
        help="Disable TLS certificate verification",
    )
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    try:
        cases = load_cases(args.cases_dir, args.case)
    except (OSError, ValueError, yaml.YAMLError) as exc:
        print(f"ERROR: {exc}", file=sys.stderr)
        return 2

    output_dir: Path | None = None
    if not args.no_save:
        output_dir = args.output_dir or (DEFAULT_OUTPUT_DIR / local_run_stamp())
        output_dir.mkdir(parents=True, exist_ok=True)
        print(f"Transcripts will be saved under: {output_dir}")

    passed_count = 0
    failed_count = 0
    with httpx.Client(
        base_url=args.base_url.rstrip("/"),
        timeout=args.timeout,
        verify=not args.insecure,
    ) as client:
        for case in cases:
            transcript = run_case(client, case)
            status = "PASS" if transcript.passed else "FAIL"
            print(f"{status} {transcript.name} clientMode={transcript.client_mode}")
            if output_dir is not None:
                saved = save_transcript(output_dir, transcript)
                print(f"  saved: {saved}")
            if transcript.passed:
                passed_count += 1
            else:
                failed_count += 1

    print(
        f"\nSummary: total={len(cases)} "
        f"passed={passed_count} failed={failed_count}"
    )
    if output_dir is not None:
        print(f"Transcripts: {output_dir}")
    return 1 if failed_count else 0


if __name__ == "__main__":
    raise SystemExit(main())