from __future__ import annotations

import argparse
import base64
import hashlib
import json
import os
import platform
import queue
import re
import subprocess
import sys
import threading
import time
from datetime import datetime
from pathlib import Path
from typing import Any
from urllib.error import HTTPError
from urllib.error import URLError
from urllib.request import Request
from urllib.request import urlopen


ARTIFACT_NAMES = (
    "lstm_forecaster.json",
    "torch_retrain_guard.json",
    "torch_threshold_calibration.json",
)


def main() -> None:
    args = parse_args()
    repo_root = Path(args.repo_root).resolve()
    runtime_dir = repo_root / "runtime"
    runtime_dir.mkdir(parents=True, exist_ok=True)
    log_path = Path(args.log_file).resolve() if args.log_file else runtime_dir / "windows_training_agent.log"

    log(log_path, f"TradeBot Windows training agent started for {args.api_base_url}")
    while True:
        try:
            poll_once(args, repo_root, runtime_dir, log_path)
        except Exception as exc:  # noqa: BLE001 - agent must keep running.
            log(log_path, f"ERROR: {exc}")
        if args.once:
            break
        time.sleep(max(5, args.poll_seconds))


def poll_once(args: argparse.Namespace, repo_root: Path, runtime_dir: Path, log_path: Path) -> None:
    worker = worker_payload(args, repo_root)
    api_json(args, "/api/training/heartbeat", worker)
    claim = api_json(args, "/api/training/claim", worker)
    if not claim.get("claimed"):
        return
    job = claim.get("job") if isinstance(claim.get("job"), dict) else {}
    job_id = str(job.get("id") or "")
    if not job_id:
        return
    log(log_path, f"Claimed retrain job {job_id}")
    report_progress(args, job_id, "running", "claimed", 2, "Задание получено Windows-agent")
    success = False
    message = ""
    summary: dict[str, Any] = {}
    try:
        run_retrain(args, job_id, job, repo_root, log_path)
        summary = read_json(runtime_dir / "torch_retrain_guard.json")
        report_progress(args, job_id, "running", "uploading", 72, "Обучение завершено, загружаю артефакты")
        for name in ARTIFACT_NAMES:
            path = runtime_dir / name
            if path.is_file():
                upload_artifact(args, job_id, path, log_path)
        success = True
        message = "training completed"
        log(log_path, f"Completed retrain job {job_id}")
    except Exception as exc:  # noqa: BLE001 - report failure to the bot.
        message = str(exc)
        log(log_path, f"Job {job_id} failed: {message}")
    finally:
        payload = {"success": success, "message": message, "summary": summary}
        api_json(args, f"/api/training/jobs/{job_id}/complete", payload)


