#!/usr/bin/env python3
from __future__ import annotations

import argparse
import json
import os
import re
import socket
import subprocess
import sys
import time
import tomllib
import uuid
from dataclasses import dataclass
from pathlib import Path
from typing import Optional

import pymysql
from buganalysis_handoff import Outbox, settings as handoff_settings

WORKER_VERSION = "0.4"
ARCHIVE_RE = re.compile(r"^Archive created:\s*(.+?)\s*$", re.MULTILINE)
WOWDIAG_PACKAGE_RE = re.compile(r"^Diagnostic package created:\s*(.+?)\s*$", re.MULTILINE)
WOWSRC_PACKAGE_RE = re.compile(r"^wowsrc package created:\s*(.+?)\s*$", re.MULTILINE)


@dataclass(frozen=True)
class Config:
    host: str
    port: int
    user: str
    password: str
    database: str
    poll_seconds: int
    timeout: int
    python: str
    wowdiag_script: Path
    wowdiag_config: Path
    wowsrc_script: Path
    wowsrc_config: Path
    log_dir: Path
    buganalysis: Optional[dict] = None


@dataclass(frozen=True)
class ToolContext:
    diagnostics_root: Path
    source_context_root: Path
    core_source: Path


@dataclass(frozen=True)
class Job:
    run_id: int
    bug_id: int
    claim_token: str


def load_config(path: Path) -> Config:
    with path.open("rb") as f:
        d = tomllib.load(f)

    db, w, p = d["database"], d.get("worker", {}), d["paths"]
    return Config(
        str(db.get("host", "127.0.0.1")),
        int(db.get("port", 3306)),
        str(db["user"]),
        str(db["password"]),
        str(db.get("database", "vcore_dev")),
        max(1, int(w.get("poll_seconds", 10))),
        max(30, int(w.get("job_timeout_seconds", 900))),
        str(p.get("python", sys.executable)),
        Path(p["wowdiag_script"]).resolve(),
        Path(p["wowdiag_config"]).resolve(),
        Path(p["wowsrc_script"]).resolve(),
        Path(p["wowsrc_config"]).resolve(),
        Path(p["log_dir"]).resolve(),
        handoff_settings(d.get('buganalysis'),path.parent),
    )


def load_tool_context(cfg: Config) -> ToolContext:
    with cfg.wowdiag_config.open("rb") as f:
        diag = tomllib.load(f)
    with cfg.wowsrc_config.open("rb") as f:
        src = tomllib.load(f)

    diag_paths = diag["paths"]
    src_paths = src["paths"]

    diag_core = Path(diag_paths["core_source"]).resolve()
    src_core = Path(src_paths["core_source"]).resolve()
    if diag_core != src_core:
        raise RuntimeError(
            "wowdiag and wowsrc use different core_source paths: "
            f"wowdiag={diag_core}, wowsrc={src_core}"
        )

    return ToolContext(
        diagnostics_root=Path(diag_paths["diagnostics"]).resolve(),
        source_context_root=Path(src_paths["output_root"]).resolve(),
        core_source=diag_core,
    )


def connect(cfg: Config):
    return pymysql.connect(
        host=cfg.host,
        port=cfg.port,
        user=cfg.user,
        password=cfg.password,
        database=cfg.database,
        charset="utf8mb4",
        autocommit=False,
        cursorclass=pymysql.cursors.DictCursor,
    )


def identity() -> str:
    return f"{socket.gethostname()}:{os.getpid()}"


def claim_one(conn, claimed_by: str) -> Optional[Job]:
    token = uuid.uuid4().hex
    try:
        with conn.cursor() as cur:
            cur.execute(
                """
                UPDATE diagnostic_runs
                SET status='running', claim_token=%s, claimed_by=%s,
                    claimed_at=NOW(), started_at=NOW(), error_message=NULL
                WHERE status='queued'
                ORDER BY diagnostic_run_id ASC
                LIMIT 1
                """,
                (token, claimed_by),
            )
            if cur.rowcount != 1:
                conn.rollback()
                return None

        conn.commit()

        with conn.cursor() as cur:
            cur.execute(
                """SELECT diagnostic_run_id, bug_id, claim_token
                   FROM diagnostic_runs
                   WHERE claim_token=%s AND status='running' LIMIT 1""",
                (token,),
            )
            row = cur.fetchone()

        if not row:
            raise RuntimeError("Claim succeeded but claimed row could not be resolved")

        return Job(
            int(row["diagnostic_run_id"]),
            int(row["bug_id"]),
            str(row["claim_token"]),
        )
    except Exception:
        conn.rollback()
        raise


