"""
Shared utilities for Honcho benchmark test runners and evaluators.

Contains the BaseRunner abstract class and RunnerConfig dataclass that provide
a common framework for all benchmark runners (longmem, beam, locomo), as well
as shared utilities for trace-based evaluators (molecular, coverage).
"""

import argparse
import asyncio
import json
import logging
import os
import time
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from datetime import datetime
from logging import Logger
from pathlib import Path
from typing import Any, Generic, TypeVar, cast

from anthropic import AsyncAnthropic
from honcho import Honcho
from honcho.api_types import SessionConfiguration, SummaryConfiguration
from openai import AsyncOpenAI

from src.config import REASONING_LEVELS, ReasoningLevel
from src.telemetry.metrics_collector import MetricsCollector

_logger = logging.getLogger(__name__)

# Type variable for result types
ResultT = TypeVar("ResultT")


@dataclass
class RunnerConfig:
    """Configuration shared across all benchmark runners."""

    base_api_port: int = 8000
    pool_size: int = 1
    timeout_seconds: int = 600
    batch_size: int = 10
    cleanup_workspace: bool = False
    use_get_context: bool = False
    reasoning_level: ReasoningLevel | None = None
    base_url: str | None = None
    api_key: str | None = None
    skip_dream: bool = False
    json_output: Path | None = None
    max_concurrent: int | None = None  # None means no limit (use batch_size)

    @classmethod
    def from_args(
        cls, args: argparse.Namespace, default_timeout: int = 600
    ) -> "RunnerConfig":
        """Create config from parsed CLI arguments."""
        return cls(
            base_api_port=args.base_api_port,
            pool_size=args.pool_size,
            timeout_seconds=args.timeout
            if args.timeout is not None
            else default_timeout,
            batch_size=args.batch_size,
            cleanup_workspace=args.cleanup_workspace,
            use_get_context=args.use_get_context,
            reasoning_level=args.reasoning_level,
            base_url=args.base_url,
            api_key=args.api_key,
            skip_dream=args.skip_dream,
            json_output=args.json_output,
            max_concurrent=args.max_concurrent,
        )


@dataclass
class ItemContext:
    """Context for executing a single benchmark item."""

    workspace_id: str
    honcho_client: Honcho
    honcho_url: str
    session_id: str
    peers: dict[str, Any] = field(default_factory=dict)
    session: Any = None


def add_common_arguments(parser: argparse.ArgumentParser) -> None:
    """
    Add common command line arguments shared across all benchmark runners.

    Args:
        parser: ArgumentParser to add arguments to
    """
    parser.add_argument(
        "--base-url",
        type=str,
        default=None,
        help="Base URL for remote Honcho instance (e.g., https://api.example.com). Overrides --base-api-port.",
    )

    parser.add_argument(
        "--api-key",
        type=str,
        default=None,
        help="API key for remote Honcho instance authentication",
    )

    parser.add_argument(
        "--base-api-port",
        type=int,
        default=8000,
        help="Base port for Honcho API instances (default: 8000)",
    )

    parser.add_argument(
        "--pool-size",
        type=int,
        default=1,
        help="Number of Honcho instances in the pool (default: 1)",
    )

    parser.add_argument(
        "--timeout",
        type=int,
        default=None,
        help="Timeout for deriver queue to empty in seconds (default: 10 minutes)",
    )

    parser.add_argument(
        "--batch-size",
        type=int,
        default=10,
        help="Number of items to run concurrently in each batch (default: 10)",
    )

    parser.add_argument(
        "--json-output",
        type=Path,
        help="Path to write JSON summary results for analytics (optional)",
    )

    parser.add_argument(
        "--cleanup-workspace",
        action="store_true",
        help="Delete workspace after executing each test (default: False)",
    )

    parser.add_argument(
        "--use-get-context",
        action="store_true",
        help="Use get_context + judge LLM instead of dialectic .chat endpoint (default: False)",
    )

    parser.add_argument(
        "--reasoning-level",
        type=str,
        choices=REASONING_LEVELS,
        default=None,
        help="Reasoning level for dialectic chat: minimal, low, medium, high, max (default: None)",
    )

    parser.add_argument(
        "--skip-dream",
        action="store_true",
        help="Skip the dream consolidation step (default: False)",
    )

    parser.add_argument(
        "--max-concurrent",
        type=int,
        default=None,
        help="Maximum concurrent items executing at once (default: unlimited, use for rate-limited remote instances)",
    )