def run_retrain(args: argparse.Namespace, job_id: str, job: dict[str, Any], repo_root: Path, log_path: Path) -> None:
    script = repo_root / "tools" / "run_torch_retrain.ps1"
    if not script.is_file():
        raise RuntimeError(f"retrain script not found: {script}")
    cmd = [
        "powershell.exe",
        "-NoProfile",
        "-ExecutionPolicy",
        "Bypass",
        "-File",
        str(script),
    ]
    parameters = job.get("parameters") if isinstance(job.get("parameters"), dict) else {}
    arg_map = {
        "symbols": "-Symbols",
        "limit": "-Limit",
        "lookbacks": "-Lookbacks",
        "architectures": "-Architectures",
        "hidden_sizes": "-HiddenSizes",
        "layers": "-Layers",
        "dropouts": "-Dropouts",
        "epochs": "-Epochs",
    }
    for key, ps_arg in arg_map.items():
        value = parameters.get(key)
        if value not in (None, ""):
            cmd.extend([ps_arg, str(value)])
    log(log_path, "Running retrain: " + " ".join(quote_for_log(part) for part in cmd))
    report_progress(args, job_id, "running", "training", 8, "PyTorch retrain запущен")
    line_count = 0
    output_queue: queue.Queue[str] = queue.Queue()

    def read_output() -> None:
        assert process.stdout is not None
        for raw_line in process.stdout:
            output_queue.put(raw_line.rstrip())

    with subprocess.Popen(
        cmd,
        cwd=str(repo_root),
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        text=True,
        encoding="utf-8",
        errors="replace",
        **hidden_subprocess_kwargs(),
    ) as process:
        reader = threading.Thread(target=read_output, name="training-output-reader", daemon=True)
        reader.start()
        last_report_at = 0.0
        started_at = time.monotonic()
        last_output_at = started_at
        last_message = "PyTorch retrain выполняется"
        while True:
            got_line = False
            try:
                message = output_queue.get(timeout=5)
                got_line = True
                last_output_at = time.monotonic()
                log(log_path, message)
                line_count += 1
                if message:
                    last_message = friendly_training_message(message)
            except queue.Empty:
                pass

            progress = min(70, 8 + line_count // 3)
            now = time.monotonic()
            if got_line or now - last_report_at >= 30:
                report_message = last_message
                if not got_line:
                    report_message = training_heartbeat_message(now, started_at, last_output_at, last_message)
                safe_report_progress(args, job_id, "running", "training", progress, report_message, log_path)
                last_report_at = now

            if process.poll() is not None and output_queue.empty():
                break

        reader.join(timeout=2)
        code = process.wait()
    if code != 0:
        raise RuntimeError(f"retrain failed with exit code {code}")
    report_progress(args, job_id, "running", "guard", 70, "Guard завершён, подготавливаю артефакты")


def friendly_training_message(message: str) -> str:
    cleaned = message.strip()
    if not cleaned:
        return "PyTorch обучает модель"

    if "Starting PyTorch recurrent retrain:" in cleaned:
        return "PyTorch LSTM/GRU запущен: готовлю данные и варианты модели"

    started = re.search(
        r"training started: symbols=(?P<symbols>\d+) interval=(?P<interval>\d+) "
        r"limit=(?P<limit>\d+) epochs=(?P<epochs>\d+)",
        cleaned,
    )
    if started:
        interval = started.group("interval")
        timeframe = "1h" if interval == "60" else f"{interval}m"
        return (
            f"Старт обучения: {started.group('symbols')} пар, таймфрейм {timeframe}, "
            f"история {started.group('limit')} свечей, до {started.group('epochs')} эпох"
        )

    pair_started = re.search(r"^(?P<symbol>[A-Z0-9]+): training started \((?P<index>\d+)/(?P<total>\d+)\)", cleaned)
    if pair_started:
        return (
            f"{pair_started.group('symbol')}: обучение пары "
            f"{pair_started.group('index')}/{pair_started.group('total')}"
        )

    preparing = re.search(r"^(?P<symbol>[A-Z0-9]+): preparing lookback=(?P<lookback>\d+)", cleaned)
    if preparing:
        return f"{preparing.group('symbol')}: готовлю окно {preparing.group('lookback')} свечей"

    fitting = re.search(
        r"^(?P<symbol>[A-Z0-9]+): fitting (?P<arch>lstm|gru) "
        r"lookback=(?P<lookback>\d+) hidden=(?P<hidden>\d+) "
        r"layers=(?P<layers>\d+) dropout=(?P<dropout>[0-9.]+)",
        cleaned,
    )
    if fitting:
        return (
            f"{fitting.group('symbol')}: обучаю {fitting.group('arch').upper()}, "
            f"окно {fitting.group('lookback')}, нейронов {fitting.group('hidden')}, "
            f"слоёв {fitting.group('layers')}, dropout {fitting.group('dropout')}"
        )

    model = re.search(
        r"^(?P<symbol>[A-Z0-9]+): model=torch_(?P<arch>lstm|gru).*?"
        r"mae=(?P<mae>[0-9.]+)%.*?skill=(?P<skill>-?[0-9.]+).*?dir=(?P<direction>[0-9.]+)",
        cleaned,
    )
    if model:
        direction = float(model.group("direction")) * 100
        skill = float(model.group("skill")) * 100
        return (
            f"{model.group('symbol')}: выбран {model.group('arch').upper()}, "
            f"ошибка {model.group('mae')}%, skill {skill:.1f}%, направление {direction:.1f}%"
        )

    if "Calibrating current artifact" in cleaned:
        return "Проверяю текущую модель на replay"
    if "Calibrating candidate artifact" in cleaned:
        return "Проверяю новую модель на replay"
    if "Running retrain guard" in cleaned:
        return "Gate сравнивает новую модель с текущей"
    if "Candidate rejected by guard" in cleaned:
        return "Новая модель обучилась, но gate не дал ей ходу"
    if "Candidate accepted by guard" in cleaned:
        return "Новая модель прошла gate и стала активной"

    return cleaned[-220:]


def training_heartbeat_message(now: float, started_at: float, last_output_at: float, last_message: str) -> str:
    elapsed = format_duration(now - started_at)
    idle_seconds = max(0.0, now - last_output_at)
    if idle_seconds >= 45:
        return (
            f"PyTorch обучает модель: процесс активен {elapsed}; "
            f"последний лог {format_duration(idle_seconds)} назад: {last_message[:140]}"
        )
    return last_message or f"PyTorch обучает модель: процесс активен {elapsed}"


def format_duration(seconds: float) -> str:
    total_seconds = max(0, int(seconds))
    minutes, seconds_part = divmod(total_seconds, 60)
    hours, minutes_part = divmod(minutes, 60)
    if hours:
        return f"{hours}ч {minutes_part}м"
    if minutes:
        return f"{minutes}м {seconds_part}с"
    return f"{seconds_part}с"


def upload_artifact(args: argparse.Namespace, job_id: str, path: Path, log_path: Path) -> None:
    digest = hashlib.sha256(path.read_bytes()).hexdigest()
    size = path.stat().st_size
    chunk_size = max(64 * 1024, args.chunk_size)
    total = max(1, (size + chunk_size - 1) // chunk_size)
    log(log_path, f"Uploading {path.name}: {size} bytes, {total} chunks")
    with path.open("rb") as source:
        for index in range(total):
            data = source.read(chunk_size)
            payload = {
                "name": path.name,
                "index": index,
                "total": total,
                "sha256": digest,
                "data_base64": base64.b64encode(data).decode("ascii"),
            }
            api_json(args, f"/api/training/jobs/{job_id}/artifacts/chunk", payload, timeout=120)
            if index == 0 or index == total - 1 or index % 10 == 0:
                progress = 72 + int(((index + 1) / total) * 23)
                report_progress(args, job_id, "running", "uploading", progress, f"Загружаю {path.name}: {index + 1}/{total}")


def report_progress(
    args: argparse.Namespace,
    job_id: str,
    status: str,
    phase: str,
    progress_percent: int,
    message: str,
) -> None:
    api_json(
        args,
        f"/api/training/jobs/{job_id}/progress",
        {
            "status": status,
            "phase": phase,
            "progress_percent": progress_percent,
            "message": message,
            "worker": worker_payload(args, Path(args.repo_root).resolve()),
        },
    )


def safe_report_progress(
    args: argparse.Namespace,
    job_id: str,
    status: str,
    phase: str,
    progress_percent: int,
    message: str,
    log_path: Path,
) -> None:
    last_error: Exception | None = None
    for attempt in range(1, 4):
        try:
            report_progress(args, job_id, status, phase, progress_percent, message)
            return
        except Exception as exc:  # noqa: BLE001 - keep the local training process alive.
            last_error = exc
            if attempt < 3:
                time.sleep(attempt * 2)
    log(log_path, f"Temporary progress upload error; training continues: {last_error}")


def api_json(args: argparse.Namespace, path: str, payload: dict[str, Any], timeout: int = 30) -> dict[str, Any]:
    url = args.api_base_url.rstrip("/") + path
    body = json.dumps(payload, ensure_ascii=False).encode("utf-8")
    headers = {"Content-Type": "application/json", "Accept": "application/json"}
    token = args.api_auth or os.environ.get("TRADEBOT_API_AUTH", "")
    headers.update(auth_headers(token))
    request = Request(url, data=body, headers=headers, method="POST")
    try:
        with urlopen(request, timeout=timeout) as response:
            text = response.read().decode("utf-8")
    except HTTPError as exc:
        detail = exc.read().decode("utf-8", errors="replace")
        raise RuntimeError(f"HTTP {exc.code} {path}: {detail[:300]}") from exc
    except URLError as exc:
        raise RuntimeError(f"network error {path}: {exc.reason}") from exc
    return json.loads(text) if text.strip() else {}


def auth_headers(token: str) -> dict[str, str]:
    value = token.strip()
    if not value:
        return {}
    headers = {"X-TradeBot-Token": value}
    if value.lower().startswith(("basic ", "bearer ")):
        headers["Authorization"] = value
    elif ":" in value:
        encoded = base64.b64encode(value.encode("utf-8")).decode("ascii")
        headers["Authorization"] = f"Basic {encoded}"
    else:
        headers["Authorization"] = f"Bearer {value}"
    return headers


def worker_payload(args: argparse.Namespace, repo_root: Path) -> dict[str, Any]:
    name = args.worker_name or platform.node() or "Windows training host"
    return {
        "worker_id": args.worker_id or f"{name}:{repo_root}",
        "name": name,
        "path": str(repo_root),
        "version": "1",
    }


def log(path: Path, message: str) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    stamp = datetime.now().astimezone().isoformat(timespec="seconds")
    line = f"[{stamp}] {message}"
    if sys.stdout is not None:
        try:
            print(line, flush=True)
        except OSError:
            pass
    with path.open("a", encoding="utf-8") as handle:
        handle.write(line + "\n")


def read_json(path: Path) -> dict[str, Any]:
    try:
        data = json.loads(path.read_text(encoding="utf-8"))
    except (OSError, json.JSONDecodeError):
        return {}
    return data if isinstance(data, dict) else {}


def hidden_subprocess_kwargs() -> dict[str, Any]:
    if os.name != "nt":
        return {}
    startupinfo = subprocess.STARTUPINFO()
    startupinfo.dwFlags |= subprocess.STARTF_USESHOWWINDOW
    startupinfo.wShowWindow = 0
    return {
        "creationflags": getattr(subprocess, "CREATE_NO_WINDOW", 0),
        "startupinfo": startupinfo,
    }


def quote_for_log(value: str) -> str:
    return f'"{value}"' if " " in value else value


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Poll TradeBot for retrain jobs and execute them on Windows.")
    parser.add_argument("--api-base-url", default=os.environ.get("TRADEBOT_API_BASE_URL", "https://tb.kusoft.xyz"))
    parser.add_argument("--api-auth", default=os.environ.get("TRADEBOT_API_AUTH", ""))
    parser.add_argument("--repo-root", default=str(Path(__file__).resolve().parents[1]))
    parser.add_argument("--worker-id", default=os.environ.get("TRADEBOT_TRAINING_WORKER_ID", ""))
    parser.add_argument("--worker-name", default=os.environ.get("TRADEBOT_TRAINING_WORKER_NAME", ""))
    parser.add_argument("--poll-seconds", type=int, default=int(os.environ.get("TRADEBOT_TRAINING_POLL_SECONDS", "60")))
    parser.add_argument("--chunk-size", type=int, default=int(os.environ.get("TRADEBOT_TRAINING_CHUNK_SIZE", str(512 * 1024))))
    parser.add_argument("--log-file", default=os.environ.get("TRADEBOT_TRAINING_LOG", ""))
    parser.add_argument("--once", action="store_true")
    return parser.parse_args()


if __name__ == "__main__":
    main()