"""Streaming merge and normalization for per-PID strace file traces."""

from __future__ import annotations

import ast
from collections import Counter
from contextlib import ExitStack
from dataclasses import asdict, dataclass
import heapq
import json
from pathlib import Path
import re
from typing import Iterator, TextIO


_TIMESTAMP_RE = re.compile(r"^(?P<timestamp>[0-9]+(?:\.[0-9]+)?)\s+(?P<body>.*)$")
_SYSCALL_RE = re.compile(r"^(?P<syscall>[A-Za-z_][A-Za-z0-9_]*)\(")
_ERRNO_RE = re.compile(
	r"=\s+-1\s+(?P<errno>E[A-Z0-9]+)\s+\((?P<description>[^)]*)\)"
)
_QUOTED_RE = re.compile(r'"(?:\\.|[^"\\])*"')

_ERROR_CLASSIFICATIONS = {
	"ENOENT": "missing",
	"ENOTDIR": "path-component-not-directory",
	"EACCES": "permission-denied",
	"EPERM": "operation-not-permitted",
	"ELOOP": "symlink-loop",
	"ENAMETOOLONG": "name-too-long",
	"EROFS": "read-only-filesystem",
	"EIO": "io-error",
}


@dataclass(frozen=True)
class TraceSummary:
	shards: int
	lines: int
	normalized_records: int
	focused_lines: int
	error_counts: dict[str, int]
	top_errors: list[dict[str, object]]
	merged_path: Path
	normalized_path: Path
	focused_path: Path
	focused_json_path: Path

	def to_dict(self) -> dict[str, object]:
		result = asdict(self)
		for key in (
			"merged_path",
			"normalized_path",
			"focused_path",
			"focused_json_path",
		):
			result[key] = str(result[key])
		return result


def process_trace_shards(
	raw_directory: Path,
	result_directory: Path,
	*,
	trigger_epoch: float,
	focus_before_seconds: float,
	focus_after_seconds: float,
	live_game: Path,
) -> TraceSummary:
	shards = sorted(
		(path for path in raw_directory.glob("trace.*") if path.is_file()),
		key=_shard_pid,
	)
	merged_path = result_directory / "merged.txt"
	normalized_path = result_directory / "filesystem.jsonl"
	focused_path = result_directory / "focused.txt"
	focused_json_path = result_directory / "focused.jsonl"
	error_counts: Counter[str] = Counter()
	error_paths: Counter[tuple[str, str, str, str | None]] = Counter()
	line_count = 0
	normalized_count = 0
	focused_count = 0
	focus_start = trigger_epoch - focus_before_seconds
	focus_end = trigger_epoch + focus_after_seconds

	with (
		merged_path.open("w", encoding="utf-8") as merged,
		normalized_path.open("w", encoding="utf-8") as normalized,
		focused_path.open("w", encoding="utf-8") as focused,
		focused_json_path.open("w", encoding="utf-8") as focused_json,
	):
		for timestamp, pid, raw in _merge_shards(shards):
			line_count += 1
			body = _body(raw)
			merged_line = f"{timestamp:.6f} [{pid}] {body}\n"
			merged.write(merged_line)
			record = parse_trace_line(pid, raw, live_game)
			if record is not None:
				normalized_count += 1
				normalized.write(json.dumps(record, sort_keys=True))
				normalized.write("\n")
				errno = record.get("errno")
				classification = record.get("error_classification")
				if isinstance(errno, str) and isinstance(classification, str):
					error_counts[classification] += 1
					paths = record.get("paths")
					path = paths[0] if isinstance(paths, list) and paths else ""
					game_relative = record.get("game_relative")
					error_paths[
						(
							classification,
							errno,
							str(path),
							str(game_relative) if game_relative is not None else None,
						)
					] += 1
			if focus_start <= timestamp <= focus_end:
				focused_count += 1
				focused.write(merged_line)
				if record is not None:
					focused_json.write(json.dumps(record, sort_keys=True))
					focused_json.write("\n")

	top_errors = [
		{
			"classification": key[0],
			"errno": key[1],
			"path": key[2],
			"game_relative": key[3],
			"count": count,
		}
		for key, count in error_paths.most_common(200)
	]
	return TraceSummary(
		shards=len(shards),
		lines=line_count,
		normalized_records=normalized_count,
		focused_lines=focused_count,
		error_counts=dict(sorted(error_counts.items())),
		top_errors=top_errors,
		merged_path=merged_path,
		normalized_path=normalized_path,
		focused_path=focused_path,
		focused_json_path=focused_json_path,
	)


