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

import hashlib
import json
import os
import sys
import tempfile
import time
from dataclasses import dataclass
from datetime import datetime
from pathlib import Path
from typing import TYPE_CHECKING, Mapping

from runtime_paths import runtime_dir

if TYPE_CHECKING:
    from health_client import AccountSnapshot

STATE_DIR_ENV = "SYSTRAY_STATUSLINE_STATE_DIR"
ACCOUNT_SLUG_ENV = "SYSTRAY_CODEX_ACCOUNT_SLUG"
ACCOUNT_HOME_ENV = "SYSTRAY_CODEX_ACCOUNT_HOME"
CLAUDE_ACCOUNT_SLUG_ENV = "SYSTRAY_CLAUDE_ACCOUNT_SLUG"
CLAUDE_ACCOUNT_HOME_ENV = "SYSTRAY_CLAUDE_ACCOUNT_HOME"
CLAUDE_ACCOUNT_HOME_NAME = "CLAUDE_HOME"
CLAUDE_ACCOUNTS_DIR_NAME = "claude-accounts"
STALE_AFTER_SECONDS = 15 * 60
RESET = "\033[00m"
DIM = "\033[02m"
GREEN = "\033[01;32m"
BLUE = "\033[01;34m"
CYAN = "\033[00;36m"
YELLOW = "\033[00;33m"
MAGENTA = "\033[00;35m"
BOLD_RED = "\033[01;31m"
ACTIVITY_STALE_SECONDS = 3600


def claude_accounts_dir() -> Path:
    """Read at call time: `runtime_dir()` follows an environment variable, and a
    module-level constant freezes whatever it was at import."""
    return runtime_dir() / CLAUDE_ACCOUNTS_DIR_NAME


def _percentage_color(used: int) -> str:
    if used >= 90:
        return "\033[01;31m"
    if used >= 75:
        return "\033[00;31m"
    if used >= 60:
        return YELLOW
    if used >= 40:
        return "\033[01;33m"
    return "\033[00;32m"


@dataclass(frozen=True, slots=True)
class SessionMetadata:
    turn_count: int
    first_activity: float | None
    last_activity: float | None


@dataclass(frozen=True, slots=True)
class SessionPin:
    session_id: str
    account_slug: str
    account_home: Path
    pinned_at: float


def _state_dir(state_dir: Path | None = None) -> Path:
    return state_dir or Path(os.environ.get(STATE_DIR_ENV, runtime_dir() / "statusline"))


def _pin_path(state_dir: Path, session_id: str) -> Path:
    digest = hashlib.sha256(session_id.encode("utf-8")).hexdigest()
    return state_dir / "session-pins" / f"{digest}.json"


def _atomic_write(path: Path, payload: Mapping[str, object]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True, mode=0o700)
    os.chmod(path.parent, 0o700)
    with tempfile.NamedTemporaryFile("w", encoding="utf-8", dir=path.parent, delete=False) as handle:
        json.dump(payload, handle, separators=(",", ":"))
        handle.write("\n")
        temporary = Path(handle.name)
    os.chmod(temporary, 0o600)
    os.replace(temporary, path)


def write_session_pin(state_dir: Path, pin: SessionPin) -> None:
    _atomic_write(
        _pin_path(state_dir, pin.session_id),
        {
            "schema_version": 1,
            "session_id": pin.session_id,
            "account_slug": pin.account_slug,
            "account_home": str(pin.account_home),
            "pinned_at": pin.pinned_at,
        },
    )


def _read_json(path: Path) -> dict[str, object] | None:
    try:
        payload = json.loads(path.read_text(encoding="utf-8"))
    except (OSError, json.JSONDecodeError):
        return None
    return payload if isinstance(payload, dict) else None


def read_session_pin(state_dir: Path, session_id: str) -> SessionPin | None:
    payload = _read_json(_pin_path(state_dir, session_id))
    if payload is None:
        return None
    slug = payload.get("account_slug")
    account_home = payload.get("account_home")
    pinned_at = payload.get("pinned_at")
    if (
        payload.get("session_id") != session_id
        or not isinstance(slug, str)
        or not slug
        or not isinstance(account_home, str)
        or not account_home
        or not isinstance(pinned_at, (int, float))
    ):
        return None
    return SessionPin(session_id, slug, Path(account_home), float(pinned_at))