def mark_completed(
    conn,
    job: Job,
    wowdiag_archive: Path,
    wowdiag_version: str,
    core_revision: str,
):
    """
    Keep the existing diagnostic_runs schema contract intact.

    output_path remains the wowdiag archive. The paired wowsrc archive is
    generated in the configured wowsrc output_root and is recorded in the
    worker log. A future DB migration can add dedicated wowsrc columns without
    changing the pipeline execution semantics introduced here.
    """
    with conn.cursor() as cur:
        cur.execute(
            """
            UPDATE diagnostic_runs
            SET status='completed', finished_at=NOW(), output_path=%s,
                error_message=NULL, wowdiag_version=%s, core_revision=%s
            WHERE diagnostic_run_id=%s AND claim_token=%s AND status='running'
            """,
            (
                str(wowdiag_archive),
                wowdiag_version[:64],
                core_revision[:64],
                job.run_id,
                job.claim_token,
            ),
        )
        if cur.rowcount != 1:
            conn.rollback()
            raise RuntimeError("Could not finalize job: claim ownership/status changed")
    conn.commit()


def mark_failed(conn, job: Job, message: str):
    with conn.cursor() as cur:
        cur.execute(
            """
            UPDATE diagnostic_runs
            SET status='failed', finished_at=NOW(), error_message=%s
            WHERE diagnostic_run_id=%s AND claim_token=%s AND status='running'
            """,
            (message[:65535], job.run_id, job.claim_token),
        )
        if cur.rowcount != 1:
            conn.rollback()
            raise RuntimeError("Could not mark job failed: claim ownership/status changed")
    conn.commit()


def validated_path_from_match(
    stdout: str,
    regex: re.Pattern[str],
    root: Path,
    description: str,
    must_be_dir: bool,
) -> Optional[Path]:
    m = regex.search(stdout)
    if not m:
        return None

    p = Path(m.group(1).strip()).resolve()
    try:
        p.relative_to(root)
    except ValueError as exc:
        raise RuntimeError(f"{description} outside configured output root: {p}") from exc

    if must_be_dir:
        if not p.is_dir():
            raise RuntimeError(f"{description} directory does not exist: {p}")
    elif not p.is_file():
        raise RuntimeError(f"{description} file does not exist: {p}")

    return p


def package_from_stdout(
    stdout: str,
    regex: re.Pattern[str],
    root: Path,
    description: str,
) -> Optional[Path]:
    return validated_path_from_match(
        stdout, regex, root, description, must_be_dir=True
    )


def archive_from_stdout(stdout: str, root: Path, description: str) -> Optional[Path]:
    p = validated_path_from_match(
        stdout, ARCHIVE_RE, root, description, must_be_dir=False
    )
    if p and p.suffixes[-2:] != [".tar", ".gz"]:
        raise RuntimeError(f"unexpected archive type: {p}")
    return p


def fallback_archive(root: Path, pattern: str, started: float) -> Optional[Path]:
    found = []
    for p in root.glob(pattern):
        try:
            st = p.stat()
        except OSError:
            continue
        if st.st_mtime >= started - 2:
            found.append((st.st_mtime, p.resolve()))
    return sorted(found, reverse=True)[0][1] if found else None


def manifest_metadata(package_dir: Path, expected_tool: str) -> tuple[str, str]:
    manifest_path = package_dir / "manifest.json"
    if not manifest_path.is_file():
        raise RuntimeError(f"{expected_tool} manifest not found: {manifest_path}")

    try:
        manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
    except Exception as exc:
        raise RuntimeError(f"Could not read {expected_tool} manifest: {exc}") from exc

    tool = str(manifest.get("tool") or "").strip()
    if tool and tool != expected_tool:
        raise RuntimeError(
            f"Unexpected manifest tool in {manifest_path}: {tool!r}"
        )

    version = str(manifest.get("tool_version") or "").strip()
    revision = str(manifest.get("core_revision") or "").strip()

    if not version:
        raise RuntimeError(f"{expected_tool} manifest contains no tool_version")
    if not revision:
        raise RuntimeError(f"{expected_tool} manifest contains no core_revision")

    return version, revision


def run_command(cmd: list[str], cwd: Path, timeout: int) -> tuple[int, str, str]:
    r = subprocess.run(
        cmd,
        cwd=str(cwd),
        capture_output=True,
        text=True,
        timeout=timeout,
        check=False,
    )
    return r.returncode, r.stdout or "", r.stderr or ""


