from __future__ import annotations

import re
from pathlib import Path
from typing import Any

from . import db


ANSI_RE = re.compile(r"\x1b\[[0-?]*[ -/]*[@-~]")
ULTRALYTICS_EPOCH_RE = re.compile(r"^\s*(?P<current>\d+)\s*/\s*(?P<total>\d+)(?=\s)")
EPOCH_WORD_RE = re.compile(
    r"\bepoch\b[^0-9]{0,12}(?P<current>\d+)\s*(?:/|of)\s*(?P<total>\d+)",
    re.IGNORECASE,
)
MMENGINE_EPOCH_RE = re.compile(
    r"Epoch\((?P<stage>train|val|test)\)\s*\[(?P<epoch>\d+)\]\s*\[(?P<current>\d+)\s*/\s*(?P<total>\d+)\]",
    re.IGNORECASE,
)
MMENGINE_ITER_RE = re.compile(
    r"Iter\((?P<stage>train|val|test)\)\s*\[(?P<current>\d+)\s*/\s*(?P<total>\d+)\]",
    re.IGNORECASE,
)
TQDM_PERCENT_RE = re.compile(r"(?P<percent>\d{1,3}(?:\.\d+)?)%\|")
COMMAND_EPOCHS_RE = re.compile(r"(?:--epochs|--max-epochs|max_epochs)\D{0,20}(?P<total>\d+)", re.IGNORECASE)


def strip_ansi(value: str) -> str:
    return ANSI_RE.sub("", value.replace("\r", "\n"))


def progress_from_log_path(log_path: str | Path, status: str | None = None, max_bytes: int = 65536) -> dict[str, Any]:
    return parse_progress(db.log_tail(log_path, max_bytes=max_bytes), status=status)


def parse_progress(log_text: str, status: str | None = None) -> dict[str, Any]:
    text = strip_ansi(log_text or "")
    parsed = _parse_training_progress(text)
    if parsed is None:
        parsed = _status_progress(status)

    if status == "success":
        return {
            **parsed,
            "percent": 100.0,
            "stage": "completed",
            "label": _terminal_label("已完成", parsed),
            "source": "status",
        }
    if status == "failed":
        return {
            **parsed,
            "stage": "failed",
            "label": _terminal_label("失败", parsed),
            "source": parsed.get("source", "status"),
        }
    if status == "cancelled":
        return {
            **parsed,
            "stage": "cancelled",
            "label": _terminal_label("已取消", parsed),
            "source": parsed.get("source", "status"),
        }
    return parsed