def _default_pin(session_id: str, now: float) -> SessionPin | None:
    base_dir = runtime_dir()
    try:
        slug = (base_dir / "default_slug").read_text(encoding="utf-8").strip()
    except OSError:
        return None
    account_home = base_dir / "accounts" / slug / "CODEX_HOME"
    if not slug or not account_home.is_dir():
        return None
    return SessionPin(session_id, slug, account_home, now)


def pin_session(
    payload: Mapping[str, object],
    environment: Mapping[str, str] | None = None,
    state_dir: Path | None = None,
    *,
    now: float | None = None,
) -> SessionPin | None:
    session_id = payload.get("session_id")
    if not isinstance(session_id, str) or not session_id:
        return None
    directory = _state_dir(state_dir)
    existing = read_session_pin(directory, session_id)
    current_time = time.time() if now is None else now
    environment = os.environ if environment is None else environment
    pin: SessionPin | None = None
    slug = environment.get(ACCOUNT_SLUG_ENV)
    account_home_value = environment.get(ACCOUNT_HOME_ENV)
    if isinstance(slug, str) and slug and isinstance(account_home_value, str):
        account_home = Path(account_home_value)
        if account_home.is_dir():
            if existing is not None and existing.account_slug == slug:
                return existing
            pin = SessionPin(session_id, slug, account_home, current_time)
            write_session_pin(directory, pin)
            return pin
    if existing is not None:
        return existing
    pin = _default_pin(session_id, current_time)
    if pin is not None:
        write_session_pin(directory, pin)
    return pin


def claude_account(environment: Mapping[str, str] | None = None) -> tuple[str, Path] | None:
    environment = os.environ if environment is None else environment
    slug = environment.get(CLAUDE_ACCOUNT_SLUG_ENV) or ""
    account_home_value = (
        environment.get(CLAUDE_ACCOUNT_HOME_ENV) or environment.get("CLAUDE_CONFIG_DIR") or ""
    )
    account_home = Path(account_home_value) if account_home_value else None
    if not slug and account_home is not None and account_home.name == CLAUDE_ACCOUNT_HOME_NAME:
        slug = account_home.parent.name
    if account_home is None and slug:
        account_home = claude_accounts_dir() / slug / CLAUDE_ACCOUNT_HOME_NAME
    if not slug or account_home is None or not account_home.is_dir():
        return None
    return slug, account_home


def ingest_claude_usage(
    payload: Mapping[str, object], environment: Mapping[str, str] | None = None
) -> bool:
    account = claude_account(environment)
    rate_limits = payload.get("rate_limits")
    if account is None or not isinstance(rate_limits, dict):
        return False
    account_home = account[1]
    mapped: dict[str, object] = {}
    for source, target in (("five_hour", "five_hour"), ("seven_day", "seven_day")):
        window = rate_limits.get(source)
        if not isinstance(window, dict):
            continue
        used = window.get("used_percentage")
        resets_at = window.get("resets_at")
        if isinstance(used, bool) or not isinstance(used, (int, float)):
            continue
        mapped[target] = {"percent": used, "resets_at": resets_at}
    if not mapped:
        return False
    from claude_health_client import ClaudeHealthClient

    ClaudeHealthClient._write_rate_limits_payload(account_home, mapped)
    return True


def codex_snapshot(slug: str) -> AccountSnapshot | None:
    from command_router import _tool_health_snapshots
    from health_store import HealthSnapshotStore

    path = runtime_dir() / "health_cache.json"
    snapshots = HealthSnapshotStore(path, stale_after_s=0).read()
    return _tool_health_snapshots("codex", snapshots).get(slug)


def _format_remaining(resets_at: object, now: float) -> str | None:
    if not isinstance(resets_at, (int, float)):
        return None
    seconds = max(0, int(float(resets_at) - now))
    if seconds == 0:
        return "0m"
    if seconds < 3600:
        return f"{max(1, (seconds + 59) // 60)}m"
    if seconds < 86400:
        hours, minutes = divmod(seconds, 3600)
        return f"{hours}h" if minutes < 60 else f"{hours}h{minutes // 60}m"
    days, remainder = divmod(seconds, 86400)
    return f"{days}d" if remainder < 3600 else f"{days}d{remainder // 3600}h"


def _format_window(label: str, window: object, now: float, stale: bool) -> str | None:
    if not isinstance(window, dict):
        return None
    used = window.get("used_percentage")
    remaining = _format_remaining(window.get("resets_at"), now)
    if not isinstance(used, int) or not 0 <= used <= 100 or remaining is None:
        return None
    suffix = f" ({remaining} left)"
    if stale:
        suffix = f" ({remaining} left; stale)"
    return f"{MAGENTA}{label}:{_percentage_color(used)}{used}%{RESET}{suffix}"


