import asyncio
from dataclasses import dataclass
from pathlib import Path

from claude_agent_sdk import (
    AssistantMessage,
    ClaudeAgentOptions,
    ClaudeSDKClient,
    TextBlock,
)
from rich.console import Group
from rich.live import Live
from rich.progress import (
    BarColumn,
    MofNCompleteColumn,
    Progress,
    SpinnerColumn,
    TaskID,
    TextColumn,
    TimeElapsedColumn,
    TimeRemainingColumn,
)

from harbor.models.trial.config import TrialConfig
from harbor.models.trial.paths import TrialPaths
from harbor.models.trial.result import TrialResult


@dataclass
class SummarizeResult:
    """Result of a summarization operation."""

    summary_path: str
    n_trials_summarized: int
    job_summary_created: bool


class Summarizer:
    SUMMARIZE_TRIAL_PROMPT = (Path(__file__).parent / "summarize-trial.txt").read_text()
    SUMMARIZE_JOB_PROMPT = (Path(__file__).parent / "summarize-job.txt").read_text()

    def __init__(
        self,
        job_dir: Path,
        n_concurrent: int = 5,
        model: str | None = None,
        only_failed: bool = False,
        overwrite: bool = False,
    ):
        """
        Initialize the summarizer.

        Args:
            job_dir: Path to the job directory containing trial subdirectories
            n_concurrent: Maximum number of concurrent summarization queries
            model: Model to use for summarization (e.g., 'haiku', 'sonnet', 'opus')
            only_failed: Only analyze trials that failed (exceptions/errors or reward = 0)
            overwrite: Overwrite existing summary.md files (by default, skip trials with existing summaries)
        """
        self.job_dir = job_dir
        self.n_concurrent = n_concurrent
        self.model = model
        self.only_failed = only_failed
        self.overwrite = overwrite

        # Get all trial directories with results
        all_trial_dirs = [
            d for d in job_dir.iterdir() if d.is_dir() and (d / "result.json").exists()
        ]

        # Filter trials based on criteria
        self.trial_dirs = self._filter_trials(all_trial_dirs)

    def _filter_trials(self, trial_dirs: list[Path]) -> list[Path]:
        """
        Filter trials based on the specified criteria.

        Args:
            trial_dirs: List of trial directories to filter

        Returns:
            Filtered list of trial directories
        """
        # First, filter out trials with existing summaries (unless overwrite=True)
        if not self.overwrite:
            trial_dirs = [d for d in trial_dirs if not (d / "summary.md").exists()]

        if not self.only_failed:
            # No additional filters, return trials
            return trial_dirs

        filtered_dirs = []
        for trial_dir in trial_dirs:
            trial_paths = TrialPaths(trial_dir)
            if not trial_paths.result_path.exists():
                continue

            try:
                trial_result = TrialResult.model_validate_json(
                    trial_paths.result_path.read_text()
                )

                # Check if trial has reward = 1 (successful)
                has_reward_one = False
                if (
                    trial_result.verifier_result
                    and trial_result.verifier_result.rewards
                ):
                    # Check if any reward is 1
                    has_reward_one = any(
                        reward == 1
                        for reward in trial_result.verifier_result.rewards.values()
                    )

                # Skip trials with reward = 1 even if they have exceptions
                if has_reward_one:
                    continue

                # Include trials with exceptions or reward = 0
                has_exception = trial_result.exception_info is not None
                has_reward_zero = False

                if (
                    trial_result.verifier_result
                    and trial_result.verifier_result.rewards
                ):
                    # Check if any reward is 0
                    has_reward_zero = any(
                        reward == 0
                        for reward in trial_result.verifier_result.rewards.values()
                    )

                if has_exception or has_reward_zero:
                    filtered_dirs.append(trial_dir)

            except Exception as e:
                print(f"Warning: Could not parse result for {trial_dir.name}: {e}")
                continue

        return filtered_dirs

    async def summarize_trial(self, trial_dir: Path) -> tuple[str, str]:
        """
        Summarize a single trial using Claude Agent SDK with full read-only access.

        Args:
            trial_dir: Path to the trial directory

        Returns:
            Tuple of (trial_name, summary_text)
        """

        trial_name = trial_dir.name

        trial_paths = TrialPaths(trial_dir)

        task_dir = None

        if trial_paths.config_path.exists():
            trial_config = TrialConfig.model_validate_json(
                trial_paths.config_path.read_text()
            )
            task_dir = trial_config.task.get_task_id().get_local_path()

        prompt = self.SUMMARIZE_TRIAL_PROMPT.format(
            trial_dir=trial_dir,
            task_dir=task_dir or "N/A",
        )

        options = ClaudeAgentOptions(
            permission_mode="bypassPermissions",
            allowed_tools=["Read", "Glob"],
            cwd=str(trial_dir),
            add_dirs=[str(task_dir)] if task_dir else [],
            model=self.model,
        )

        summary_parts = []
        async with ClaudeSDKClient(options=options) as client:
            await client.query(prompt)

            async for message in client.receive_response():
                if isinstance(message, AssistantMessage):
                    for block in message.content:
                        if isinstance(block, TextBlock):
                            summary_parts.append(block.text)

        summary = "\n".join(summary_parts)

        summary_path = trial_dir / "summary.md"
        summary_path.write_text(f"# Summary: {trial_name}\n\n{summary}\n")

        return trial_name, summary

    async def _summarize_trial_with_progress(
        self,
        trial_dir: Path,
        semaphore: asyncio.Semaphore,
        loading_progress: Progress,
        loading_progress_task: TaskID,
        running_progress: Progress,
    ) -> tuple[str, str]:
        """
        Summarize a single trial with progress tracking.

        Args:
            trial_dir: Path to the trial directory
            semaphore: Semaphore to control concurrency
            loading_progress: Overall progress bar
            loading_progress_task: Task ID for the loading progress
            running_progress: Progress bar for currently running trials

        Returns:
            Tuple of (trial_name, summary_text)
        """
        async with semaphore:
            trial_name = trial_dir.name
            trial_progress_task = running_progress.add_task(
                f"{trial_name}: summarizing...", total=None
            )

            try:
                trial_name, summary = await self.summarize_trial(trial_dir)
                return trial_name, summary
            finally:
                running_progress.remove_task(trial_progress_task)
                loading_progress.advance(loading_progress_task)

    async def _summarize_all_trials(self) -> dict[str, str]:
        """
        Summarize all trials in parallel using a semaphore to control concurrency.

        Returns:
            Dictionary mapping trial names to their summaries
        """
        semaphore = asyncio.Semaphore(self.n_concurrent)
        summaries = {}

        # Create progress bars
        loading_progress = Progress(
            SpinnerColumn(),
            MofNCompleteColumn(),
            TextColumn("[progress.description]{task.description}"),
            BarColumn(),
            TimeElapsedColumn(),
            TimeRemainingColumn(),
        )

        running_progress = Progress(
            SpinnerColumn(),
            TimeElapsedColumn(),
            TextColumn("[progress.description]{task.description}"),
        )

        with Live(Group(loading_progress, running_progress), refresh_per_second=10):
            progress_task = loading_progress.add_task(
                "Summarizing trials...", total=len(self.trial_dirs)
            )

            async with asyncio.TaskGroup() as tg:
                tasks = [
                    tg.create_task(
                        self._summarize_trial_with_progress(
                            trial_dir,
                            semaphore,
                            loading_progress,
                            progress_task,
                            running_progress,
                        )
                    )
                    for trial_dir in self.trial_dirs
                ]

            # Collect results
            for task in tasks:
                trial_name, summary = task.result()
                summaries[trial_name] = summary

        return summaries

    def _get_relevant_trial_dirs(self) -> list[Path]:
        """
        Get all trial directories that are relevant for job summary.

        Returns only failed trials if only_failed=True, otherwise all with results.
        """
        all_trial_dirs = [
            d
            for d in self.job_dir.iterdir()
            if d.is_dir() and (d / "result.json").exists()
        ]

        if not self.only_failed:
            return all_trial_dirs

        # Filter for failed trials
        failed_dirs = []
        for trial_dir in all_trial_dirs:
            trial_paths = TrialPaths(trial_dir)
            if not trial_paths.result_path.exists():
                continue

            try:
                trial_result = TrialResult.model_validate_json(
                    trial_paths.result_path.read_text()
                )

                # Check if trial has reward = 1 (successful)
                has_reward_one = False
                if (
                    trial_result.verifier_result
                    and trial_result.verifier_result.rewards
                ):
                    has_reward_one = any(
                        reward == 1
                        for reward in trial_result.verifier_result.rewards.values()
                    )

                # Skip trials with reward = 1
                if has_reward_one:
                    continue

                # Include trials with exceptions or reward = 0
                has_exception = trial_result.exception_info is not None
                has_reward_zero = False

                if (
                    trial_result.verifier_result
                    and trial_result.verifier_result.rewards
                ):
                    has_reward_zero = any(
                        reward == 0
                        for reward in trial_result.verifier_result.rewards.values()
                    )

                if has_exception or has_reward_zero:
                    failed_dirs.append(trial_dir)

            except Exception:
                continue

        return failed_dirs

    async def _create_job_summary(self) -> str:
        """
        Create a job-level summary by reading ALL relevant trial summaries from disk.

        Returns:
            Job-level summary text
        """
        # Get all relevant trial directories (failed if only_failed=True, else all)
        relevant_trial_dirs = self._get_relevant_trial_dirs()

        # Read summaries from disk for all relevant trials
        trial_summaries: dict[str, str] = {}
        for trial_dir in relevant_trial_dirs:
            summary_path = trial_dir / "summary.md"
            if summary_path.exists():
                try:
                    trial_summaries[trial_dir.name] = summary_path.read_text()
                except Exception:
                    continue

        if not trial_summaries:
            return ""

        # Build aggregated summary from individual trial summaries
        summary_text_parts = []
        summary_text_parts.append(f"Total Trials: {len(trial_summaries)}\n")
        summary_text_parts.append("## Individual Trial Summaries:\n")

        for trial_name, summary in sorted(trial_summaries.items()):
            summary_text_parts.append(f"\n### {trial_name}\n")
            summary_text_parts.append(summary)
            summary_text_parts.append("")

        aggregated_summaries = "\n".join(summary_text_parts)

        prompt = self.SUMMARIZE_JOB_PROMPT.format(
            summaries=aggregated_summaries,
        )

        options = ClaudeAgentOptions(
            permission_mode="bypassPermissions",
            allowed_tools=[],
            model=self.model,
        )

        summary_parts = []
        async with ClaudeSDKClient(options=options) as client:
            await client.query(prompt)

            async for message in client.receive_response():
                if isinstance(message, AssistantMessage):
                    for block in message.content:
                        if isinstance(block, TextBlock):
                            summary_parts.append(block.text)

        summary = "\n".join(summary_parts)

        summary_path = self.job_dir / "summary.md"
        summary_path.write_text(f"{summary}\n")

        return summary

    async def summarize_async(self) -> SummarizeResult:
        """
        Asynchronously summarize all trials and create a job-level summary.

        Returns:
            SummarizeResult with summary path and counts
        """
        # Print filtering information
        all_trial_dirs = [
            d
            for d in self.job_dir.iterdir()
            if d.is_dir() and (d / "result.json").exists()
        ]
        total_trials = len(all_trial_dirs)
        filtered_count = len(self.trial_dirs)

        if filtered_count < total_trials:
            filter_parts = []
            if not self.overwrite:
                existing_summaries = sum(
                    1 for d in all_trial_dirs if (d / "summary.md").exists()
                )
                if existing_summaries > 0:
                    filter_parts.append(f"{existing_summaries} with existing summaries")
            if self.only_failed:
                filter_parts.append("only failed trials")

            filter_desc = ", ".join(filter_parts) if filter_parts else ""
            print(
                f"\nFiltering: {filtered_count}/{total_trials} trials"
                + (f" ({filter_desc})" if filter_desc else "")
            )

        n_trials_summarized = 0
        job_summary_created = False
        summary_path = self.job_dir / "summary.md"

        # Summarize trials that need it
        if filtered_count > 0:
            await self._summarize_all_trials()
            n_trials_summarized = filtered_count

        # Check if there are any relevant trial summaries on disk
        relevant_trial_dirs = self._get_relevant_trial_dirs()
        has_any_summaries = any(
            (d / "summary.md").exists() for d in relevant_trial_dirs
        )

        if has_any_summaries:
            # Create job-level summary with progress display
            job_progress = Progress(
                SpinnerColumn(),
                TextColumn("[progress.description]{task.description}"),
                TimeElapsedColumn(),
            )

            with job_progress:
                job_task = job_progress.add_task(
                    "Creating job-level summary...", total=None
                )
                result = await self._create_job_summary()
                job_progress.remove_task(job_task)

            if result:
                job_summary_created = True
                print(f"\n✓ Job summary written to: {summary_path}")
        else:
            msg = "\n⚠ No trial summaries to aggregate."
            if not self.overwrite:
                msg += " Use --overwrite to regenerate existing summaries."
            print(msg)

        return SummarizeResult(
            summary_path=str(summary_path) if job_summary_created else "",
            n_trials_summarized=n_trials_summarized,
            job_summary_created=job_summary_created,
        )

    def summarize(self) -> str:
        """
        Synchronous wrapper for summarize_async.

        Returns:
            Path to the job-level summary file
        """
        result = asyncio.run(self.summarize_async())
        return result.summary_path