def _parse_training_progress(text: str) -> dict[str, Any] | None:
    lines = [line.strip() for line in text.splitlines() if line.strip()]
    if not lines:
        return None
    total_epochs = _find_total_epochs(text)

    for line in reversed(lines[-400:]):
        mmengine_epoch = MMENGINE_EPOCH_RE.search(line)
        if mmengine_epoch:
            epoch = int(mmengine_epoch.group("epoch"))
            current = int(mmengine_epoch.group("current"))
            total = int(mmengine_epoch.group("total"))
            stage = _stage_name(mmengine_epoch.group("stage"))
            if total_epochs:
                percent = ((max(epoch - 1, 0) + current / max(total, 1)) / total_epochs) * 100
                return _progress(
                    percent=percent,
                    label=f"Epoch {epoch}/{total_epochs} · iter {current}/{total}",
                    stage=stage,
                    current=epoch,
                    total=total_epochs,
                    unit="epoch",
                    source="mmengine_epoch",
                )
            return _progress(
                percent=(current / max(total, 1)) * 100,
                label=f"Epoch {epoch} · iter {current}/{total}",
                stage=stage,
                current=current,
                total=total,
                unit="iter",
                source="mmengine_iter_in_epoch",
            )

        epoch_word = EPOCH_WORD_RE.search(line)
        if epoch_word:
            current = int(epoch_word.group("current"))
            total = int(epoch_word.group("total"))
            return _progress(
                percent=(current / max(total, 1)) * 100,
                label=f"Epoch {current}/{total}",
                stage=_stage_from_line(line),
                current=current,
                total=total,
                unit="epoch",
                source="epoch_word",
            )

        ultralytics_epoch = ULTRALYTICS_EPOCH_RE.search(line)
        if ultralytics_epoch:
            current = int(ultralytics_epoch.group("current"))
            total = int(ultralytics_epoch.group("total"))
            if total > 1 or "loss" in line.lower() or "gpu" in line.lower():
                return _progress(
                    percent=(current / max(total, 1)) * 100,
                    label=f"Epoch {current}/{total}",
                    stage=_stage_from_line(line),
                    current=current,
                    total=total,
                    unit="epoch",
                    source="ultralytics_epoch",
                )

        mmengine_iter = MMENGINE_ITER_RE.search(line)
        if mmengine_iter:
            current = int(mmengine_iter.group("current"))
            total = int(mmengine_iter.group("total"))
            return _progress(
                percent=(current / max(total, 1)) * 100,
                label=f"Iter {current}/{total}",
                stage=_stage_name(mmengine_iter.group("stage")),
                current=current,
                total=total,
                unit="iter",
                source="mmengine_iter",
            )

        tqdm = TQDM_PERCENT_RE.search(line)
        if tqdm:
            percent = float(tqdm.group("percent"))
            if 0 <= percent <= 100:
                return _progress(
                    percent=percent,
                    label=f"{percent:g}%",
                    stage=_stage_from_line(line),
                    current=None,
                    total=None,
                    unit="percent",
                    source="tqdm_percent",
                )

    tail = "\n".join(lines[-20:]).lower()
    if "results saved" in tail or "save_dir=" in tail:
        return _progress(
            percent=99.0,
            label="保存结果",
            stage="saving",
            current=None,
            total=None,
            unit=None,
            source="completion_hint",
        )
    if "starting training" in tail or "train:" in tail:
        return _progress(
            percent=None,
            label="训练中",
            stage="training",
            current=None,
            total=total_epochs,
            unit="epoch" if total_epochs else None,
            source="training_hint",
        )
    return None


def _find_total_epochs(text: str) -> int | None:
    values = [int(match.group("total")) for match in COMMAND_EPOCHS_RE.finditer(text)]
    return max(values) if values else None


def _status_progress(status: str | None) -> dict[str, Any]:
    if status == "queued":
        return _progress(0.0, "排队中", "queued", None, None, None, "status")
    if status == "running":
        return _progress(None, "运行中", "running", None, None, None, "status")
    if status == "success":
        return _progress(100.0, "已完成", "completed", None, None, None, "status")
    if status == "failed":
        return _progress(None, "失败", "failed", None, None, None, "status")
    if status == "cancelled":
        return _progress(None, "已取消", "cancelled", None, None, None, "status")
    return _progress(None, "未开始", "unknown", None, None, None, "status")


def _progress(
    percent: float | None,
    label: str,
    stage: str,
    current: int | None,
    total: int | None,
    unit: str | None,
    source: str,
) -> dict[str, Any]:
    if percent is not None:
        percent = round(max(0.0, min(100.0, float(percent))), 2)
    return {
        "percent": percent,
        "label": label,
        "stage": stage,
        "current": current,
        "total": total,
        "unit": unit,
        "source": source,
    }


def _terminal_label(prefix: str, parsed: dict[str, Any]) -> str:
    label = parsed.get("label")
    if label and label not in {prefix, "运行中", "排队中", "未开始"}:
        return f"{prefix} · {label}"
    return prefix


def _stage_from_line(line: str) -> str:
    lower = line.lower()
    if "val" in lower or "validat" in lower:
        return "validation"
    if "predict" in lower or "infer" in lower:
        return "prediction"
    if "save" in lower:
        return "saving"
    return "training"


def _stage_name(value: str) -> str:
    lower = value.lower()
    if lower == "val":
        return "validation"
    if lower == "test":
        return "evaluation"
    return "training"