from __future__ import annotations

import json
from dataclasses import dataclass
from typing import Any, Dict, List, Optional

# Local import placed inside TYPE_CHECKING block to avoid runtime cycles.
from typing import TYPE_CHECKING

if TYPE_CHECKING:  # pragma: no cover – for static analysers only
    from ..api import _Result  # noqa: F401


@dataclass
class ResearchTaskId:
    """Structured research task ID.

    Attributes
    ----------
    id:
        Unique identifier for the research task.
    """

    id: str

    # ---------------------------------------------------------------------
    # Pretty representation helpers
    # ---------------------------------------------------------------------
    def __str__(self) -> str:  # pragma: no cover – convenience only
        return f"ID: {self.id}\n"


@dataclass
class ResearchTask:
    """Structured research task.

    Attributes
    ----------
    id:
        Unique identifier for the research task.
    status:
        Current task status
    instructions:
        Instructions for the task
    schema:
        Output schema defining the task
    data:
        JSON-serialisable answer generated by Exa (may be ``None`` until the task
        completes).
    citations:
        Mapping from *root field* in the output schema to the list of search
        results that were used to generate that part of the answer.
    """

    id: str
    status: str
    instructions: str
    schema: Dict[str, Any]
    data: Optional[Dict[str, Any]]
    citations: Dict[str, List["_Result"]]

    # ---------------------------------------------------------------------
    # Pretty representation helpers
    # ---------------------------------------------------------------------
    def __str__(self) -> str:  # pragma: no cover – convenience only
        """Human-readable representation including *all* relevant fields."""
        schema_repr = json.dumps(self.schema, indent=2, ensure_ascii=False)
        data_repr = (
            json.dumps(self.data, indent=2, ensure_ascii=False)
            if self.data is not None
            else "None"
        )

        # Render citations grouped by the root field they belong to.
        if self.citations:
            # Each key is a root field, each value is a list of _Result objects.
            citations_lines = []
            for field, sources in self.citations.items():
                rendered_sources = "\n    ".join(str(src) for src in sources)
                citations_lines.append(f"{field}:\n    {rendered_sources}")
            citations_str = "\n\n".join(citations_lines)
        else:
            citations_str = "None"

        return (
            f"ID: {self.id}\n"
            f"Status: {self.status}\n"
            f"Instructions: {self.instructions}\n"
            f"Schema:\n{schema_repr}\n"
            f"Data:\n{data_repr}\n\n"
            f"Citations:\n{citations_str}"
        )


__all__ = [
    "ResearchTaskId",
    "ResearchTask",
]
