from __future__ import annotations

import uuid
from collections.abc import Mapping, Sequence
from typing import Any, Literal, Union

from pydantic import AliasChoices
from pydantic import BaseModel, ConfigDict, Field
from pydantic import field_validator
from pydantic import model_validator


class RuntimeRerankControl(BaseModel):
    enabled: bool | None = None


class RuntimeQueryExpansionControl(BaseModel):
    enabled: bool | None = None


class RuntimeSubquestionHitlControl(BaseModel):
    enabled: bool = False


class RuntimeQueryExpansionHitlControl(BaseModel):
    enabled: bool = False


class RuntimeHitlControl(BaseModel):
    enabled: bool = False
    subquestions: RuntimeSubquestionHitlControl | None = None
    query_expansion: RuntimeQueryExpansionHitlControl | None = None

    @model_validator(mode="after")
    def normalize_enabled(self) -> "RuntimeHitlControl":
        if (
            (self.subquestions is not None and self.subquestions.enabled)
            or (self.query_expansion is not None and self.query_expansion.enabled)
        ):
            self.enabled = True
        elif self.enabled and self.subquestions is None and self.query_expansion is None:
            self.subquestions = RuntimeSubquestionHitlControl(enabled=True)
        return self


class RuntimeAgentRunControls(BaseModel):
    rerank: RuntimeRerankControl | None = None
    query_expansion: RuntimeQueryExpansionControl | None = None
    hitl: RuntimeHitlControl | None = None


class RuntimeAgentRunRuntimeConfig(BaseModel):
    rerank: RuntimeRerankControl | None = None
    query_expansion: RuntimeQueryExpansionControl | None = None


class RuntimeCustomPrompts(BaseModel):
    model_config = ConfigDict(extra="ignore")

    subanswer: str | None = None
    synthesis: str | None = None


class RuntimeAgentRunRequest(BaseModel):
    query: str = Field(min_length=1)
    thread_id: str | None = None
    controls: RuntimeAgentRunControls | None = None
    runtime_config: RuntimeAgentRunRuntimeConfig | None = None
    custom_prompts: RuntimeCustomPrompts | None = Field(
        default=None,
        validation_alias=AliasChoices("custom_prompts", "custom-prompts"),
    )

    @field_validator("thread_id")
    @classmethod
    def validate_thread_id(cls, value: str | None) -> str | None:
        if value is None:
            return None
        normalized_value = value.strip()
        if not normalized_value:
            raise ValueError("thread_id must be a non-empty UUID string.")
        try:
            return str(uuid.UUID(normalized_value))
        except ValueError as exc:
            raise ValueError("thread_id must be a valid UUID string.") from exc


class SubQuestionAnswer(BaseModel):
    sub_question: str
    sub_answer: str
    tool_call_input: str = ""
    expanded_query: str = ""
    sub_agent_response: str = ""
    answerable: bool = False
    verification_reason: str = ""


class RuntimeAgentRunResponse(BaseModel):
    main_question: str = ""
    thread_id: str = ""
    sub_qa: list[SubQuestionAnswer] = Field(default_factory=list)
    sub_answers: list[SubQuestionAnswer] = Field(default_factory=list)
    output: str
    final_citations: list["CitationSourceRow"] = Field(default_factory=list)


class HitlResumeDecision(BaseModel):
    item_id: str = Field(min_length=1)
    action: Literal["approve", "edit", "reject"]
    replacement_text: str | None = None

    @model_validator(mode="after")
    def validate_replacement_text(self) -> "HitlResumeDecision":
        if self.action == "edit":
            if self.replacement_text is None or not self.replacement_text.strip():
                raise ValueError("replacement_text is required when action='edit'.")
            self.replacement_text = self.replacement_text.strip()
            return self
        if self.replacement_text is not None:
            self.replacement_text = self.replacement_text.strip() or None
        return self


class HitlReviewItem(BaseModel):
    item_id: str = Field(min_length=1)
    text: str
    index: int = Field(ge=0)

    def approve(self) -> HitlResumeDecision:
        return HitlResumeDecision(item_id=self.item_id, action="approve")

    def edit(self, replacement_text: str) -> HitlResumeDecision:
        return HitlResumeDecision(item_id=self.item_id, action="edit", replacement_text=replacement_text)

    def reject(self) -> HitlResumeDecision:
        return HitlResumeDecision(item_id=self.item_id, action="reject")


