from __future__ import annotations

import hashlib
import json
from collections import defaultdict
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Literal

from agentic_pd_hybrid.trace import TraceRequest, load_trace


SampleProfile = Literal["default", "small-append"]


@dataclass(frozen=True)
class SessionSampleConfig:
    trace_path: Path
    output_path: Path
    target_duration_s: float = 600.0
    start_time_s: float = 0.0
    session_sample_rate: float = 1.0
    min_turns: int = 1
    max_requests: int | None = None
    profile: SampleProfile = "default"
    min_initial_input_tokens: int | None = None
    max_initial_input_tokens: int | None = None
    max_append_input_tokens: int | None = None
    max_output_tokens: int | None = None
    min_overlap_ratio: float | None = None


@dataclass(frozen=True)
class SessionSampleSummary:
    input_trace_path: str
    output_trace_path: str
    request_count: int
    session_count: int
    multi_turn_session_count: int
    start_time_s: float
    end_time_s: float
    sampled_duration_s: float
    session_sample_rate: float
    min_turns: int
    profile: str
    min_initial_input_tokens: int | None
    max_initial_input_tokens: int | None
    max_append_input_tokens: int | None
    max_output_tokens: int | None
    min_overlap_ratio: float | None
    mean_append_input_tokens: float | None
    mean_turn_overlap_ratio: float | None


def sample_trace_sessions(config: SessionSampleConfig) -> SessionSampleSummary:
    requests = load_trace(config.trace_path)
    sessions: dict[str, list[TraceRequest]] = defaultdict(list)
    for request in requests:
        sessions[request.session_id].append(request)

    filters = _resolve_filters(config)
    eligible_sessions = {
        session_id: session_requests
        for session_id, session_requests in sessions.items()
        if len(session_requests) >= filters.min_turns
        and _session_matches_filters(session_requests, filters)
        and _keep_session(session_id, config.session_sample_rate)
    }
    ordered_sessions = sorted(
        eligible_sessions.values(),
        key=lambda session_requests: session_requests[0].timestamp_s,
    )

    selected_requests: list[TraceRequest] = []
    sampled_start: float | None = None
    sampled_end: float | None = None
    for session_requests in ordered_sessions:
        session_first = session_requests[0].timestamp_s
        if session_first < config.start_time_s:
            continue

        if sampled_start is None:
            sampled_start = session_first

        selected_requests.extend(session_requests)
        sampled_end = max(request.timestamp_s for request in session_requests)

        if config.max_requests is not None and len(selected_requests) >= config.max_requests:
            break
        if sampled_end - sampled_start >= config.target_duration_s:
            break

    selected_requests.sort(key=lambda request: request.timestamp_s)
    if config.max_requests is not None:
        selected_requests = selected_requests[: config.max_requests]

    if not selected_requests:
        raise ValueError("Sampling produced no requests; adjust the sampling arguments")

    config.output_path.parent.mkdir(parents=True, exist_ok=True)
    with config.output_path.open("w", encoding="utf-8") as handle:
        for request in selected_requests:
            payload = {
                "request_id": request.request_id,
                "session_id": request.session_id,
                "chat_id": request.chat_id,
                "parent_chat_id": request.parent_chat_id,
                "timestamp": request.timestamp_s,
                "input_length": request.input_length,
                "output_length": request.output_length,
                "type": request.request_type,
                "turn": request.turn_id,
                "hash_ids": list(request.hash_ids),
            }
            handle.write(json.dumps(payload, sort_keys=True) + "\n")

    selected_session_ids = {request.session_id for request in selected_requests}
    selected_session_requests = [
        eligible_sessions[session_id] for session_id in selected_session_ids
    ]
    append_lengths = [
        length
        for session_requests in selected_session_requests
        for length in _turn_append_lengths(session_requests)
    ]
    overlap_ratios = [
        ratio
        for session_requests in selected_session_requests
        for ratio in _turn_overlap_ratios(session_requests)
    ]
    summary = SessionSampleSummary(
        input_trace_path=str(config.trace_path),
        output_trace_path=str(config.output_path),
        request_count=len(selected_requests),
        session_count=len(selected_session_ids),
        multi_turn_session_count=sum(
            1
            for session_id in selected_session_ids
            if len(eligible_sessions[session_id]) > 1
        ),
        start_time_s=selected_requests[0].timestamp_s,
        end_time_s=selected_requests[-1].timestamp_s,
        sampled_duration_s=selected_requests[-1].timestamp_s
        - selected_requests[0].timestamp_s,
        session_sample_rate=config.session_sample_rate,
        min_turns=filters.min_turns,
        profile=config.profile,
        min_initial_input_tokens=filters.min_initial_input_tokens,
        max_initial_input_tokens=filters.max_initial_input_tokens,
        max_append_input_tokens=filters.max_append_input_tokens,
        max_output_tokens=filters.max_output_tokens,
        min_overlap_ratio=filters.min_overlap_ratio,
        mean_append_input_tokens=_mean(append_lengths),
        mean_turn_overlap_ratio=_mean(overlap_ratios),
    )
    summary_path = config.output_path.with_suffix(config.output_path.suffix + ".summary.json")
    with summary_path.open("w", encoding="utf-8") as handle:
        json.dump(asdict(summary), handle, indent=2, sort_keys=True)
    return summary


