"""Token usage logger for MCP tool calls.

Appends a record to token_logs.txt in a temp directory per application:
  <tempdir>/<app_name>/token_logs.txt

Measures input argument size and output response size as a proxy for token
usage (~3.5 chars = 1 token, same estimate used by generate_tests internally).
"""

from __future__ import annotations

import json
import tempfile
import threading
from datetime import datetime
from pathlib import Path


_CHARS_PER_TOKEN = 3.5

# Resolved lazily on first tool call — set once per server process
_LOG_FILE: Path | None = None
_JSONL_FILE: Path | None = None
_SESSION_ID: str | None = None
_SESSION_LOCK = threading.Lock()
_WRITE_LOCK = threading.Lock()

# Running per-session totals {tool_name: {"calls", "input_chars", "output_chars", "total_duration"}}
_SESSION_STATS: dict[str, dict] = {}


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _est(chars: int) -> int:
    return max(0, round(chars / _CHARS_PER_TOKEN))


def _serialize_size(obj: object) -> int:
    try:
        return len(json.dumps(obj, default=str))
    except Exception:
        return len(str(obj))


def _resolve_app_name() -> str:
    """Read the active application name from the token store. Falls back to 'maeris'."""
    try:
        from maeris_mcp.auth.token_store import TokenStore
        name = TokenStore().get_app_name()
        if name:
            # Sanitise for use as a directory name
            safe = "".join(c if c.isalnum() or c in "-_ " else "_" for c in name).strip()
            return safe or "maeris"
    except Exception:
        pass
    return "maeris"


def _get_log_file() -> Path:
    """Return (and lazily resolve) the log file path for this server process."""
    global _LOG_FILE
    if _LOG_FILE is None:
        app_name = _resolve_app_name()
        _LOG_FILE = Path(tempfile.gettempdir()) / "maeris_logs" / app_name / "token_logs.txt"
    return _LOG_FILE


def _get_jsonl_file() -> Path:
    """Return (and lazily resolve) the structured JSONL log file path for this server process."""
    global _JSONL_FILE
    if _JSONL_FILE is None:
        app_name = _resolve_app_name()
        _JSONL_FILE = Path(tempfile.gettempdir()) / "maeris_logs" / app_name / "tool_calls.jsonl"
    return _JSONL_FILE


# ---------------------------------------------------------------------------
# Session management
# ---------------------------------------------------------------------------

def _get_or_init_session() -> str:
    global _SESSION_ID
    with _SESSION_LOCK:
        if _SESSION_ID is None:
            _SESSION_ID = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
            _write_session_header(_SESSION_ID)
        return _SESSION_ID


def _write_session_header(session_id: str) -> None:
    log_file = _get_log_file()
    line = (
        "\n"
        + "=" * 80 + "\n"
        + f"SESSION START : {session_id}\n"
        + f"LOG FILE      : {log_file}\n"
        + "=" * 80 + "\n"
        + f"{'TOOL':<35} {'CALLS':>6}  {'IN_CHARS':>10}  {'IN_TOK':>8}  "
          f"{'OUT_CHARS':>10}  {'OUT_TOK':>8}  {'TOTAL_TOK':>10}  {'AVG_DUR_S':>9}\n"
        + "-" * 80 + "\n"
    )
    with _WRITE_LOCK:
        log_file.parent.mkdir(parents=True, exist_ok=True)
        with open(log_file, "a", encoding="utf-8") as f:
            f.write(line)


# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------