class HitlReview(BaseModel):
    kind: Literal["subquestion_review", "query_expansion_review"]
    stage: str = ""
    checkpoint_id: str = Field(min_length=1)
    items: list[HitlReviewItem] = Field(default_factory=list)

    @classmethod
    def from_interrupt_payload(cls, payload: Any) -> "HitlReview":
        if not isinstance(payload, Mapping):
            raise ValueError("HITL interrupt payload must be a mapping.")
        kind = str(payload.get("kind") or "").strip()
        checkpoint_id = str(payload.get("checkpoint_id") or "").strip()
        if not kind:
            raise ValueError("HITL interrupt payload is missing kind.")
        if kind not in {"subquestion_review", "query_expansion_review"}:
            raise ValueError(f"Unsupported HITL review kind '{kind}'.")
        if not checkpoint_id:
            raise ValueError("HITL interrupt payload is missing checkpoint_id.")
        return cls(
            kind=kind,
            stage=str(payload.get("stage") or ""),
            checkpoint_id=checkpoint_id,
            items=cls._normalize_items(kind=kind, payload=payload),
        )

    @staticmethod
    def _normalize_items(*, kind: str, payload: Mapping[str, Any]) -> list[HitlReviewItem]:
        raw_items_key = "subquestions" if kind == "subquestion_review" else "expansions"
        raw_items = payload.get(raw_items_key)
        if not isinstance(raw_items, Sequence) or isinstance(raw_items, (str, bytes, bytearray)):
            return []
        item_id_key = "subquestion_id" if kind == "subquestion_review" else "expansion_id"
        text_key = "sub_question" if kind == "subquestion_review" else "query"
        items: list[HitlReviewItem] = []
        for index, raw_item in enumerate(raw_items):
            if not isinstance(raw_item, Mapping):
                continue
            raw_index = raw_item.get("index")
            items.append(
                HitlReviewItem(
                    item_id=str(raw_item.get(item_id_key) or f"item-{index + 1}").strip(),
                    text=str(raw_item.get(text_key) or "").strip(),
                    index=raw_index if isinstance(raw_index, int) and raw_index >= 0 else index,
                )
            )
        return items

    def approve_all(self) -> "HitlResumeRequest":
        return HitlResumeRequest(
            checkpoint_id=self.checkpoint_id,
            review_kind=self.kind,
            decisions=[item.approve() for item in self.items],
        )

    def with_decisions(self, *decisions: HitlResumeDecision) -> "HitlResumeRequest":
        return HitlResumeRequest(
            checkpoint_id=self.checkpoint_id,
            review_kind=self.kind,
            decisions=list(decisions),
        )


class HitlResumeRequest(BaseModel):
    checkpoint_id: str = Field(min_length=1)
    review_kind: Literal["subquestion_review", "query_expansion_review"]
    decisions: list[HitlResumeDecision] = Field(min_length=1)

    @classmethod
    def approve(cls, review: HitlReview, item_id: str) -> "HitlResumeRequest":
        return review.with_decisions(HitlResumeDecision(item_id=item_id, action="approve"))

    @classmethod
    def edit(cls, review: HitlReview, item_id: str, replacement_text: str) -> "HitlResumeRequest":
        return review.with_decisions(
            HitlResumeDecision(item_id=item_id, action="edit", replacement_text=replacement_text)
        )

    @classmethod
    def reject(cls, review: HitlReview, item_id: str) -> "HitlResumeRequest":
        return review.with_decisions(HitlResumeDecision(item_id=item_id, action="reject"))


class RuntimeAgentRunResult(BaseModel):
    status: Literal["completed", "paused"]
    thread_id: str = ""
    checkpoint_id: str | None = None
    review: HitlReview | None = Field(
        default=None,
        validation_alias=AliasChoices("review", "interrupt_payload"),
    )
    response: RuntimeAgentRunResponse | None = None

    @field_validator("review", mode="before")
    @classmethod
    def validate_review(cls, value: Any) -> HitlReview | Any | None:
        if value is None or isinstance(value, HitlReview):
            return value
        if isinstance(value, Mapping):
            return HitlReview.from_interrupt_payload(value)
        return value


class AgentRunStageMetadata(BaseModel):
    stage: str
    status: str
    sub_question: str = ""
    lane_index: int = 0
    lane_total: int = 0
    emitted_at: float | None = None


