from __future__ import annotations

import json
from datetime import datetime, timedelta
from pathlib import Path
from typing import Iterable

from .config import HISTORY_DAYS, HISTORY_HOURS, SESSIONS_DIR, SESSIONS_INDEX_PATH, TIMEZONE, USAGE_COST_CACHE_PATH
from .models import SessionIndexReport, UsageCostCacheReport, UsageReport
from .utils import as_int, open_session_file, sanitize_label, now_in_tz, session_id_from_path


def iter_session_paths() -> Iterable[Path]:
    for path in sorted(SESSIONS_DIR.glob("*.jsonl*")):
        if ".trajectory." in path.name:
            continue
        yield path


def collect_usage() -> UsageReport:
    report = UsageReport()
    now = now_in_tz()
    today = now.date()
    day_cutoff = today - timedelta(days=HISTORY_DAYS - 1)
    hour_cutoff = now.replace(minute=0, second=0, microsecond=0) - timedelta(hours=HISTORY_HOURS - 1)

    for path in iter_session_paths():
        report.files_scanned += 1
        session_id = session_id_from_path(path)

        try:
            with open_session_file(path) as handle:
                for line in handle:
                    report.lines_scanned += 1
                    if not line.strip():
                        continue

                    try:
                        record = json.loads(line)
                    except json.JSONDecodeError:
                        report.parse_errors += 1
                        continue

                    message = record.get("message") or {}
                    timestamp_ms = message.get("timestamp", record.get("timestamp"))
                    if not isinstance(timestamp_ms, (int, float)):
                        continue

                    dt = datetime.fromtimestamp(timestamp_ms / 1000, TIMEZONE)
                    usage_date = dt.date()
                    is_today = usage_date == today

                    if record.get("type") != "message":
                        report.non_message_records_all_time += 1
                        if is_today:
                            report.non_message_records_today += 1
                        continue

                    report.message_records_all_time += 1
                    if is_today:
                        report.message_records_today += 1

                    role = sanitize_label(message.get("role"))
                    model = sanitize_label(message.get("model"))
                    provider = sanitize_label(message.get("provider"))
                    stop_reason = sanitize_label(message.get("stopReason"))
                    tool_name = sanitize_label(message.get("toolName"))
                    error_type = sanitize_label(message.get("errorType"))
                    is_error = bool(message.get("isError"))
                    usage = message.get("usage")

                    if is_today:
                        report.role_counts_today[role] += 1
                        if model != "unknown":
                            report.unique_models_today.add(model)
                        if provider != "unknown":
                            report.unique_providers_today.add(provider)
                        if session_id:
                            report.unique_session_ids_today.add(session_id)
                            report.session_message_counts_today[(session_id, role)] += 1
                        if stop_reason != "unknown":
                            report.stop_reason_counts_today[(stop_reason, model, provider)] += 1
                        if tool_name != "unknown":
                            report.tool_call_counts_today[tool_name] += 1
                            if session_id:
                                report.session_tool_call_counts_today[(session_id, tool_name)] += 1
                        if is_error:
                            report.error_type_counts_today[error_type] += 1
                            if session_id:
                                report.session_error_counts_today[session_id] += 1

                    if not isinstance(usage, dict):
                        continue

                    total_tokens, input_tokens, output_tokens, _, _, reasoning_tokens, skipped_negative_costs = report.all_time.add_usage(usage)
                    if skipped_negative_costs:
                        report.negative_cost_entries_skipped += 1
                        for cost_field, raw_value in skipped_negative_costs.items():
                            report.negative_cost_fields_skipped[cost_field] += 1
                            report.negative_cost_amounts_skipped[cost_field] += abs(raw_value)
                    report.tokens_by_model_all_time[(model, provider)] += total_tokens
                    report.input_tokens_by_model_all_time[(model, provider)] += input_tokens
                    report.output_tokens_by_model_all_time[(model, provider)] += output_tokens
                    if session_id:
                        report.session_tokens_all_time[session_id] += total_tokens

                    if usage_date >= day_cutoff:
                        report.day_buckets[usage_date.isoformat()].add_usage(usage)

                    hour_bucket = dt.replace(minute=0, second=0, microsecond=0)
                    if hour_bucket >= hour_cutoff:
                        report.hour_buckets[hour_bucket.isoformat()].add_usage(usage)

                    if is_today:
                        report.today.add_usage(usage)
                        report.tokens_by_model_today[(model, provider)] += total_tokens
                        report.input_tokens_by_model_today[(model, provider)] += input_tokens
                        report.output_tokens_by_model_today[(model, provider)] += output_tokens
                        report.reasoning_tokens_by_model_today[(model, provider)] += reasoning_tokens
                        if session_id:
                            report.session_tokens_today[session_id] += total_tokens
                            report.session_input_tokens_today[session_id] += input_tokens
                            report.session_output_tokens_today[session_id] += output_tokens
        except OSError:
            report.parse_errors += 1
            continue

    return report


def collect_session_index() -> SessionIndexReport:
    report = SessionIndexReport()
    try:
        data = json.loads(SESSIONS_INDEX_PATH.read_text(encoding="utf-8"))
    except (OSError, json.JSONDecodeError):
        return report

    if not isinstance(data, dict):
        return report

    if isinstance(data.get("sessions"), list):
        sessions = [session for session in data.get("sessions", []) if isinstance(session, dict)]
        report.visible_count = as_int(data.get("count")) or len(sessions)
        report.total_count = as_int(data.get("totalCount")) or len(sessions)
        report.has_more = bool(data.get("hasMore"))
        report.sessions = sessions
        return report

    sessions = [session for session in data.values() if isinstance(session, dict)]
    report.visible_count = len(sessions)
    report.total_count = len(sessions)
    report.has_more = False
    report.sessions = sessions
    return report


def collect_usage_cost_cache() -> UsageCostCacheReport:
    report = UsageCostCacheReport()
    try:
        data = json.loads(USAGE_COST_CACHE_PATH.read_text(encoding="utf-8"))
    except (OSError, json.JSONDecodeError):
        return report

    if not isinstance(data, dict):
        return report

    files = data.get("files")
    report.version = as_int(data.get("version")) if data.get("version") is not None else None
    report.updated_at_ms = as_int(data.get("updatedAt")) if data.get("updatedAt") is not None else None
    report.files_tracked = len(files) if isinstance(files, dict) else 0
    return report