def write_log(
    cfg: Config,
    job: Job,
    diag_cmd: list[str],
    diag_rc: Optional[int],
    diag_out: str,
    diag_err: str,
    src_cmd: Optional[list[str]] = None,
    src_rc: Optional[int] = None,
    src_out: str = "",
    src_err: str = "",
    wowdiag_archive: Optional[Path] = None,
    wowsrc_archive: Optional[Path] = None,
    failure: Optional[BaseException] = None,
):
    cfg.log_dir.mkdir(parents=True, exist_ok=True)
    p = cfg.log_dir / f"DIAG-{job.run_id:06d}.log"

    parts = [
        f"worker_version={WORKER_VERSION}",
        f"diagnostic_run_id={job.run_id}",
        f"bug_id={job.bug_id}",
        f"claim_token={job.claim_token}",
        f"wowdiag_archive={wowdiag_archive or ''}",
        f"wowsrc_archive={wowsrc_archive or ''}",
        "",
        "WOWDIAG COMMAND:",
        repr(diag_cmd),
        f"returncode={diag_rc}",
        "",
        "WOWDIAG STDOUT:",
        diag_out or "",
        "",
        "WOWDIAG STDERR:",
        diag_err or "",
    ]

    if src_cmd is not None:
        parts += [
            "",
            "WOWSRC COMMAND:",
            repr(src_cmd),
            f"returncode={src_rc}",
            "",
            "WOWSRC STDOUT:",
            src_out or "",
            "",
            "WOWSRC STDERR:",
            src_err or "",
        ]

    if failure:
        parts += ["", "ERROR:", str(failure)]

    p.write_text("\n".join(parts), encoding="utf-8")


def execute_pipeline(cfg: Config, ctx: ToolContext, job: Job):
    if not cfg.wowdiag_script.is_file():
        raise RuntimeError(f"wowdiag script not found: {cfg.wowdiag_script}")
    if not cfg.wowdiag_config.is_file():
        raise RuntimeError(f"wowdiag config not found: {cfg.wowdiag_config}")
    if not cfg.wowsrc_script.is_file():
        raise RuntimeError(f"wowsrc script not found: {cfg.wowsrc_script}")
    if not cfg.wowsrc_config.is_file():
        raise RuntimeError(f"wowsrc config not found: {cfg.wowsrc_config}")

    ctx.diagnostics_root.mkdir(parents=True, exist_ok=True)
    ctx.source_context_root.mkdir(parents=True, exist_ok=True)

    diag_cmd = [
        cfg.python,
        str(cfg.wowdiag_script),
        "--config",
        str(cfg.wowdiag_config),
        "bug",
        str(job.bug_id),
    ]

    src_cmd: Optional[list[str]] = None
    diag_out = diag_err = src_out = src_err = ""
    diag_rc: Optional[int] = None
    src_rc: Optional[int] = None
    wowdiag_archive: Optional[Path] = None
    wowsrc_archive: Optional[Path] = None

    try:
        diag_started = time.time()
        diag_rc, diag_out, diag_err = run_command(
            diag_cmd, cfg.wowdiag_script.parent, cfg.timeout
        )
        if diag_rc != 0:
            raise RuntimeError(
                f"wowdiag exited with code {diag_rc}; "
                f"stderr: {diag_err.strip() or '<empty>'}"
            )

        diag_package = package_from_stdout(
            diag_out,
            WOWDIAG_PACKAGE_RE,
            ctx.diagnostics_root,
            "wowdiag package",
        )
        if not diag_package:
            raise RuntimeError(
                "wowdiag exited successfully but no package directory could be identified"
            )

        wowdiag_archive = (
            archive_from_stdout(diag_out, ctx.diagnostics_root, "wowdiag archive")
            or fallback_archive(
                ctx.diagnostics_root,
                f"BUG-{job.bug_id:06d}-*.tar.gz",
                diag_started,
            )
        )
        if not wowdiag_archive:
            raise RuntimeError(
                "wowdiag exited successfully but no archive could be identified"
            )

        wowdiag_version, diag_revision = manifest_metadata(diag_package, "wowdiag")

        # wowsrc deliberately consumes the exact wowdiag archive generated by
        # this same queue job. This prevents it from silently picking up an
        # older package for the same bug.
        src_cmd = [
            cfg.python,
            str(cfg.wowsrc_script),
            "--config",
            str(cfg.wowsrc_config),
            "bug",
            str(job.bug_id),
            "--diag",
            str(wowdiag_archive),
        ]

        src_started = time.time()
        src_rc, src_out, src_err = run_command(
            src_cmd, cfg.wowsrc_script.parent, cfg.timeout
        )
        if src_rc != 0:
            raise RuntimeError(
                f"wowsrc exited with code {src_rc}; "
                f"stderr: {src_err.strip() or '<empty>'}"
            )

        src_package = package_from_stdout(
            src_out,
            WOWSRC_PACKAGE_RE,
            ctx.source_context_root,
            "wowsrc package",
        )
        if not src_package:
            raise RuntimeError(
                "wowsrc exited successfully but no package directory could be identified"
            )

        wowsrc_archive = (
            archive_from_stdout(src_out, ctx.source_context_root, "wowsrc archive")
            or fallback_archive(
                ctx.source_context_root,
                f"WOWSRC-BUG-{job.bug_id:06d}-*.tar.gz",
                src_started,
            )
        )
        if not wowsrc_archive:
            raise RuntimeError(
                "wowsrc exited successfully but no archive could be identified"
            )

        _, src_revision = manifest_metadata(src_package, "wowsrc")
        if diag_revision != src_revision:
            raise RuntimeError(
                "wowdiag and wowsrc were generated from different core revisions: "
                f"wowdiag={diag_revision}, wowsrc={src_revision}"
            )

        write_log(
            cfg,
            job,
            diag_cmd,
            diag_rc,
            diag_out,
            diag_err,
            src_cmd,
            src_rc,
            src_out,
            src_err,
            wowdiag_archive,
            wowsrc_archive,
        )

        return wowdiag_archive, wowsrc_archive, wowdiag_version, diag_revision

    except Exception as exc:
        write_log(
            cfg,
            job,
            diag_cmd,
            diag_rc,
            diag_out,
            diag_err,
            src_cmd,
            src_rc,
            src_out,
            src_err,
            wowdiag_archive,
            wowsrc_archive,
            exc,
        )
        raise