def log_tool_call(
    tool_name: str,
    arguments: dict,
    result: object,
    duration_s: float,
) -> None:
    """Append one tool-call record.

    Called from server.py after every MCP tool dispatch.
    """
    session_id = _get_or_init_session()
    log_file = _get_log_file()
    jsonl_file = _get_jsonl_file()

    input_chars = _serialize_size(arguments)
    output_chars = _serialize_size(result)
    input_tokens = _est(input_chars)
    output_tokens = _est(output_chars)
    total_tokens = input_tokens + output_tokens

    ts = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    duration_ms = max(0, round(duration_s * 1000))

    with _SESSION_LOCK:
        s = _SESSION_STATS.setdefault(tool_name, {
            "calls": 0, "input_chars": 0, "output_chars": 0, "total_duration": 0.0,
        })
        s["calls"] += 1
        s["input_chars"] += input_chars
        s["output_chars"] += output_chars
        s["total_duration"] += duration_s

    line = (
        f"[{ts}] {tool_name}\n"
        f"  input : {input_chars:>10,} chars  (~{input_tokens:>7,} tokens)\n"
        f"  output: {output_chars:>10,} chars  (~{output_tokens:>7,} tokens)\n"
        f"  total : {input_chars + output_chars:>10,} chars  (~{total_tokens:>7,} tokens)"
        f"  |  duration: {duration_s:.2f}s\n\n"
    )

    with _WRITE_LOCK:
        log_file.parent.mkdir(parents=True, exist_ok=True)
        with open(log_file, "a", encoding="utf-8") as f:
            f.write(line)
        # Structured record (one JSON object per line).
        try:
            record = {
                "timestamp": ts,
                "sessionId": session_id,
                "tool": tool_name,
                "durationMs": duration_ms,
                "inputChars": input_chars,
                "outputChars": output_chars,
                "inputEstimatedTokens": input_tokens,
                "outputEstimatedTokens": output_tokens,
                "totalEstimatedTokens": total_tokens,
            }
            with open(jsonl_file, "a", encoding="utf-8") as jf:
                jf.write(json.dumps(record, ensure_ascii=False) + "\n")
        except Exception:
            # Never fail tool execution because of logging.
            pass


def log_files_read(
    session_id: str,
    call_number: int,
    phase: str,
    files_read: list[str],
    token_stats: dict,
    duration_ms: int,
) -> None:
    """Append a per-batch file-read record to token_logs.txt.

    Called from generate_tests for each scan batch so the user can see
    exactly which source files were sent to the LLM in each call.
    Never raises.
    """
    try:
        _get_or_init_session()
        log_file = _get_log_file()
        this_call = token_stats.get("thisCall", {})
        in_tok  = this_call.get("inputEstimatedTokens", 0)
        out_tok = this_call.get("outputEstimatedTokens", 0)
        tot_tok = this_call.get("totalEstimatedTokens", 0)
        cumul   = token_stats.get("sessionCumulativeTokens", 0)
        ts = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        files_section = (
            "\n".join(f"    {f}" for f in files_read)
            if files_read else "    (none)"
        )
        line = (
            f"[{ts}] generate_tests | call={call_number} | {phase} | session={session_id[:8]}\n"
            f"  files read ({len(files_read)}):\n"
            f"{files_section}\n"
            f"  tokens this call : ~{tot_tok:,}  (in: ~{in_tok:,} / out: ~{out_tok:,})\n"
            f"  session cumulative: ~{cumul:,}\n"
            f"  duration         : {duration_ms:,}ms\n\n"
        )
        with _WRITE_LOCK:
            log_file.parent.mkdir(parents=True, exist_ok=True)
            with open(log_file, "a", encoding="utf-8") as f:
                f.write(line)
    except Exception:
        pass  # Never fatal


def write_session_summary() -> None:
    """Write a per-tool summary table sorted by total tokens (highest first).

    Called on server shutdown from main.py.
    """
    with _SESSION_LOCK:
        stats = dict(_SESSION_STATS)
        session_id = _SESSION_ID or "unknown"

    if not stats:
        return

    log_file = _get_log_file()
    rows = sorted(stats.items(), key=lambda kv: -(kv[1]["input_chars"] + kv[1]["output_chars"]))

    lines = [
        "\n" + "=" * 100 + "\n",
        f"SESSION SUMMARY: {session_id}\n",
        "=" * 100 + "\n",
        f"{'TOOL':<35} {'CALLS':>6}  {'IN_CHARS':>10}  {'IN_TOK':>8}  "
        f"{'OUT_CHARS':>10}  {'OUT_TOK':>8}  {'TOTAL_TOK':>10}  {'AVG_DUR_S':>9}\n",
        "-" * 100 + "\n",
    ]
    for tool_name, s in rows:
        calls = s["calls"]
        in_c = s["input_chars"]
        out_c = s["output_chars"]
        in_t = _est(in_c)
        out_t = _est(out_c)
        total_t = in_t + out_t
        avg_dur = s["total_duration"] / calls if calls else 0.0
        lines.append(
            f"{tool_name:<35} {calls:>6}  {in_c:>10,}  {in_t:>8,}  "
            f"{out_c:>10,}  {out_t:>8,}  {total_t:>10,}  {avg_dur:>9.2f}\n"
        )
    lines.append("-" * 100 + "\n\n")

    with _WRITE_LOCK:
        log_file.parent.mkdir(parents=True, exist_ok=True)
        with open(log_file, "a", encoding="utf-8") as f:
            f.writelines(lines)
