from __future__ import annotations

import ast

from .shared import *

class LLMStructuredExtractObjectMixin:
    @staticmethod
    def _structured_candidate_preference(field: str) -> int:
        normalized = str(field or "").strip().lower()
        if normalized == "message.content":
            return 4
        if normalized == "response.output_text":
            return 3
        if normalized == "choice.text":
            return 2
        if normalized == "tool_calls.function.arguments":
            return 1
        return 0

    @classmethod
    def _structured_payload_richness(cls, payload: Any) -> int:
        if isinstance(payload, dict):
            score = 0
            for key, value in payload.items():
                score += 1 if str(key or "").strip() else 0
                score += cls._structured_payload_richness(value)
            return score
        if isinstance(payload, list):
            return len(payload) + sum(cls._structured_payload_richness(item) for item in payload)
        if isinstance(payload, str):
            return len(str(payload).strip())
        if isinstance(payload, (int, float)) and not isinstance(payload, bool):
            return 1
        if isinstance(payload, bool):
            return 1
        return 0

    def _extract_structured_object_from_response(
        self,
        *,
        function_name: str,
        response: Any,
        fallback_text: str,
        provider: str,
        model: str,
    ) -> StructuredExtractionResult:
        message = {}
        choice: dict[str, Any] = {}
        if isinstance(response, dict):
            choices = response.get("choices") if isinstance(response.get("choices"), list) else []
            choice = choices[0] if choices and isinstance(choices[0], dict) else {}
            message = choice.get("message") if isinstance(choice.get("message"), dict) else {}

        last_nonempty_field = ""
        last_nonempty_text = ""
        last_invalid_result: Optional[StructuredExtractionResult] = None
        valid_results: list[StructuredExtractionResult] = []
        for field, raw_value in self._iter_structured_candidates(
            message=message,
            choice=choice,
            response=response,
            include_reasoning=False,
        ):
            text = self._sanitize_visible_answer(self._extract_text_from_value(raw_value))
            normalized = self._strip_code_fence(text)
            if normalized:
                last_nonempty_field = field
                last_nonempty_text = normalized[:4000]
            if not self._looks_like_json_container(normalized):
                continue
            candidate = self._extract_json_candidate(normalized)
            if not candidate:
                continue
            try:
                parsed = json.loads(candidate)
            except Exception:
                parsed = self._parse_loose_structured_candidate(candidate)
                if parsed is None:
                    last_invalid_result = StructuredExtractionResult(
                        payload=None,
                        source_field=field,
                        provider=provider,
                        model=model,
                        error_code="structured_json_invalid",
                        raw_text=candidate[:4000],
                        validation="invalid",
                    )
                    continue
            parsed = self._unwrap_structured_payload_once(parsed)
            parsed = self._coerce_structured_payload(
                function_name=function_name,
                payload=parsed,
            )
            validation_error = self._validate_structured_payload(function_name=function_name, payload=parsed)
            if validation_error is not None:
                last_invalid_result = StructuredExtractionResult(
                    payload=None,
                    source_field=field,
                    provider=provider,
                    model=model,
                    error_code=validation_error,
                    raw_text=candidate[:4000],
                    validation="invalid",
                )
                continue
            valid_results.append(
                StructuredExtractionResult(
                    payload=parsed if isinstance(parsed, dict) else None,
                    source_field=field,
                    provider=provider,
                    model=model,
                    raw_text=candidate[:4000],
                    validation="ok",
                )
            )

        fallback_normalized = self._strip_code_fence(fallback_text)
        fallback_candidate = (
            self._extract_json_candidate(fallback_normalized)
            if self._looks_like_json_container(fallback_normalized)
            else ""
        )
        if fallback_candidate:
            parsed: Any = None
            try:
                parsed = json.loads(fallback_candidate)
            except Exception:
                parsed = self._parse_loose_structured_candidate(fallback_candidate)
            if parsed is not None:
                parsed = self._unwrap_structured_payload_once(parsed)
                parsed = self._coerce_structured_payload(
                    function_name=function_name,
                    payload=parsed,
                )
                validation_error = self._validate_structured_payload(function_name=function_name, payload=parsed)
                if validation_error is None and isinstance(parsed, dict):
                    valid_results.append(
                        StructuredExtractionResult(
                            payload=parsed,
                            source_field="generated_text",
                            provider=provider,
                            model=model,
                            raw_text=fallback_candidate[:4000],
                            validation="ok",
                        )
                    )
                elif validation_error is not None:
                    last_invalid_result = StructuredExtractionResult(
                        payload=None,
                        source_field="generated_text",
                        provider=provider,
                        model=model,
                        error_code=validation_error or "structured_schema_invalid",
                        raw_text=fallback_candidate[:4000],
                        validation="invalid",
                    )
            else:
                last_invalid_result = StructuredExtractionResult(
                    payload=None,
                    source_field="generated_text",
                    provider=provider,
                    model=model,
                    error_code="structured_json_invalid",
                    raw_text=fallback_candidate[:4000],
                    validation="invalid",
                )

        if valid_results:
            return max(
                valid_results,
                key=lambda item: (
                    self._structured_payload_richness(item.payload),
                    self._structured_candidate_preference(item.source_field or ""),
                ),
            )

        if last_invalid_result is not None:
            return last_invalid_result

        if not last_nonempty_field and str(fallback_text or "").strip():
            last_nonempty_field = "generated_text"
            last_nonempty_text = str(fallback_text or "").strip()[:4000]
        error_code = "structured_json_not_found" if last_nonempty_field else "structured_empty"
        return StructuredExtractionResult(
            payload=None,
            source_field=last_nonempty_field,
            provider=provider,
            model=model,
            error_code=error_code,
            raw_text=last_nonempty_text,
            validation="invalid",
        )

    @staticmethod
    def _parse_loose_structured_candidate(candidate: str) -> Any:
        text = str(candidate or "").strip()
        if not text:
            return None
        try:
            parsed = ast.literal_eval(text)
        except Exception:
            return None
        if isinstance(parsed, str):
            try:
                return json.loads(parsed)
            except Exception:
                return None
        if isinstance(parsed, (dict, list)):
            return parsed
        return None

    @staticmethod
    def _classify_structured_exception(exc: Exception) -> str:
        if isinstance(exc, (asyncio.TimeoutError, httpx.TimeoutException)):
            return "structured_timeout"
        if isinstance(exc, (httpx.HTTPError, httpx.RequestError)):
            return "structured_http_error"
        lowered = str(exc or "").strip().lower()
        if "timeout" in lowered:
            return "structured_timeout"
        if "status" in lowered or "http" in lowered:
            return "structured_http_error"
        return "structured_empty"