def parse_trace_line(pid: int, line: str, live_game: Path) -> dict[str, object] | None:
	match = _TIMESTAMP_RE.match(line.rstrip("\n"))
	if match is None:
		return None
	timestamp = float(match.group("timestamp"))
	body = match.group("body")
	syscall_match = _SYSCALL_RE.match(body)
	if syscall_match is None:
		return None
	paths = [_decode_quoted(item) for item in _QUOTED_RE.findall(body)]
	paths = [item for item in paths if item is not None]
	errno_match = _ERRNO_RE.search(body)
	errno = errno_match.group("errno") if errno_match is not None else None
	classification = (
		_ERROR_CLASSIFICATIONS.get(errno, "other-error") if errno is not None else None
	)
	game_relative = next(
		(
			_relative_game_path(path, live_game)
			for path in paths
			if _relative_game_path(path, live_game) is not None
		),
		None,
	)
	return {
		"timestamp": timestamp,
		"pid": pid,
		"syscall": syscall_match.group("syscall"),
		"paths": paths,
		"errno": errno,
		"error_classification": classification,
		"game_relative": game_relative,
		"raw": body,
	}


def _merge_shards(paths: list[Path]) -> Iterator[tuple[float, int, str]]:
	with ExitStack() as stack:
		iterators: list[Iterator[tuple[float, int, str]]] = []
		for path in paths:
			handle = stack.enter_context(path.open("r", encoding="utf-8", errors="replace"))
			iterators.append(_iter_shard(handle, _shard_pid(path)))
		heap: list[tuple[float, int, int, str]] = []
		for index, iterator in enumerate(iterators):
			try:
				timestamp, pid, raw = next(iterator)
			except StopIteration:
				continue
			heapq.heappush(heap, (timestamp, pid, index, raw))
		while heap:
			timestamp, pid, index, raw = heapq.heappop(heap)
			yield timestamp, pid, raw
			try:
				next_timestamp, next_pid, next_raw = next(iterators[index])
			except StopIteration:
				continue
			heapq.heappush(heap, (next_timestamp, next_pid, index, next_raw))


def _iter_shard(handle: TextIO, pid: int) -> Iterator[tuple[float, int, str]]:
	last_timestamp = 0.0
	for raw in handle:
		match = _TIMESTAMP_RE.match(raw)
		if match is not None:
			last_timestamp = float(match.group("timestamp"))
		yield last_timestamp, pid, raw


def _shard_pid(path: Path) -> int:
	try:
		return int(path.name.rsplit(".", 1)[1])
	except (IndexError, ValueError):
		return 0


def _body(raw: str) -> str:
	match = _TIMESTAMP_RE.match(raw.rstrip("\n"))
	return match.group("body") if match is not None else raw.rstrip("\n")


def _decode_quoted(value: str) -> str | None:
	try:
		decoded = ast.literal_eval(value)
	except (SyntaxError, ValueError):
		return None
	return decoded if isinstance(decoded, str) else None


def _relative_game_path(value: str, live_game: Path) -> str | None:
	normalized = value.replace("\\", "/")
	root = str(live_game).replace("\\", "/").rstrip("/")
	index = normalized.casefold().find(root.casefold())
	if index < 0:
		return None
	relative = normalized[index + len(root) :].lstrip("/")
	return relative or "."