"""LLM norm extraction: chunk, call model, merge raw norm dicts (no DB / NormCreate)."""
from __future__ import annotations

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

from norm_extractor.chunking import (
    dedupe_norm_dicts_by_title,
    distribute_max_norms_across_chunks,
    split_text_for_extraction,
)
from norm_extractor.exceptions import BadExtractionJsonError, LlmExtractionError
from norm_extractor.parsing import parse_llm_norms_json, truncate_raw_llm
from norm_extractor.prompts import load_zh_extraction_system_prompt
from norm_extractor.settings import get_extractor_settings

from llm.runtime import get_default_llm_provider

_JSON_RETRY_ATTEMPTS = 3


@dataclass(frozen=True)
class RawExtractionOutcome:
    """Immutable result of LLM extraction before NormCreate / DB."""

    norms_raw: List[Dict[str, Any]]
    chunk_results: List[Dict[str, Any]]
    chunk_count: int
    max_content_chars: int
    overlap_chars: int


async def _llm_extract_one_chunk(
    *,
    system_prompt: str,
    default_lang: str,
    field_id: Optional[UUID],
    created_by: Optional[str],
    chunk_max_norms: int,
    default_authority_level: Optional[str],
    default_formality_level: Optional[str],
    doc_id: str,
    doc_title: str,
    content: str,
    chunk_index: int,
    chunk_total: int,
    settings: Any,
) -> str:
    user_payload: Dict[str, Any] = {
        "task_mode": "extract_or_generate",
        "default_lang": default_lang,
        "target_field_id": str(field_id) if field_id else None,
        "default_created_by": created_by,
        "default_authority_level": default_authority_level,
        "default_formality_level": default_formality_level,
        "source_document": {
            "doc_id": doc_id,
            "title": doc_title,
            "version": "1",
            "language": default_lang,
            "content": content,
            "chunk_index": chunk_index,
            "chunk_total": chunk_total,
        },
        "extraction_scope": {
            "max_norms": chunk_max_norms,
            "include_permissions": True,
            "include_prohibitions": True,
            "include_obligations": True,
            "chunk_index": chunk_index,
            "chunk_total": chunk_total,
        },
    }
    user_message = (
        "请根据以下输入萃取并生成 norm 数据，严格遵循系统提示词中的字段合同与规则。\n\n"
        "输入：\n\n"
        f"{json.dumps(user_payload, ensure_ascii=False, indent=2)}\n\n"
        '输出要求：仅输出 JSON；顶层为 { "norms": [...] }；不要输出任何额外说明。'
    )
    llm = get_default_llm_provider()
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": user_message},
    ]
    chat_kwargs: Dict[str, Any] = {}
    if settings.extraction_llm_max_output_tokens > 0:
        chat_kwargs["max_tokens"] = settings.extraction_llm_max_output_tokens
    if settings.extraction_llm_temperature is not None:
        chat_kwargs["temperature"] = settings.extraction_llm_temperature
    try:
        if chat_kwargs:
            try:
                return await llm.chat(messages, **chat_kwargs)
            except TypeError:
                return await llm.chat(messages)
        return await llm.chat(messages)
    except Exception as e:
        raise LlmExtractionError(str(e)) from e


async def run_raw_extraction(
    text: str,
    *,
    default_lang: str,
    field_id: Optional[UUID],
    created_by: Optional[str],
    max_norms: int,
    default_authority_level: Optional[str],
    default_formality_level: Optional[str],
    doc_id: str,
    doc_title: str,
) -> RawExtractionOutcome:
    """
    Run chunked LLM extraction; return raw norm dicts (not validated against NormCreate).

    Raises:
        LlmExtractionError: provider failure.
        BadExtractionJsonError: invalid JSON or missing ``norms`` array for a chunk.
    """
    settings = get_extractor_settings()
    system_prompt = load_zh_extraction_system_prompt()
    max_chars = settings.extraction_llm_max_content_chars
    overlap = settings.extraction_chunk_overlap_chars
    chunks = split_text_for_extraction(text or "", max_chars, overlap)
    per_chunk_limits = distribute_max_norms_across_chunks(max_norms, len(chunks))

    merged_norms_raw: List[Dict[str, Any]] = []
    chunk_results: List[Dict[str, Any]] = []
    raw_trunc_limit = settings.extraction_artifact_raw_llm_max_chars
    chunk_total = len(chunks)

    for i, chunk in enumerate(chunks):
        budget = per_chunk_limits[i]
        raw_reply = ""
        norms_part: List[Any] = []
        last_json_err: Exception | None = None
        requested_max = budget

        for attempt in range(_JSON_RETRY_ATTEMPTS):
            chunk_max = max(1, budget // (2**attempt))
            requested_max = chunk_max
            raw_reply = await _llm_extract_one_chunk(
                system_prompt=system_prompt,
                default_lang=default_lang,
                field_id=field_id,
                created_by=created_by,
                chunk_max_norms=chunk_max,
                default_authority_level=default_authority_level,
                default_formality_level=default_formality_level,
                doc_id=doc_id,
                doc_title=doc_title,
                content=chunk,
                chunk_index=i,
                chunk_total=chunk_total,
                settings=settings,
            )
            try:
                parsed = parse_llm_norms_json(raw_reply)
            except (json.JSONDecodeError, ValueError) as e:
                last_json_err = e
                if attempt + 1 >= _JSON_RETRY_ATTEMPTS:
                    raise BadExtractionJsonError(
                        f"LLM output is not valid JSON (chunk {i + 1}/{chunk_total}) "
                        f"after {_JSON_RETRY_ATTEMPTS} attempts "
                        f"(max_norms per try down to {chunk_max}): {e}"
                    ) from e
                continue

            np = parsed.get("norms")
            if not isinstance(np, list):
                last_json_err = ValueError("top-level 'norms' is not an array")
                if attempt + 1 >= _JSON_RETRY_ATTEMPTS:
                    raise BadExtractionJsonError(
                        f"LLM JSON must contain a 'norms' array (chunk {i + 1}/{chunk_total}) "
                        f"after {_JSON_RETRY_ATTEMPTS} attempts."
                    ) from last_json_err
                continue

            norms_part = np
            break

        cap = budget
        slice_norms = [
            x for x in norms_part[: cap + 10] if isinstance(x, dict)
        ][:cap]
        merged_norms_raw.extend(slice_norms)

        tr, was_trunc = truncate_raw_llm(raw_reply, raw_trunc_limit)
        chunk_results.append(
            {
                "index": i,
                "content_chars": len(chunk),
                "raw_llm": tr,
                "raw_llm_truncated": was_trunc,
                "norm_count": len(slice_norms),
                "max_norms_requested": requested_max,
                "json_parse_attempts": attempt + 1,
            }
        )

    norms_raw = dedupe_norm_dicts_by_title(merged_norms_raw)
    norms_raw = [x for x in norms_raw if isinstance(x, dict)][:max_norms]

    return RawExtractionOutcome(
        norms_raw=norms_raw,
        chunk_results=chunk_results,
        chunk_count=chunk_total,
        max_content_chars=max_chars,
        overlap_chars=overlap,
    )