def _iso_timestamp(value: object) -> float | None:
    if not isinstance(value, str):
        return None
    try:
        return datetime.fromisoformat(value.replace("Z", "+00:00")).timestamp()
    except ValueError:
        return None


def _session_metadata(payload: Mapping[str, object]) -> SessionMetadata:
    transcript_path = payload.get("transcript_path")
    if not isinstance(transcript_path, str) or not transcript_path:
        return SessionMetadata(0, None, None)
    prompt_ids: set[str] = set()
    activities: list[float] = []
    try:
        with Path(transcript_path).open(encoding="utf-8") as transcript:
            for line in transcript:
                try:
                    record = json.loads(line)
                except json.JSONDecodeError:
                    continue
                if not isinstance(record, dict):
                    continue
                if record.get("type") == "user" and not record.get("toolUseResult"):
                    prompt_id = record.get("promptId")
                    if isinstance(prompt_id, str) and prompt_id:
                        prompt_ids.add(prompt_id)
                if record.get("type") == "assistant":
                    timestamp = _iso_timestamp(record.get("timestamp"))
                    if timestamp is not None:
                        activities.append(timestamp)
    except OSError:
        return SessionMetadata(0, None, None)
    return SessionMetadata(
        len(prompt_ids),
        min(activities) if activities else None,
        max(activities) if activities else None,
    )


def _configured_effort(payload: Mapping[str, object]) -> str:
    effort_field = payload.get("effort")
    effort = effort_field.get("level") if isinstance(effort_field, dict) else None
    if not isinstance(effort, str) or not effort:
        settings = _read_json(Path.home() / ".claude" / "settings.json")
        effort = settings.get("effortLevel") if settings is not None else None
    return {"medium": "med"}.get(effort, effort) if isinstance(effort, str) else ""


def _format_duration(seconds: float) -> str:
    total = max(0, int(seconds))
    if total >= 86400:
        days, remainder = divmod(total, 86400)
        hours, remainder = divmod(remainder, 3600)
        return f"{days}d{hours:02d}h{remainder // 60:02d}m"
    if total >= 3600:
        hours, remainder = divmod(total, 3600)
        return f"{hours}h{remainder // 60}m"
    if total >= 60:
        minutes, seconds = divmod(total, 60)
        return f"{minutes}m{seconds}s"
    return f"{total}s"


def _core_line(payload: Mapping[str, object], now: float) -> str:
    model = payload.get("model")
    model_name = model.get("display_name") if isinstance(model, dict) else None
    if not isinstance(model_name, str) or not model_name:
        model_name = model.get("id") if isinstance(model, dict) else ""
    effort = _configured_effort(payload)
    effort_label = f"({effort})" if effort else ""
    workspace = payload.get("workspace")
    cwd = workspace.get("current_dir") if isinstance(workspace, dict) else ""
    if isinstance(cwd, str):
        cwd = cwd.replace(str(Path.home()), "~", 1)
    else:
        cwd = ""
    parts = [f"{GREEN}{model_name}{effort_label}{RESET}:{BLUE}{cwd}{RESET}"]
    metadata = _session_metadata(payload)
    if metadata.turn_count:
        parts.append(f"{CYAN}{metadata.turn_count}t{RESET}")
    if metadata.last_activity is not None:
        activity_color = BOLD_RED if now - metadata.last_activity > ACTIVITY_STALE_SECONDS else YELLOW
        timestamp = datetime.fromtimestamp(metadata.last_activity).strftime("%H:%M:%S %d/%m")
        parts.append(f"{activity_color}{timestamp}{RESET}")
    if metadata.first_activity is not None and metadata.last_activity is not None:
        parts.append(f"{DIM}{_format_duration(metadata.last_activity - metadata.first_activity)}{RESET}")
    context = payload.get("context_window")
    used = context.get("used_percentage") if isinstance(context, dict) else None
    if isinstance(used, (int, float)):
        rounded = round(used)
        parts.append(f"{YELLOW}ctx:{_percentage_color(rounded)}{rounded}%{RESET}")
    return "  ".join(parts)


def _usage_windows(
    five_hour: object, seven_day: object, now: float, stale: bool
) -> str:
    return " ".join(
        window
        for window in (
            _format_window("5h", five_hour, now, stale),
            _format_window("7d", seven_day, now, stale),
        )
        if window is not None
    )