class RuntimeAgentRunAsyncStartResponse(BaseModel):
    job_id: str
    run_id: str
    thread_id: str = ""
    status: str


class RuntimeAgentRunAsyncStatusResponse(BaseModel):
    job_id: str
    run_id: str = ""
    thread_id: str = ""
    status: str
    message: str = ""
    stage: str = ""
    stages: list[AgentRunStageMetadata] = Field(default_factory=list)
    decomposition_sub_questions: list[str] = Field(default_factory=list)
    sub_question_artifacts: list["SubQuestionArtifacts"] = Field(default_factory=list)
    sub_qa: list[SubQuestionAnswer] = Field(default_factory=list)
    sub_answers: list[SubQuestionAnswer] = Field(default_factory=list)
    output: str = ""
    result: RuntimeAgentRunResponse | None = None
    error: str | None = None
    cancel_requested: bool = False
    review: HitlReview | None = Field(
        default=None,
        validation_alias=AliasChoices("review", "interrupt_payload"),
    )
    checkpoint_id: str | None = None
    started_at: float | None = None
    finished_at: float | None = None
    elapsed_ms: int | None = None

    @field_validator("review", mode="before")
    @classmethod
    def validate_review(cls, value: Any) -> HitlReview | Any | None:
        if value is None or isinstance(value, HitlReview):
            return value
        if isinstance(value, Mapping):
            return HitlReview.from_interrupt_payload(value)
        return value


class RuntimeAgentRunAsyncCancelResponse(BaseModel):
    status: Literal["success"]
    message: str


class RuntimeSubquestionDecision(BaseModel):
    subquestion_id: str = Field(min_length=1)
    action: Literal["approve", "edit", "deny", "skip"]
    edited_text: str | None = None

    @model_validator(mode="after")
    def validate_edit_payload(self) -> "RuntimeSubquestionDecision":
        if self.action == "edit":
            if self.edited_text is None or not self.edited_text.strip():
                raise ValueError("edited_text is required when action='edit'.")
        elif self.edited_text is not None:
            self.edited_text = self.edited_text.strip() or None
        return self


class RuntimeSubquestionResumeEnvelope(BaseModel):
    checkpoint_id: str = Field(min_length=1)
    decisions: list[RuntimeSubquestionDecision] = Field(min_length=1)


class RuntimeQueryExpansionDecision(BaseModel):
    expansion_id: str = Field(min_length=1)
    action: Literal["approve", "edit", "deny", "skip"]
    edited_query: str | None = None

    @model_validator(mode="after")
    def validate_edit_payload(self) -> "RuntimeQueryExpansionDecision":
        if self.action == "edit":
            if self.edited_query is None or not self.edited_query.strip():
                raise ValueError("edited_query is required when action='edit'.")
        elif self.edited_query is not None:
            self.edited_query = self.edited_query.strip() or None
        return self


class RuntimeQueryExpansionResumeEnvelope(BaseModel):
    checkpoint_id: str = Field(min_length=1)
    decisions: list[RuntimeQueryExpansionDecision] = Field(min_length=1)


class RuntimeAgentRunResumeRequest(BaseModel):
    resume: Union[
        bool,
        HitlResumeRequest,
        dict[str, Any],
        RuntimeSubquestionResumeEnvelope,
        RuntimeQueryExpansionResumeEnvelope,
    ] = True

    @field_validator("resume", mode="before")
    @classmethod
    def validate_resume(cls, value: Any) -> Any:
        if isinstance(value, bool):
            return value
        if isinstance(value, HitlResumeRequest):
            return value
        if not isinstance(value, dict):
            raise ValueError(
                "resume must be a boolean, SDK HITL resume request, legacy object payload, or typed decision envelope."
            )
        if "review_kind" in value and "checkpoint_id" in value and "decisions" in value:
            return HitlResumeRequest.model_validate(value)
        if "checkpoint_id" in value or "decisions" in value:
            decisions = value.get("decisions")
            if isinstance(decisions, list):
                for decision in decisions:
                    if isinstance(decision, dict) and (
                        "expansion_id" in decision or "edited_query" in decision
                    ):
                        return RuntimeQueryExpansionResumeEnvelope.model_validate(value)
            return RuntimeSubquestionResumeEnvelope.model_validate(value)
        return value


class RuntimeAgentInfo(BaseModel):
    agent_name: str


