"""Filter command."""

import asyncio
import json
import math
from collections.abc import MutableMapping
from dataclasses import dataclass
from datetime import UTC, datetime
from pathlib import Path
from typing import Any, Optional, Sequence

import click

try:  # Python 3.11+
    import tomllib as _toml  # type: ignore[attr-defined]
except Exception:  # pragma: no cover
    _toml = None  # type: ignore[assignment]

from synth_ai.core.tracing_v3 import SessionTracer  # type: ignore[import-untyped]
from synth_ai.sdk.task.config import FilterConfig


class FilterCliError(RuntimeError):
    """Base exception for filter CLI failures."""


@dataclass(slots=True)
class TomlUnavailableError(FilterCliError):
    hint: str | None = None


@dataclass(slots=True)
class FilterConfigNotFoundError(FilterCliError):
    path: str


@dataclass(slots=True)
class FilterConfigParseError(FilterCliError):
    path: str
    detail: str


@dataclass(slots=True)
class MissingFilterTableError(FilterCliError):
    """Raised when the filter config lacks a [filter] table."""


@dataclass(slots=True)
class InvalidFilterConfigError(FilterCliError):
    detail: str


@dataclass(slots=True)
class NoTracesFoundError(FilterCliError):
    db_url: str


@dataclass(slots=True)
class NoSessionsMatchedError(FilterCliError):
    hint: str | None = None


def validate_filter_options(options: MutableMapping[str, Any]) -> MutableMapping[str, Any]:
    """Validate parameters passed to the filter CLI command."""
    result: dict[str, Any] = dict(options)

    def _coerce_list(key: str) -> None:
        value = result.get(key)
        if value is None:
            result[key] = []
        elif isinstance(value, list | tuple | set):
            result[key] = [str(item).strip() for item in value if str(item).strip()]
        else:
            result[key] = [str(value).strip()] if str(value).strip() else []

    def _coerce_dict(key: str) -> None:
        value = result.get(key)
        if value is None:
            result[key] = {}
        elif isinstance(value, MutableMapping):
            normalized: dict[str, float] = {}
            for k, v in value.items():
                if k is None:
                    continue
                try:
                    number = float(v)
                    if math.isnan(number) or math.isinf(number):
                        continue
                    normalized[str(k).strip()] = number
                except Exception:
                    continue
            result[key] = normalized
        else:
            result[key] = {}

    _coerce_list("splits")
    _coerce_list("task_ids")
    _coerce_list("models")
    _coerce_dict("min_verifier_scores")
    _coerce_dict("max_verifier_scores")

    objective_value = result.get("objective")
    if objective_value is None or str(objective_value).strip() == "":
        result["objective"] = None
    else:
        result["objective"] = str(objective_value).strip()

    for duration_key in ("min_official_score", "max_official_score"):
        value = result.get(duration_key)
        if value is None or value == "":
            result[duration_key] = None
        else:
            try:
                result[duration_key] = float(value)
            except Exception:
                result[duration_key] = None

    for int_key in ("limit", "offset", "shuffle_seed"):
        value = result.get(int_key)
        if value is None or value == "":
            result[int_key] = None
        else:
            try:
                result[int_key] = int(value)
            except Exception:
                result[int_key] = None

    shuffle_value = result.get("shuffle")
    if isinstance(shuffle_value, str):
        result["shuffle"] = shuffle_value.strip().lower() in {"1", "true", "yes"}
    else:
        result["shuffle"] = bool(shuffle_value)

    return result


def _parse_datetime_for_trace(value: Any) -> datetime | None:
    if isinstance(value, datetime):
        return value if value.tzinfo else value.replace(tzinfo=UTC)
    if isinstance(value, str):
        value = value.replace("Z", "+00:00")
        try:
            dt = datetime.fromisoformat(value)
        except ValueError:
            try:
                dt = datetime.fromtimestamp(float(value), tz=UTC)
            except Exception:
                return None
        return dt if dt.tzinfo else dt.replace(tzinfo=UTC)
    if isinstance(value, int | float):
        try:
            return datetime.fromtimestamp(float(value), tz=UTC)
        except Exception:
            return None
    return None