def validate_common_arguments(args: argparse.Namespace) -> str | None:
    """
    Validate common command line arguments.

    Args:
        args: Parsed arguments

    Returns:
        Error message if validation fails, None otherwise
    """
    if args.batch_size <= 0:
        return f"Error: Batch size must be positive, got {args.batch_size}"

    if args.pool_size <= 0:
        return f"Error: Pool size must be positive, got {args.pool_size}"

    if args.max_concurrent is not None and args.max_concurrent <= 0:
        return f"Error: Max concurrent must be positive, got {args.max_concurrent}"

    return None


def configure_logging(
    level: int = logging.WARNING,
    name: str | None = None,
) -> logging.Logger:
    """
    Configure logging for benchmark runners.

    Sets up logging at the specified level and suppresses HTTP request logs.

    Args:
        level: Logging level (default: WARNING for runners, use INFO for evaluators)
        name: Logger name (default: caller's __name__)

    Returns:
        Logger instance
    """
    logging.basicConfig(level=level, format="%(asctime)s - %(levelname)s - %(message)s")
    # Suppress HTTP request logs from the Honcho SDK
    logging.getLogger("httpx").setLevel(logging.ERROR)
    logging.getLogger("httpcore").setLevel(logging.ERROR)
    return logging.getLogger(name or __name__)


def create_anthropic_client(api_key: str | None = None) -> AsyncAnthropic:
    """
    Create an AsyncAnthropic client.

    Args:
        api_key: Optional API key. If not provided, uses LLM_ANTHROPIC_API_KEY env var.

    Returns:
        AsyncAnthropic client instance

    Raises:
        ValueError: If no API key is available
    """
    if api_key:
        return AsyncAnthropic(api_key=api_key)

    env_key = os.getenv("LLM_ANTHROPIC_API_KEY")
    if not env_key:
        raise ValueError("LLM_ANTHROPIC_API_KEY is not set")
    return AsyncAnthropic(api_key=env_key)


def create_openai_client(
    api_key: str | None = None,
    base_url: str | None = None,
    env_key_name: str = "OPENAI_API_KEY",
) -> AsyncOpenAI:
    """
    Create an AsyncOpenAI client.

    Args:
        api_key: Optional API key. If not provided, uses env_key_name env var.
        base_url: Optional base URL for OpenAI-compatible APIs (e.g., OpenRouter).
        env_key_name: Name of the environment variable for the API key.

    Returns:
        AsyncOpenAI client instance

    Raises:
        ValueError: If no API key is available
    """
    key = api_key or os.getenv(env_key_name)
    if not key:
        raise ValueError(f"{env_key_name} is not set")

    if base_url:
        return AsyncOpenAI(api_key=key, base_url=base_url)
    return AsyncOpenAI(api_key=key)


