import asyncio
from dataclasses import dataclass, field
from pathlib import Path

from claude_agent_sdk import (
    ClaudeAgentOptions,
    ResultMessage,
    query,
)
from pydantic import BaseModel
from rich.console import Group
from rich.live import Live
from rich.progress import (
    BarColumn,
    MofNCompleteColumn,
    Progress,
    SpinnerColumn,
    TaskID,
    TextColumn,
    TimeElapsedColumn,
    TimeRemainingColumn,
)
from tenacity import retry, stop_after_attempt, wait_exponential

from harbor.models.task.config import TaskConfig

ANNOTATE_PROMPT = (Path(__file__).parent / "annotate-task.md").read_text()

SUPPORTED_MODELS = {"sonnet", "opus", "haiku"}


class AnnotateOutput(BaseModel):
    readme: str
    description: str


@dataclass
class AnnotateResult:
    annotated: int = 0
    skipped: int = 0
    failed: int = 0
    errors: list[str] = field(default_factory=list)


class Annotator:
    def __init__(
        self,
        task_dirs: list[Path],
        n_concurrent: int = 5,
        model: str | None = None,
        overwrite: bool = False,
    ):
        if model and model not in SUPPORTED_MODELS:
            raise ValueError(
                f"Unsupported model '{model}'. "
                f"Must be one of: {', '.join(sorted(SUPPORTED_MODELS))}"
            )

        self.task_dirs = task_dirs
        self.n_concurrent = n_concurrent
        self.model = model
        self.overwrite = overwrite

    def _build_file_tree(self, task_dir: Path) -> str:
        lines = []
        for path in sorted(task_dir.rglob("*")):
            if not path.is_file():
                continue
            rel_path = path.relative_to(task_dir)
            lines.append(str(rel_path))
        return "\n".join(lines) if lines else "No files found"

    @retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, max=30))
    async def annotate_task(self, task_dir: Path) -> AnnotateOutput:
        file_tree = self._build_file_tree(task_dir)
        prompt = ANNOTATE_PROMPT.format(task_dir=task_dir, file_tree=file_tree)

        options = ClaudeAgentOptions(
            permission_mode="bypassPermissions",
            allowed_tools=["Read", "Glob"],
            cwd=str(task_dir),
            model=self.model,
            output_format={
                "type": "json_schema",
                "schema": AnnotateOutput.model_json_schema(),
            },
        )

        structured_output = None
        async for message in query(prompt=prompt, options=options):
            if isinstance(message, ResultMessage):
                structured_output = message.structured_output

        if structured_output is None:
            raise ValueError("SDK did not return structured output")

        return AnnotateOutput.model_validate(structured_output)

    def _write_results(self, task_dir: Path, output: AnnotateOutput) -> None:
        # Write README.md
        (task_dir / "README.md").write_text(output.readme)

        # Update description in task.toml if [task] section exists
        config_path = task_dir / "task.toml"
        config = TaskConfig.model_validate_toml(config_path.read_text())
        if config.task is not None:
            config.task.description = output.description
            config_path.write_text(config.model_dump_toml())

    async def _annotate_with_progress(
        self,
        task_dir: Path,
        semaphore: asyncio.Semaphore,
        loading_progress: Progress,
        loading_task: TaskID,
        running_progress: Progress,
        result: AnnotateResult,
    ) -> None:
        readme_path = task_dir / "README.md"
        if readme_path.exists() and not self.overwrite:
            result.skipped += 1
            loading_progress.advance(loading_task)
            return

        async with semaphore:
            task_name = task_dir.name
            running_task = running_progress.add_task(
                f"{task_name}: annotating...", total=None
            )
            try:
                output = await self.annotate_task(task_dir)
                self._write_results(task_dir, output)
                result.annotated += 1
            except Exception as e:
                result.failed += 1
                result.errors.append(f"{task_name}: {e}")
            finally:
                running_progress.remove_task(running_task)
                loading_progress.advance(loading_task)

    async def run(self) -> AnnotateResult:
        result = AnnotateResult()
        semaphore = asyncio.Semaphore(self.n_concurrent)

        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):
            loading_task = loading_progress.add_task(
                "Annotating tasks...", total=len(self.task_dirs)
            )

            async with asyncio.TaskGroup() as tg:
                for task_dir in self.task_dirs:
                    tg.create_task(
                        self._annotate_with_progress(
                            task_dir,
                            semaphore,
                            loading_progress,
                            loading_task,
                            running_progress,
                            result,
                        )
                    )

        return result