def _score_ok(value: Any, min_val: Any, max_val: Any) -> bool:
    try:
        if value is None:
            return min_val is None
        value = float(value)
    except Exception:
        return False
    if min_val is not None and value < float(min_val):
        return False
    return not (max_val is not None and value > float(max_val))


def _load_filter_config(config_path: Path) -> tuple[FilterConfig, dict[str, Any]]:
    if _toml is None:
        raise TomlUnavailableError(hint="Install tomli or use Python 3.11+")

    if not config_path.exists():
        raise FilterConfigNotFoundError(path=str(config_path))

    try:
        config_data = _toml.loads(config_path.read_text(encoding="utf-8"))
    except Exception as exc:  # pragma: no cover - validation tests cover common cases
        raise FilterConfigParseError(path=str(config_path), detail=str(exc)) from exc

    filter_cfg_dict = config_data.get("filter") if isinstance(config_data, dict) else None
    if not isinstance(filter_cfg_dict, dict):
        raise MissingFilterTableError()

    try:
        normalized = validate_filter_options(filter_cfg_dict)
        normalized_dict = dict(normalized)
        filter_cfg = FilterConfig.from_dict(normalized_dict)
    except (ValueError, TypeError) as validation_error:
        raise InvalidFilterConfigError(detail=str(validation_error)) from validation_error

    click.echo(f"✓ Config validated: db={filter_cfg.db}, output={filter_cfg.output}")
    if filter_cfg.min_official_score is not None:
        click.echo(f"  → Filtering for official score >= {filter_cfg.min_official_score}")
    if filter_cfg.objective:
        click.echo(f"  → Filtering objective: {filter_cfg.objective}")
    if filter_cfg.limit:
        click.echo(f"  → Limiting to {filter_cfg.limit} examples")

    return filter_cfg, normalized_dict


def _extract_content(content: Any) -> Any:
    if isinstance(content, dict) and "content" in content:
        return content["content"]
    return content


def _extract_text(content: Any) -> str:
    if isinstance(content, str):
        return content
    if isinstance(content, dict):
        payload = content.get("payload") if isinstance(content.get("payload"), dict) else None
        if payload and "content" in payload:
            return _extract_text(payload["content"])
        for key in ("text", "content", "content_text"):
            if key in content:
                value = content[key]
                if isinstance(value, str):
                    return value
        try:
            return json.dumps(content)
        except Exception:  # pragma: no cover - defensive
            return str(content)
    if isinstance(content, list):
        parts = []
        for item in content:
            if isinstance(item, dict) and item.get("type") == "text":
                parts.append(item.get("text", ""))
        return " ".join(parts) if parts else str(content)
    return str(content)


def _select_messages(message_rows: Sequence[dict[str, Any]]) -> list[dict[str, Any]]:
    records: list[dict[str, Any]] = []
    for index, msg_row in enumerate(message_rows):
        msg_type = msg_row.get("message_type")
        content_raw = msg_row.get("content")
        if msg_type not in {"user", "policy_user_prompt"}:
            continue

        # Look backwards for system prompt
        system_msg = None
        for prev in range(index - 1, -1, -1):
            prev_type = message_rows[prev].get("message_type")
            if prev_type == "policy_system_prompt":
                system_msg = message_rows[prev]
                break

        assistant_msg = None
        tool_call_msg = None
        for follow in range(index + 1, len(message_rows)):
            next_type = message_rows[follow].get("message_type")
            if next_type == "assistant":
                assistant_msg = message_rows[follow]
                break
            elif next_type == "policy_tool_call":
                tool_call_msg = message_rows[follow]
                break

        try:
            user_content = json.loads(content_raw) if isinstance(content_raw, str) else content_raw
        except Exception:
            user_content = content_raw

        user_content = _extract_content(user_content)
        user_text = _extract_text(user_content)
        if not user_text:
            continue

        messages = []

        # Add system prompt if found
        if system_msg is not None:
            try:
                system_content_raw = system_msg.get("content")
                system_content = (
                    json.loads(system_content_raw)
                    if isinstance(system_content_raw, str)
                    else system_content_raw
                )
                system_content = _extract_content(system_content)
                system_text = _extract_text(system_content)
                if system_text:
                    messages.append({"role": "system", "content": system_text})
            except Exception:
                pass

        # Add user message
        user_payload = user_content if isinstance(user_content, list) else user_text
        messages.append({"role": "user", "content": user_payload})

        # Add assistant/tool call response
        assistant_content = None
        if tool_call_msg is not None:
            raw = tool_call_msg.get("content")
            try:
                assistant_content = json.loads(raw) if isinstance(raw, str) else raw
            except Exception:
                assistant_content = raw
            assistant_content = _extract_content(assistant_content)
        elif assistant_msg is not None:
            raw = assistant_msg.get("content")
            try:
                assistant_content = json.loads(raw) if isinstance(raw, str) else raw
            except Exception:
                assistant_content = raw
            assistant_content = _extract_content(assistant_content)

        assistant_payload = (
            assistant_content
            if isinstance(assistant_content, list)
            else (
                _extract_text(assistant_content)
                if assistant_content is not None
                else "[no response recorded]"
            )
        )
        messages.append({"role": "assistant", "content": assistant_payload})

        records.append({"messages": messages})
    return records