@dataclass(frozen=True)
class _ResolvedFilters:
    min_turns: int
    min_initial_input_tokens: int | None
    max_initial_input_tokens: int | None
    max_append_input_tokens: int | None
    max_output_tokens: int | None
    min_overlap_ratio: float | None


def _resolve_filters(config: SessionSampleConfig) -> _ResolvedFilters:
    if config.profile == "default":
        return _ResolvedFilters(
            min_turns=config.min_turns,
            min_initial_input_tokens=config.min_initial_input_tokens,
            max_initial_input_tokens=config.max_initial_input_tokens,
            max_append_input_tokens=config.max_append_input_tokens,
            max_output_tokens=config.max_output_tokens,
            min_overlap_ratio=config.min_overlap_ratio,
        )

    if config.profile != "small-append":
        raise ValueError(f"Unsupported sample profile: {config.profile}")

    return _ResolvedFilters(
        min_turns=max(config.min_turns, 2),
        min_initial_input_tokens=(
            2048
            if config.min_initial_input_tokens is None
            else config.min_initial_input_tokens
        ),
        max_initial_input_tokens=(
            16000
            if config.max_initial_input_tokens is None
            else config.max_initial_input_tokens
        ),
        max_append_input_tokens=(
            2048
            if config.max_append_input_tokens is None
            else config.max_append_input_tokens
        ),
        max_output_tokens=(
            2048 if config.max_output_tokens is None else config.max_output_tokens
        ),
        min_overlap_ratio=(
            0.75 if config.min_overlap_ratio is None else config.min_overlap_ratio
        ),
    )


def _session_matches_filters(
    session_requests: list[TraceRequest],
    filters: _ResolvedFilters,
) -> bool:
    ordered = sorted(
        session_requests,
        key=lambda request: (request.timestamp_s, request.turn_id, request.chat_id),
    )
    if not ordered:
        return False

    initial = ordered[0]
    if (
        filters.min_initial_input_tokens is not None
        and initial.input_length < filters.min_initial_input_tokens
    ):
        return False
    if (
        filters.max_initial_input_tokens is not None
        and initial.input_length > filters.max_initial_input_tokens
    ):
        return False
    if filters.max_output_tokens is not None and any(
        request.output_length > filters.max_output_tokens for request in ordered
    ):
        return False

    append_lengths = _turn_append_lengths(ordered)
    if filters.max_append_input_tokens is not None and any(
        append_length <= 0 or append_length > filters.max_append_input_tokens
        for append_length in append_lengths
    ):
        return False

    overlap_ratios = _turn_overlap_ratios(ordered)
    if filters.min_overlap_ratio is not None and any(
        overlap_ratio < filters.min_overlap_ratio for overlap_ratio in overlap_ratios
    ):
        return False

    return True


def _turn_append_lengths(session_requests: list[TraceRequest]) -> list[int]:
    ordered = sorted(
        session_requests,
        key=lambda request: (request.timestamp_s, request.turn_id, request.chat_id),
    )
    return [
        current.input_length - (previous.input_length + previous.output_length)
        for previous, current in zip(ordered, ordered[1:], strict=False)
    ]


def _turn_overlap_ratios(session_requests: list[TraceRequest]) -> list[float]:
    ordered = sorted(
        session_requests,
        key=lambda request: (request.timestamp_s, request.turn_id, request.chat_id),
    )
    ratios: list[float] = []
    for previous, current in zip(ordered, ordered[1:], strict=False):
        if not current.hash_ids:
            ratios.append(0.0)
            continue
        previous_blocks = set(previous.hash_ids)
        overlap = sum(1 for block in current.hash_ids if block in previous_blocks)
        ratios.append(overlap / len(current.hash_ids))
    return ratios


def _mean(values: list[int] | list[float]) -> float | None:
    if not values:
        return None
    return sum(values) / len(values)


def _keep_session(session_id: str, sample_rate: float) -> bool:
    if sample_rate >= 1.0:
        return True
    if sample_rate <= 0.0:
        return False
    digest = hashlib.blake2b(session_id.encode("utf-8"), digest_size=8).digest()
    bucket = int.from_bytes(digest, byteorder="big", signed=False) / 2**64
    return bucket < sample_rate