import re
import json
from dataclasses import dataclass, field, asdict
from datetime import datetime
from typing import Optional

# ── Preprocessing ────────────────────────────────────────────────────────────

_TS_RE = re.compile(r'^\[\d{4}-\d{2}-\d{2} \d{2}:\d{2}:\d{2}\]\s*')
_CHILD_RE = re.compile(r'^\[(\d+)\]\s*')


def strip_ts(line: str) -> str:
    """Remove optional `ts`-injected timestamp prefix from a log line."""
    return _TS_RE.sub('', line.strip())


def classify_line(line: str) -> tuple[str, Optional[int], str]:
    """
    Return (kind, port, content) where kind is 'router' or 'child'.
    `line` must already have the ts prefix stripped.
    """
    m = _CHILD_RE.match(line)
    if m:
        return ('child', int(m.group(1)), line[m.end():])
    return ('router', None, line)


# ── Data model ───────────────────────────────────────────────────────────────

@dataclass
class LoadingInfo:
    stages: list
    current: str
    progress: float  # 0.0 – 1.0


@dataclass
class Metrics:
    prompt_speed: Optional[float] = None   # tokens/sec during prompt eval
    prompt_tokens: Optional[int] = None    # total prompt tokens
    gen_speed: Optional[float] = None      # tokens/sec during generation
    n_decoded: Optional[int] = None        # tokens decoded so far


@dataclass
class ServerState:
    state: str = "offline"
    model: Optional[str] = None
    loading: Optional[LoadingInfo] = None
    metrics: Metrics = field(default_factory=Metrics)
    request_count: int = 0

    def to_dict(self) -> dict:
        return {
            "state": self.state,
            "model": self.model,
            "loading": asdict(self.loading) if self.loading else None,
            "metrics": asdict(self.metrics),
            "request_count": self.request_count,
            "timestamp": datetime.now().isoformat(timespec="seconds"),
        }