def format_duration(seconds: float) -> str:
    """Format a duration in seconds to a human-readable string."""
    if seconds < 60:
        return f"{seconds:.2f}s"
    elif seconds < 3600:
        minutes = int(seconds // 60)
        secs = seconds % 60
        return f"{minutes}m {secs:.1f}s"
    else:
        hours = int(seconds // 3600)
        minutes = int((seconds % 3600) // 60)
        return f"{hours}h {minutes}m"


# ---------------------------------------------------------------------------
# Trace parsing utilities (shared by trace-based evaluators)
# ---------------------------------------------------------------------------


def load_traces(path: Path) -> list[dict[str, Any]]:
    """Load traces from JSON or JSONL file.

    Attempts JSONL (one object per line) first; falls back to standard JSON
    (array or single object).  Returns an empty list when the file cannot be
    read or parsed.
    """
    traces: list[dict[str, Any]] = []

    # --- Attempt 1: try JSONL (one JSON object per line) ---
    try:
        with open(path) as f:
            first_line = f.readline().strip()
            if first_line and not first_line.startswith("["):
                f.seek(0)
                for line_no, line in enumerate(f, 1):
                    line = line.strip()
                    if not line:
                        continue
                    try:
                        traces.append(cast(dict[str, Any], json.loads(line)))
                    except json.JSONDecodeError:
                        _logger.debug(
                            "%s:%d: skipping malformed JSON line", path, line_no
                        )
                        continue
                if traces:
                    return traces
    except json.JSONDecodeError:
        _logger.warning("Failed to parse %s as JSONL", path)
    except OSError:
        _logger.exception("Could not read trace file %s", path)
        return []

    # --- Attempt 2: try standard JSON (array or single object) ---
    try:
        with open(path) as f:
            data = cast(dict[str, Any] | list[dict[str, Any]], json.load(f))
    except json.JSONDecodeError:
        _logger.exception("Failed to parse %s as JSON", path)
        return []
    except OSError:
        _logger.exception("Could not read trace file %s", path)
        return []

    if isinstance(data, list):
        return data
    return [data]


def extract_propositions(trace: dict[str, Any]) -> list[str]:
    """Extract propositions from trace output."""
    output = trace.get("output")
    if not isinstance(output, dict):
        return []
    output_d = cast(dict[str, Any], output)
    content = output_d.get("content")
    if not isinstance(content, dict):
        return []
    content_d = cast(dict[str, Any], content)
    explicit = content_d.get("explicit")
    if not isinstance(explicit, list):
        return []
    results: list[str] = []
    for raw_item in cast(list[Any], explicit):
        item = cast(dict[str, Any], raw_item)
        if isinstance(raw_item, dict) and "content" in item:
            results.append(str(item["content"]))
    return results


def extract_messages(trace: dict[str, Any]) -> list[dict[str, Any]]:
    """Extract source messages from trace input."""
    input_data = trace.get("input")
    if not isinstance(input_data, dict):
        return []
    input_d = cast(dict[str, Any], input_data)
    prompt = input_d.get("prompt", "")
    if not isinstance(prompt, str):
        return []

    messages: list[dict[str, Any]] = []
    if "<messages>" in prompt:
        section = prompt.split("<messages>")[1].split("</messages>")[0]
        for line in section.strip().split("\n"):
            line = line.strip()
            if not line:
                continue
            parts = line.split(" ", 3)
            if len(parts) >= 3:
                speaker = parts[2].rstrip(":")
                text = parts[3] if len(parts) > 3 else ""
                if "->->" in text:
                    text = text.split("->->")[0].strip()
                messages.append({"speaker": speaker, "text": text})
    return messages


def extract_peer_name(trace: dict[str, Any]) -> str:
    """Extract peer name from propositions."""
    for prop in extract_propositions(trace):
        prop_lower = prop.lower()
        for pattern in ["name is", "is named", "called"]:
            if pattern in prop_lower:
                parts = prop_lower.split(pattern)
                if len(parts) > 1:
                    name = parts[1].strip().rstrip(".").split()[0]
                    return name.capitalize()
    return "user"


def extract_conversation_id(trace: dict[str, Any], index: int) -> str:
    """Extract or generate conversation ID."""
    return trace.get("conversation_id", trace.get("id", f"trace_{index:04d}"))


class BaseRunner(ABC, Generic[ResultT]):
    """
    Abstract base class for benchmark runners.

    Provides a template method pattern for executing benchmarks with common
    infrastructure for Honcho client management, queue waiting, and dream triggering.

    Subclasses must implement:
    - get_metrics_prefix(): Return the metrics prefix (e.g., "longmem")
    - load_items(): Load and return the items to process
    - get_workspace_id(item): Return workspace ID for an item
    - get_session_id(item): Return session ID for an item
    - setup_peers(ctx, item): Create and configure peers
    - setup_session(ctx, item): Create and configure session with peers
    - ingest_messages(ctx, item): Ingest messages into the session
    - get_dream_observers(item): Return list of peer IDs to trigger dreams for
    - execute_questions(ctx, item): Execute questions and return result
    - print_summary(results, duration): Print summary of results
    - generate_output(results, duration): Generate JSON output
    """

    def __init__(self, config: RunnerConfig):
        """
        Initialize the runner with configuration.

        Args:
            config: Runner configuration
        """
        self.config: RunnerConfig = config
        self.metrics_collector: MetricsCollector = MetricsCollector()
        self.metrics_collector.start_collection(
            f"{self.get_metrics_prefix()}_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
        )
        self.logger: Logger = configure_logging()

    # -------------------------------------------------------------------------
    # Abstract methods - must be implemented by subclasses
    # -------------------------------------------------------------------------

    @abstractmethod
    def get_metrics_prefix(self) -> str:
        """Return the metrics prefix for this runner (e.g., 'longmem', 'beam')."""
        ...

    @abstractmethod
    def load_items(self) -> list[Any]:
        """Load and return the list of items to process."""
        ...

    @abstractmethod
    def get_workspace_id(self, item: Any) -> str:
        """Return the workspace ID for a given item."""
        ...

    @abstractmethod
    def get_session_id(self, item: Any, workspace_id: str) -> str:
        """Return the session ID for a given item."""
        ...

    @abstractmethod
    async def setup_peers(self, ctx: ItemContext, item: Any) -> None:
        """
        Create and configure peers for the item.

        Should populate ctx.peers with peer objects.
        """
        ...

    @abstractmethod
    async def setup_session(self, ctx: ItemContext, item: Any) -> None:
        """
        Create and configure the session with peers.

        Should set ctx.session and add peers to the session.
        """
        ...

    @abstractmethod
    async def ingest_messages(self, ctx: ItemContext, item: Any) -> int:
        """
        Ingest messages into the session.

        Returns:
            Number of messages ingested
        """
        ...

    @abstractmethod
    def get_dream_observers(self, item: Any) -> list[str]:
        """Return list of peer IDs to trigger dreams for."""
        ...

    def get_dream_session_ids(self, ctx: ItemContext, _item: Any) -> list[str]:
        """Return session IDs to use for dream scheduling.

        Subclasses can override this when ingestion stores messages across
        multiple sessions and each session should be included in dream
        scheduling.
        """
        return [ctx.session_id]

    @abstractmethod
    async def execute_questions(self, ctx: ItemContext, item: Any) -> ResultT:
        """
        Execute questions/queries for the item.

        Returns:
            Result object for this item
        """
        ...

    @abstractmethod
    def print_summary(self, results: list[ResultT], total_duration: float) -> None:
        """Print a summary of all results."""
        ...

    @abstractmethod
    def generate_output(self, results: list[ResultT], total_duration: float) -> None:
        """Generate JSON output file."""
        ...

    # -------------------------------------------------------------------------
    # Template method - the main execution flow
    # -------------------------------------------------------------------------

    async def run(self) -> tuple[list[ResultT], float]:
        """
        Run the benchmark.

        This is the main template method that orchestrates the execution flow.

        Returns:
            Tuple of (list of results, total duration in seconds)
        """
        items = self.load_items()
        if not items:
            return [], 0.0

        print(f"Found {len(items)} items to process")
        if self.config.pool_size > 1:
            print(
                f"Distributing across {self.config.pool_size} Honcho instances "
                + f"(ports {self.config.base_api_port}-{self.config.base_api_port + self.config.pool_size - 1})"
            )
        if self.config.max_concurrent:
            print(f"Limiting to {self.config.max_concurrent} concurrent item(s)")

        overall_start = time.time()
        all_results: list[ResultT | None] = [None] * len(items)

        # Two-level concurrency:
        # - inflight_sem limits how many items may be in the pipeline at once
        # - active_sem limits how many items may actively hit Honcho at once
        # Items release active_sem while waiting on queue polling so other work
        # can progress, but inflight_sem prevents an unlimited thundering herd.
        concurrency = self.config.max_concurrent or self.config.batch_size
        inflight_sem = asyncio.Semaphore(concurrency)
        active_sem = asyncio.Semaphore(concurrency)

        async def _run_item(index: int, item: Any) -> None:
            async with inflight_sem:
                result = await self.execute_item(
                    item,
                    self._get_honcho_url(index),
                    active_sem=active_sem,
                )
                all_results[index] = result

        tasks = [
            asyncio.create_task(_run_item(index, item))
            for index, item in enumerate(items)
        ]
        await asyncio.gather(*tasks)

        overall_duration = time.time() - overall_start

        # Finalize metrics
        self.metrics_collector.finalize_collection()

        missing_indexes = [
            index for index, result in enumerate(all_results) if result is None
        ]
        if missing_indexes:
            raise RuntimeError(
                f"Missing benchmark results for item indexes: {missing_indexes}"
            )

        return [cast(ResultT, result) for result in all_results], overall_duration

    async def execute_item(
        self,
        item: Any,
        honcho_url: str,
        active_sem: asyncio.Semaphore | None = None,
    ) -> ResultT:
        """
        Execute a single benchmark item.

        Active work (setup, ingest, dream scheduling, query execution) acquires
        ``active_sem`` when provided. Idle queue polling releases that slot so
        other items can continue making forward progress.

        Args:
            item: The item to process
            honcho_url: URL of the Honcho instance to use
            active_sem: Optional semaphore limiting active I/O phases

        Returns:
            Result for this item
        """
        workspace_id = self.get_workspace_id(item)
        session_id = self.get_session_id(item, workspace_id)

        print(f"\n{'=' * 80}")
        print(f"Executing {workspace_id}")
        print(f"Using Honcho instance: {honcho_url}")
        print(f"{'=' * 80}")

        # Create context
        ctx = ItemContext(
            workspace_id=workspace_id,
            honcho_client=self._create_honcho_client(workspace_id, honcho_url),
            honcho_url=honcho_url,
            session_id=session_id,
        )

        start_time = time.time()

        try:
            # Setup peers/session and ingest under the active semaphore.
            if active_sem:
                await active_sem.acquire()
            try:
                await self.setup_peers(ctx, item)
                await self.setup_session(ctx, item)

                print(f"[{workspace_id}] Ingesting messages...")
                message_count = await self.ingest_messages(ctx, item)
                print(f"[{workspace_id}] Ingested {message_count} messages")
            finally:
                if active_sem:
                    active_sem.release()

            # Wait for deriver queue
            print(f"[{workspace_id}] Waiting for deriver queue to empty...")
            queue_empty = await self._wait_for_queue_empty(ctx.honcho_client)
            if not queue_empty:
                raise TimeoutError(
                    f"Deriver queue timeout after {self.config.timeout_seconds}s"
                )

            # Trigger dreams
            dream_observers = self.get_dream_observers(item)
            dream_session_ids = self.get_dream_session_ids(ctx, item)
            if not dream_session_ids:
                raise ValueError(
                    f"No dream session IDs available for {workspace_id}. "
                    + "Dream scheduling requires at least one session id."
                )

            print(
                f"[{workspace_id}] Deriver queue empty. Triggering dreams for "
                + f"{len(dream_observers)} observer(s) across "
                + f"{len(dream_session_ids)} session(s)..."
            )

            if self.config.skip_dream:
                print(f"[{workspace_id}] Skipping dreams (--skip-dream)")
            else:

                async def _schedule_dream(
                    observer: str,
                    session_id: str,
                ) -> bool:
                    try:
                        if active_sem:
                            await active_sem.acquire()
                        try:
                            await ctx.honcho_client.aio.schedule_dream(
                                observer=observer,
                                session=session_id,
                                observed=observer,
                            )
                        finally:
                            if active_sem:
                                active_sem.release()
                        print(
                            f"[{workspace_id}] Dream triggered for "
                            + f"{observer}/{observer} in {session_id}"
                        )
                        return True
                    except Exception as e:
                        print(
                            f"[{workspace_id}] ERROR: Dream trigger exception "
                            + f"for {observer} in {session_id}: {e}"
                        )
                        return False

                dream_results = await asyncio.gather(
                    *[
                        _schedule_dream(observer, dream_session_id)
                        for observer in dream_observers
                        for dream_session_id in dream_session_ids
                    ]
                )

                if all(dream_results):
                    success = await self._wait_for_queue_empty(ctx.honcho_client)
                    if success:
                        print(f"[{workspace_id}] All dreams completed")
                    else:
                        print(f"[{workspace_id}] Dreams timed out")
                elif any(dream_results):
                    failed = [i for i, ok in enumerate(dream_results) if not ok]
                    print(
                        f"[{workspace_id}] Warning: {len(failed)} of "
                        + f"{len(dream_results)} dream schedules failed"
                    )
                    await self._wait_for_queue_empty(ctx.honcho_client)
                else:
                    print(f"[{workspace_id}] Warning: No dreams were scheduled")

            # Execute questions
            print(f"[{workspace_id}] Executing questions...")
            if active_sem:
                await active_sem.acquire()
            try:
                result = await self.execute_questions(ctx, item)
            finally:
                if active_sem:
                    active_sem.release()

            # Cleanup
            if self.config.cleanup_workspace:
                try:
                    await ctx.honcho_client.aio.delete_workspace(workspace_id)
                    print(f"[{workspace_id}] Cleaned up workspace")
                except Exception as e:
                    print(f"[{workspace_id}] Failed to delete workspace: {e}")

            duration = time.time() - start_time
            print(f"[{workspace_id}] Completed in {format_duration(duration)}")

            return result

        except Exception as e:
            self.logger.error(f"Error executing {workspace_id}: {e}")
            # Let subclass handle error result creation
            raise

    def run_and_summarize(self) -> int:
        """
        Run the benchmark, print summary, and generate output.

        This is a convenience method that runs the full benchmark flow.

        Returns:
            Exit code (0 for success, 1 for failure)
        """
        try:
            results, total_duration = asyncio.run(self.run())

            self.print_summary(results, total_duration)
            self.metrics_collector.print_summary()
            self.generate_output(results, total_duration)

            # Export metrics
            metrics_output = Path(
                f"tests/bench/perf_metrics/{self.get_metrics_prefix()}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json"
            )
            self.metrics_collector.export_to_json(metrics_output)
            self.metrics_collector.cleanup_collection()

            return 0

        except KeyboardInterrupt:
            print("\nTest execution interrupted by user")
            return 1
        except Exception as e:
            print(f"Error running tests: {e}")
            import traceback

            traceback.print_exc()
            return 1

    # -------------------------------------------------------------------------
    # Infrastructure methods
    # -------------------------------------------------------------------------

    def _get_honcho_url(self, index: int) -> str:
        """Get the Honcho URL for a given index using round-robin distribution."""
        if self.config.base_url:
            return self.config.base_url
        instance_id = index % self.config.pool_size
        port = self.config.base_api_port + instance_id
        return f"http://localhost:{port}"

    def _create_honcho_client(self, workspace_id: str, honcho_url: str) -> Honcho:
        """Create a Honcho client for a specific workspace."""
        return Honcho(
            workspace_id=workspace_id,
            base_url=honcho_url,
            api_key=self.config.api_key,
        )

    def _get_session_configuration(self) -> SessionConfiguration:
        """Get default session configuration with summaries disabled."""
        return SessionConfiguration(summary=SummaryConfiguration(enabled=False))

    async def _wait_for_queue_empty(
        self, honcho_client: Honcho, session_id: str | None = None
    ) -> bool:
        """Wait for the deriver queue to be empty with exponential backoff."""
        start_time = time.time()
        delay = 0.2
        while True:
            try:
                status = await honcho_client.aio.queue_status(session=session_id)
            except Exception:
                await asyncio.sleep(delay)
                delay = min(delay * 1.5, 2.0)
                if time.time() - start_time >= self.config.timeout_seconds:
                    return False
                continue

            if status.pending_work_units == 0 and status.in_progress_work_units == 0:
                return True

            if time.time() - start_time >= self.config.timeout_seconds:
                return False
            await asyncio.sleep(delay)
            delay = min(delay * 1.5, 2.0)

    async def _trigger_dream(
        self,
        honcho_client: Honcho,
        workspace_id: str,
        observer: str,
        session_id: str,
        observed: str | None = None,
    ) -> bool:
        """
        Trigger a dream task and wait for it to complete.

        Returns:
            True if dream completed (or was skipped), False on timeout
        """
        if self.config.skip_dream:
            print(f"[{workspace_id}] Skipping dream for {observer} (--skip-dream)")
            return True

        observed = observed or observer

        try:
            await honcho_client.aio.schedule_dream(
                observer=observer,
                session=session_id,
                observed=observed,
            )
        except Exception as e:
            print(f"[{workspace_id}] ERROR: Dream trigger exception: {e}")
            return False

        print(f"[{workspace_id}] Dream triggered for {observer}/{observed}")

        success = await self._wait_for_queue_empty(honcho_client)
        if success:
            print(f"[{workspace_id}] Dream for {observer} completed")
        else:
            print(f"[{workspace_id}] Dream for {observer} timed out")
        return success