class GraphRunMetadata(BaseModel):
    run_id: str
    thread_id: str = ""
    trace_id: str = ""
    correlation_id: str = ""


class CitationSourceRow(BaseModel):
    citation_index: int = Field(ge=1)
    rank: int = Field(ge=1)
    title: str = ""
    source: str = ""
    content: str = ""
    document_id: str = ""
    score: float | None = None


class SubQuestionArtifacts(BaseModel):
    sub_question: str
    expanded_queries: list[str] = Field(default_factory=list)
    retrieved_docs: list[CitationSourceRow] = Field(default_factory=list)
    retrieval_provenance: list[dict[str, Any]] = Field(default_factory=list)
    reranked_docs: list[CitationSourceRow] = Field(default_factory=list)
    sub_answer: str = ""
    citation_rows_by_index: dict[int, CitationSourceRow] = Field(default_factory=dict)


class AgentGraphState(BaseModel):
    main_question: str
    decomposition_sub_questions: list[str] = Field(default_factory=list)
    sub_question_artifacts: list[SubQuestionArtifacts] = Field(default_factory=list)
    final_answer: str = ""
    citation_rows_by_index: dict[int, CitationSourceRow] = Field(default_factory=dict)
    run_metadata: GraphRunMetadata
    sub_qa: list[SubQuestionAnswer] = Field(default_factory=list)
    output: str = ""
    stage_snapshots: list["GraphStageSnapshot"] = Field(default_factory=list)


class DecomposeNodeInput(BaseModel):
    main_question: str
    run_metadata: GraphRunMetadata
    initial_search_context: list[dict[str, Any]] = Field(default_factory=list)


class DecomposeNodeOutput(BaseModel):
    decomposition_sub_questions: list[str] = Field(default_factory=list)


class ExpandNodeInput(BaseModel):
    main_question: str
    sub_question: str
    run_metadata: GraphRunMetadata


class ExpandNodeOutput(BaseModel):
    expanded_queries: list[str] = Field(default_factory=list)


class SearchNodeInput(BaseModel):
    sub_question: str
    expanded_queries: list[str] = Field(default_factory=list)
    run_metadata: GraphRunMetadata


class SearchNodeOutput(BaseModel):
    retrieved_docs: list[CitationSourceRow] = Field(default_factory=list)
    retrieval_provenance: list[dict[str, Any]] = Field(default_factory=list)
    citation_rows_by_index: dict[int, CitationSourceRow] = Field(default_factory=dict)


class RerankNodeInput(BaseModel):
    sub_question: str
    expanded_query: str = ""
    retrieved_docs: list[CitationSourceRow] = Field(default_factory=list)
    run_metadata: GraphRunMetadata


class RerankNodeOutput(BaseModel):
    reranked_docs: list[CitationSourceRow] = Field(default_factory=list)
    citation_rows_by_index: dict[int, CitationSourceRow] = Field(default_factory=dict)


class AnswerSubquestionNodeInput(BaseModel):
    sub_question: str
    reranked_docs: list[CitationSourceRow] = Field(default_factory=list)
    citation_rows_by_index: dict[int, CitationSourceRow] = Field(default_factory=dict)
    run_metadata: GraphRunMetadata


class AnswerSubquestionNodeOutput(BaseModel):
    sub_answer: str = ""
    citation_indices_used: list[int] = Field(default_factory=list)
    answerable: bool = False
    verification_reason: str = ""
    citation_rows_by_index: dict[int, CitationSourceRow] = Field(default_factory=dict)


class SynthesizeFinalNodeInput(BaseModel):
    main_question: str
    sub_qa: list[SubQuestionAnswer] = Field(default_factory=list)
    sub_question_artifacts: list[SubQuestionArtifacts] = Field(default_factory=list)
    run_metadata: GraphRunMetadata


class SynthesizeFinalNodeOutput(BaseModel):
    final_answer: str = ""


class GraphStageSnapshot(BaseModel):
    stage: str
    status: str = "completed"
    sub_question: str = ""
    lane_index: int = 0
    lane_total: int = 0
    decomposition_sub_questions: list[str] = Field(default_factory=list)
    sub_qa: list[SubQuestionAnswer] = Field(default_factory=list)
    sub_question_artifacts: list[SubQuestionArtifacts] = Field(default_factory=list)
    output: str = ""