def process_one(cfg: Config, ctx: ToolContext) -> bool:
    conn = connect(cfg)
    try:
        job = claim_one(conn, identity())
        if not job:
            return False

        print(
            f"Claimed DIAG-{job.run_id} for BUG-{job.bug_id} as {identity()}",
            flush=True,
        )

        try:
            diag_archive, src_archive, version, revision = execute_pipeline(
                cfg, ctx, job
            )
            if cfg.buganalysis:
                Outbox(cfg.buganalysis).stage(job.run_id,job.bug_id,diag_archive,src_archive,revision)
            mark_completed(conn, job, diag_archive, version, revision)
            print(
                f"Completed DIAG-{job.run_id}: "
                f"wowdiag={diag_archive} wowsrc={src_archive}",
                flush=True,
            )
        except Exception as exc:
            mark_failed(conn, job, str(exc))
            print(
                f"FAILED DIAG-{job.run_id}: {exc}",
                file=sys.stderr,
                flush=True,
            )

        return True
    finally:
        conn.close()


def main() -> int:
    ap = argparse.ArgumentParser(
        prog="wowdiag-worker",
        description="Vcore diagnostic_runs -> wowdiag -> wowsrc worker",
    )
    ap.add_argument(
        "--config",
        default=str(Path(__file__).with_name("config.toml")),
    )
    ap.add_argument(
        "--once",
        action="store_true",
        help="Process at most one queued job and exit",
    )
    ap.add_argument('--retry-handoffs',action='store_true',help='Retry queued API deliveries only; no capture jobs')
    ap.add_argument('--handoff-status',action='store_true',help='List API delivery status; no capture jobs')
    args = ap.parse_args()

    cfg = load_config(Path(args.config).resolve())
    outbox=Outbox(cfg.buganalysis) if cfg.buganalysis else None
    if args.handoff_status or args.retry_handoffs:
        if not outbox:
            print('BugAnalysis handoff is disabled.',file=sys.stderr);return 1
        if args.retry_handoffs:outbox.flush(force=True,maximum=100)
        print(json.dumps(outbox.status()));return 0
    ctx = load_tool_context(cfg)

    print(f"Vcore diagnostic worker v{WORKER_VERSION} ({identity()})", flush=True)

    if args.once:
        if outbox:outbox.flush()
        if not process_one(cfg, ctx):
            print("No queued diagnostic job found.", flush=True)
        if outbox:outbox.flush()
        return 0

    while True:
        try:
            if outbox:outbox.flush()
            if not process_one(cfg, ctx):
                time.sleep(cfg.poll_seconds)
        except KeyboardInterrupt:
            print("Worker stopped.", flush=True)
            return 0
        except Exception as exc:
            print(f"Worker loop error: {exc}", file=sys.stderr, flush=True)
            time.sleep(cfg.poll_seconds)


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