def _codex_line(payload: Mapping[str, object], directory: Path, now: float) -> str:
    session_id = payload.get("session_id")
    if not isinstance(session_id, str) or not session_id:
        return "codex:unavailable (missing session)"
    pin = read_session_pin(directory, session_id)
    if pin is None:
        return "codex:unavailable (no session pin)"
    if not pin.account_home.is_dir():
        replacement = _default_pin(session_id, now)
        if replacement is None:
            return "codex:unavailable (pinned account missing)"
        write_session_pin(directory, replacement)
        pin = replacement
    snapshot = codex_snapshot(pin.account_slug)
    if snapshot is None:
        return f"codex:{pin.account_slug} usage:unavailable (no health cache)"
    if snapshot.status.value == "broken":
        detail = snapshot.detail or snapshot.status.value
        return f"codex:{pin.account_slug} usage:unavailable ({detail})"
    stale = snapshot.checked_at is None or now - snapshot.checked_at > STALE_AFTER_SECONDS
    windows = {
        label: {"used_percentage": percent, "resets_at": resets_at}
        for label, percent, resets_at in (
            ("five_hour", snapshot.primary_used_pct, snapshot.primary_reset_at),
            ("seven_day", snapshot.secondary_used_pct, snapshot.secondary_reset_at),
        )
        if percent is not None
    }
    available = _usage_windows(windows.get("five_hour"), windows.get("seven_day"), now, stale)
    if not available:
        detail = snapshot.detail or snapshot.status.value
        return f"codex:{pin.account_slug} usage:unavailable ({detail})"
    return f"codex:{pin.account_slug} {available}"


def _claude_line(now: float, environment: Mapping[str, str] | None = None) -> str:
    account = claude_account(environment)
    if account is None:
        return "claude:unavailable (no account)"
    slug, account_home = account
    cache_path = account_home / "rate-limits-cache.json"
    cache = _read_json(cache_path)
    if cache is None:
        return f"claude:{slug} usage:unavailable (no cache)"
    try:
        checked_at = cache_path.stat().st_mtime
        from claude_health_client import ClaudeHealthClient, HealthStatus

        snapshot = ClaudeHealthClient._snapshot_from_payload(
            cache, HealthStatus.OK, checked_at, now
        )
    except (OSError, ImportError, AttributeError, TypeError, ValueError):
        return f"claude:{slug} usage:unavailable (unreadable)"
    windows = {
        label: {"used_percentage": percent, "resets_at": resets_at}
        for label, percent, resets_at in (
            ("five_hour", snapshot.primary_used_pct, snapshot.primary_reset_at),
            ("seven_day", snapshot.secondary_used_pct, snapshot.secondary_reset_at),
        )
        if percent is not None
    }
    stale = now - checked_at > STALE_AFTER_SECONDS
    available = _usage_windows(windows.get("five_hour"), windows.get("seven_day"), now, stale)
    return f"claude:{slug}{f' {available}' if available else ' usage:unavailable'}"


def render(
    payload: Mapping[str, object],
    state_dir: Path | None = None,
    *,
    now: float | None = None,
    environment: Mapping[str, str] | None = None,
) -> str:
    directory = _state_dir(state_dir)
    current_time = time.time() if now is None else now
    core = _core_line(payload, current_time)
    codex = _codex_line(payload, directory, current_time)
    claude = _claude_line(current_time, environment)
    return f"{core}\n{codex}\n{claude}"


def start_session(
    payload: Mapping[str, object],
    environment: Mapping[str, str] | None = None,
    state_dir: Path | None = None,
    *,
    now: float | None = None,
) -> SessionPin | None:
    return pin_session(payload, environment, state_dir, now=now)


def _read_payload() -> dict[str, object] | None:
    try:
        payload = json.load(sys.stdin)
    except (json.JSONDecodeError, OSError):
        return None
    return payload if isinstance(payload, dict) else None


def main(argv: list[str]) -> int:
    if len(argv) != 2 or argv[1] not in {"session-start", "render"}:
        return 2
    payload = _read_payload()
    if payload is None:
        return 1
    ingest_claude_usage(payload)
    if argv[1] == "session-start":
        return 0 if start_session(payload) is not None else 1
    print(render(payload))
    return 0


if __name__ == "__main__":
    raise SystemExit(main(sys.argv))