"""SelfQAMixin — autonomous answer verification and revision."""
from __future__ import annotations

from datetime import datetime, UTC

from logos_adk.context import get_run_context
from logos_adk.errors import BudgetExhaustedError
from logos_adk.self_qa import (
    SelfQAAssertion, SelfQAAttempt, SelfQACheck, SelfQAPlanModel,
    SelfQAReport, SelfQAVerdict, SelfQAVerdictModel,
)
from logos_adk.types import Message, Response, content_text


class SelfQAMixin:
    """Mixin providing self-QA verification methods for Agent."""

    def _self_qa_enabled(self) -> bool:
        return bool(getattr(self._config, "self_qa", None) and self._config.self_qa.enabled)

    def _fallback_self_qa_assertions(self, prompt_text: str, answer_text: str) -> list[SelfQAAssertion]:
        assertions = [
            SelfQAAssertion(
                assertion="The answer directly addresses the user's request without ignoring the main task.",
                rationale="Ensures the result is on-topic.",
            ),
            SelfQAAssertion(
                assertion="The answer is specific enough to be actionable and internally consistent.",
                rationale="Prevents vague or contradictory output.",
            ),
        ]
        if prompt_text.strip() and answer_text.strip():
            assertions.append(
                SelfQAAssertion(
                    assertion="The answer reflects concrete details from the request rather than a generic template.",
                    rationale="Guards against boilerplate output.",
                )
            )
        return assertions

    async def _self_qa_stage_generate(
        self,
        *,
        prompt: str,
        trace_id: str,
        session_id: str | None,
        request_id: str | None,
    ) -> Response:
        messages = [Message(role="system", content="Return exactly what was requested."), Message(role="user", content=prompt)]
        _, response = await self._provider_generate_with_reliability(
            messages,
            tools=None,
            output_schema=None,
            trace_id=trace_id,
            session_id=session_id,
            request_id=request_id,
        )
        failure = self._strict_budget_failure_response(
            run_context=get_run_context(),
            stage="self_qa",
            usage=response.usage,
            raw=response.raw,
        )
        if failure is not None:
            raise BudgetExhaustedError(failure.content)
        return response

    async def _generate_self_qa_assertions(
        self,
        *,
        prompt_text: str,
        answer_text: str,
        trace_id: str,
        session_id: str | None,
        request_id: str | None,
    ) -> list[SelfQAAssertion]:
        config = self._config.self_qa
        stage_started_at = datetime.now(UTC)
        response = await self._self_qa_stage_generate(
            prompt=(
                "You are the answering agent's internal QA planner.\n"
                f"Write {config.min_assertions}-{config.max_assertions} concise, objective assertions that the candidate answer must satisfy.\n"
                'Return strict JSON with key "assertions", where each item has "assertion" and optional "rationale".\n'
                f"Original task:\n{prompt_text}\n\nCandidate answer:\n{answer_text}"
            ),
            trace_id=trace_id,
            session_id=session_id,
            request_id=request_id,
        )
        parsed = self._parse_structured_response(response, SelfQAPlanModel).parsed
        assertions: list[SelfQAAssertion] = []
        if parsed is not None:
            assertions = [
                SelfQAAssertion(assertion=item.assertion, rationale=item.rationale)
                for item in parsed.assertions[: config.max_assertions]
            ]
        if len(assertions) < config.min_assertions:
            for item in self._fallback_self_qa_assertions(prompt_text, answer_text):
                if len(assertions) >= config.max_assertions:
                    break
                if any(existing.assertion == item.assertion for existing in assertions):
                    continue
                assertions.append(item)
        await self._emit_span(
            name="SelfQA.assertions",
            trace_id=trace_id,
            parent_span_id=None,
            started_at=stage_started_at,
            session_id=session_id,
            attributes={"assertions": len(assertions)},
        )
        return assertions[: config.max_assertions]

    async def _verify_self_qa_assertions(
        self,
        *,
        prompt_text: str,
        answer_text: str,
        assertions: list[SelfQAAssertion],
        trace_id: str,
        session_id: str | None,
        request_id: str | None,
    ) -> SelfQAVerdict:
        stage_started_at = datetime.now(UTC)
        assertion_lines = "\n".join(f"- {item.assertion}" for item in assertions)
        response = await self._self_qa_stage_generate(
            prompt=(
                "You are a strict tester agent. Verify the candidate answer against every assertion.\n"
                'Return strict JSON with keys "passed", "summary", and "checks".\n'
                'Each check must include "assertion", "passed", optional "evidence", and optional "failure_reason".\n'
                "Mark passed=false if any assertion is not clearly satisfied.\n"
                f"Original task:\n{prompt_text}\n\nAssertions:\n{assertion_lines}\n\nCandidate answer:\n{answer_text}"
            ),
            trace_id=trace_id,
            session_id=session_id,
            request_id=request_id,
        )
        parsed = self._parse_structured_response(response, SelfQAVerdictModel).parsed
        if parsed is None:
            verdict = SelfQAVerdict(
                passed=False,
                summary="Self-QA tester returned an invalid verification payload.",
                checks=[
                    SelfQACheck(
                        assertion=item.assertion,
                        passed=False,
                        failure_reason="Tester output was not valid structured JSON.",
                    )
                    for item in assertions
                ],
            )
        else:
            verdict = SelfQAVerdict(
                passed=parsed.passed,
                summary=parsed.summary,
                checks=[
                    SelfQACheck(
                        assertion=item.assertion,
                        passed=item.passed,
                        evidence=item.evidence,
                        failure_reason=item.failure_reason,
                    )
                    for item in parsed.checks
                ],
            )
        await self._emit_span(
            name="SelfQA.verify",
            trace_id=trace_id,
            parent_span_id=None,
            started_at=stage_started_at,
            session_id=session_id,
            status="ok" if verdict.passed else "error",
            attributes={"checks": len(verdict.checks), "passed": verdict.passed},
        )
        return verdict

    async def _revise_after_self_qa_failure(
        self,
        *,
        prompt_text: str,
        answer_text: str,
        assertions: list[SelfQAAssertion],
        verdict: SelfQAVerdict,
        trace_id: str,
        session_id: str | None,
        request_id: str | None,
    ) -> str:
        stage_started_at = datetime.now(UTC)
        failure_lines = "\n".join(
            f"- {item.assertion}: {item.failure_reason or item.evidence or 'failed'}"
            for item in verdict.checks
            if not item.passed
        )
        response = await self._self_qa_stage_generate(
            prompt=(
                "Revise the candidate answer so that every failed assertion becomes true.\n"
                "Preserve correct parts, remove contradictions, and return only the improved final answer.\n"
                "If the answer should remain structured JSON, keep that format valid.\n"
                f"Original task:\n{prompt_text}\n\nAssertions:\n"
                + "\n".join(f"- {item.assertion}" for item in assertions)
                + f"\n\nCurrent answer:\n{answer_text}\n\nFailed checks:\n{failure_lines}"
            ),
            trace_id=trace_id,
            session_id=session_id,
            request_id=request_id,
        )
        revised_text = content_text(response.content, response.content_parts).strip() or answer_text
        await self._emit_span(
            name="SelfQA.revise",
            trace_id=trace_id,
            parent_span_id=None,
            started_at=stage_started_at,
            session_id=session_id,
            attributes={"revised_length": len(revised_text)},
        )
        return revised_text

    async def _run_self_qa(
        self,
        *,
        prompt_text: str,
        response: Response,
        trace_id: str,
        session_id: str | None,
        request_id: str | None,
    ) -> Response:
        if not self._self_qa_enabled() or response.status != "completed":
            return response

        config = self._config.self_qa
        run_context = get_run_context()
        effective_max_retries = self._effective_self_qa_max_retries(run_context)
        answer_text = content_text(response.content, response.content_parts).strip()
        if not answer_text:
            return response

        attempts: list[SelfQAAttempt] = []
        candidate_text = answer_text
        final_verdict: SelfQAVerdict | None = None

        try:
            for attempt_number in range(1, effective_max_retries + 2):
                assertions = await self._generate_self_qa_assertions(
                    prompt_text=prompt_text,
                    answer_text=candidate_text,
                    trace_id=trace_id,
                    session_id=session_id,
                    request_id=request_id,
                )
                verdict = await self._verify_self_qa_assertions(
                    prompt_text=prompt_text,
                    answer_text=candidate_text,
                    assertions=assertions,
                    trace_id=trace_id,
                    session_id=session_id,
                    request_id=request_id,
                )
                attempt = SelfQAAttempt(
                    attempt=attempt_number,
                    candidate=candidate_text,
                    assertions=assertions,
                    verdict=verdict,
                )
                attempts.append(attempt)
                final_verdict = verdict
                if verdict.passed:
                    raw = dict(response.raw or {})
                    raw["self_qa"] = SelfQAReport(
                        passed=True,
                        attempts=attempts,
                        final_summary=verdict.summary,
                        max_retries=effective_max_retries,
                    ).to_dict()
                    return self._attach_budget_metadata(Response(
                        content=candidate_text,
                        content_parts=response.content_parts if candidate_text == answer_text else [],
                        tool_calls=response.tool_calls,
                        usage=response.usage,
                        raw=raw,
                        reasoning=response.reasoning,
                        status="completed",
                    ), run_context=run_context, self_qa_max_retries=effective_max_retries)
                if attempt_number <= effective_max_retries:
                    revised_text = await self._revise_after_self_qa_failure(
                        prompt_text=prompt_text,
                        answer_text=candidate_text,
                        assertions=assertions,
                        verdict=verdict,
                        trace_id=trace_id,
                        session_id=session_id,
                        request_id=request_id,
                    )
                    attempts[-1].revised_candidate = revised_text
                    candidate_text = revised_text
        except BudgetExhaustedError:
            failure = self._strict_budget_failure_response(
                run_context=run_context,
                stage="self_qa",
                usage=response.usage,
                raw=response.raw,
            )
            if failure is not None:
                return failure

        report = SelfQAReport(
            passed=False,
            attempts=attempts,
            final_summary=(
                final_verdict.summary
                if final_verdict is not None
                else "Self-QA failed without a tester verdict."
            ),
            max_retries=effective_max_retries,
        )
        raw = dict(response.raw or {})
        raw["self_qa"] = report.to_dict()
        raw["self_qa_last_candidate"] = candidate_text
        if not config.fail_closed:
            return self._attach_budget_metadata(Response(
                content=candidate_text,
                content_parts=[],
                tool_calls=response.tool_calls,
                usage=response.usage,
                raw=raw,
                reasoning=response.reasoning,
                status="completed",
            ), run_context=run_context, self_qa_max_retries=effective_max_retries)
        return self._attach_budget_metadata(Response(
            content=f"Self-QA failed after {len(attempts)} attempts. {report.final_summary}",
            content_parts=[],
            tool_calls=[],
            usage=response.usage,
            raw=raw,
            reasoning=response.reasoning,
            status="failed",
        ), run_context=run_context, self_qa_max_retries=effective_max_retries)