@click.command(
    help="Export filtered tracing sessions to SFT-ready JSONL based on a TOML config.",
)
@click.option(
    "--config",
    "config_path",
    type=click.Path(),
    required=True,
    help="Path to TOML config describing the input trace DB, score thresholds, and output JSONL.",
)
@click.option(
    "--objective",
    "objective_key",
    required=False,
    help="Objective key to filter in outcome_rewards (default: reward).",
)
def filter(config_path: str, objective_key: Optional[str]) -> None:
    try:
        filter_cfg, raw_cfg = _load_filter_config(Path(config_path))
    except FilterCliError as exc:
        raise click.ClickException(_format_filter_error(exc)) from exc

    db_url = filter_cfg.get_db_url()
    output_path = filter_cfg.get_output_path()

    splits = set(filter_cfg.splits)
    task_ids = set(filter_cfg.task_ids)
    models = set(filter_cfg.models)
    min_official = filter_cfg.min_official_score
    max_official = filter_cfg.max_official_score
    objective = objective_key or filter_cfg.objective or "reward"
    min_verifier_scores = filter_cfg.min_verifier_scores
    max_verifier_scores = filter_cfg.max_verifier_scores
    min_created = _parse_datetime_for_trace(raw_cfg.get("min_created_at"))
    max_created = _parse_datetime_for_trace(raw_cfg.get("max_created_at"))
    limit = filter_cfg.limit

    async def _run() -> None:
        tracer = SessionTracer(db_url=db_url, auto_save=False)
        await tracer.initialize()

        if tracer.db is None:
            raise FilterCliError("Database not initialized")

        df = await tracer.db.query_traces(
            "SELECT session_id, created_at, metadata FROM session_traces ORDER BY created_at"
        )
        if getattr(df, "empty", True):
            raise NoTracesFoundError(db_url=db_url)

        sessions = df.to_dict("records")
        accepted: list[dict[str, Any]] = []

        for row in sessions:
            metadata_raw = row.get("metadata")
            if isinstance(metadata_raw, str):
                try:
                    metadata = json.loads(metadata_raw)
                except Exception:
                    metadata = {}
            elif isinstance(metadata_raw, dict):
                metadata = dict(metadata_raw)
            else:
                metadata = {}

            created_at_raw = row.get("created_at")
            created_at_dt = _parse_datetime_for_trace(created_at_raw)
            session_id = row.get("session_id")

            if splits and metadata.get("task_split") not in splits:
                continue
            if task_ids and metadata.get("task_id") not in task_ids:
                continue
            if models and metadata.get("model") not in models:
                continue

            if min_created and (created_at_dt is None or created_at_dt < min_created):
                continue
            if max_created and (created_at_dt is None or created_at_dt > max_created):
                continue

            total_reward = None
            achievements_count = None
            if min_official is not None or max_official is not None:
                if tracer.db is None:
                    raise FilterCliError("Database not initialized")
                reward_query = (
                    "SELECT total_reward, achievements_count FROM outcome_rewards "
                    "WHERE session_id = :session_id AND objective_key = :objective_key"
                )
                reward_params = {"session_id": session_id, "objective_key": objective}
                try:
                    reward_rows = await tracer.db.query_traces(reward_query, reward_params)
                except Exception:
                    reward_rows = await tracer.db.query_traces(
                        "SELECT total_reward, achievements_count FROM outcome_rewards WHERE session_id = :session_id",
                        {"session_id": session_id},
                    )
                reward_records = (
                    reward_rows.to_dict("records") if hasattr(reward_rows, "to_dict") else []
                )
                if reward_records:
                    total_reward = reward_records[0].get("total_reward")
                    achievements_count = reward_records[0].get("achievements_count")
                    if not _score_ok(total_reward, min_official, max_official):
                        continue
                elif min_official is not None:
                    continue

            verifier_scores = metadata.get("verifier_scores") or {}
            include = True
            for verifier_name, threshold in (min_verifier_scores or {}).items():
                if not _score_ok(verifier_scores.get(verifier_name), threshold, None):
                    include = False
                    break
            if not include:
                continue
            for verifier_name, threshold in (max_verifier_scores or {}).items():
                if not _score_ok(verifier_scores.get(verifier_name), None, threshold):
                    include = False
                    break
            if not include:
                continue

            messages_query = "\n            SELECT message_type, content, timestamp \n            FROM messages \n            WHERE session_id = :session_id\n            ORDER BY timestamp ASC, id ASC\n        "
            if tracer.db is None:
                raise FilterCliError("Database not initialized")
            msg_df = await tracer.db.query_traces(messages_query, {"session_id": session_id})
            message_rows = msg_df.to_dict("records") if hasattr(msg_df, "to_dict") else []

            if not message_rows:
                prompt = metadata.get("prompt") or ""
                completion = metadata.get("completion") or ""
                if prompt and completion:
                    accepted.append(
                        {
                            "messages": [
                                {"role": "user", "content": str(prompt)},
                                {"role": "assistant", "content": str(completion)},
                            ],
                            "metadata": {
                                "session_id": session_id,
                                "env_name": metadata.get("env_name"),
                                "policy_name": metadata.get("policy_name"),
                                "seed": metadata.get("seed"),
                                "total_reward": total_reward,
                                "achievements_count": achievements_count,
                                "model": metadata.get("model"),
                                "created_at": created_at_dt.isoformat()
                                if created_at_dt
                                else created_at_raw,
                            },
                        }
                    )
                continue

            for record in _select_messages(message_rows):
                record["metadata"] = {
                    "session_id": session_id,
                    "env_name": metadata.get("env_name"),
                    "policy_name": metadata.get("policy_name"),
                    "seed": metadata.get("seed"),
                    "total_reward": total_reward,
                    "achievements_count": achievements_count,
                    "model": metadata.get("model"),
                    "created_at": created_at_dt.isoformat() if created_at_dt else created_at_raw,
                }
                accepted.append(record)

        if not accepted:
            raise NoSessionsMatchedError()

        if limit is not None and limit > 0:
            accepted[:] = accepted[:limit]

        output_path.parent.mkdir(parents=True, exist_ok=True)
        with output_path.open("w", encoding="utf-8") as handle:
            for item in accepted:
                handle.write(json.dumps(item, ensure_ascii=False))
                handle.write("\n")

        click.echo(f"Wrote {len(accepted)} examples -> {output_path}")
        if tracer.db is not None:
            await tracer.db.close()

    try:
        asyncio.run(_run())
    except FilterCliError as exc:
        raise click.ClickException(_format_filter_error(exc)) from exc


def _format_filter_error(err: FilterCliError) -> str:
    if isinstance(err, TomlUnavailableError):
        hint = err.hint or "Install tomli or use Python 3.11+."
        return f"TOML parser not available. {hint}"
    if isinstance(err, FilterConfigNotFoundError):
        return f"Filter config not found: {err.path}"
    if isinstance(err, FilterConfigParseError):
        return f"Failed to parse TOML '{err.path}': {err.detail}"
    if isinstance(err, MissingFilterTableError):
        return "Config must contain a [filter] table."
    if isinstance(err, InvalidFilterConfigError):
        return f"Invalid filter config: {err.detail}"
    if isinstance(err, NoTracesFoundError):
        return f"No traces found in database ({err.db_url})."
    if isinstance(err, NoSessionsMatchedError):
        hint = err.hint or "Adjust the filter thresholds or choose a different dataset."
        return f"No sessions matched the provided filters. {hint}"
    return str(err)
