"""
xmldev.llm
~~~~~~~~~~
LLM fallback client for XML repair.

Reads configuration from xmldev.config.env (KEY=VALUE).
LLM calls are only made when XMDEV_ALLOW_LLM=true is set.
PII redaction is ON by default.
"""
from __future__ import annotations

import json
import logging
import os
import re
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

from openai import OpenAI, OpenAIError

from .schema import Schema

logger = logging.getLogger("xmldev.llm")

# ---------------------------------------------------------------------------
# Config
# ---------------------------------------------------------------------------

@dataclass
class LLMConfig:
    base_url: str = ""
    api_key: str = ""
    model_id: str = ""
    timeout: float = 60.0
    max_tokens: int = 2048
    temperature: float = 0.0
    allow_llm: bool = False
    redact: bool = True
    redaction_tags: list[str] = field(default_factory=list)
    redaction_attrs: list[str] = field(default_factory=list)

    @classmethod
    def load(cls, config_path: str | Path | None) -> "LLMConfig | None":
        """
        Load config from an env file. Returns None if:
        - config_path is None
        - file does not exist
        - XMDEV_ALLOW_LLM != true
        """
        if config_path is None:
            return None

        path = Path(config_path)
        if not path.exists():
            logger.debug("LLM config not found at %s; LLM disabled.", path)
            return None

        kv: dict[str, str] = {}
        for line in path.read_text(encoding="utf-8").splitlines():
            line = line.strip()
            if not line or line.startswith("#"):
                continue
            if "=" in line:
                key, _, value = line.partition("=")
                kv[key.strip()] = value.strip()

        if kv.get("XMDEV_ALLOW_LLM", "false").lower() != "true":
            logger.debug("XMDEV_ALLOW_LLM not true; LLM disabled.")
            return None

        return cls(
            base_url=kv.get("LLM_BASE_URL", "https://api.openai.com/v1"),
            api_key=kv.get("LLM_API_KEY", ""),
            model_id=kv.get("LLM_MODEL_ID", "gpt-4o-mini"),
            timeout=float(kv.get("LLM_REQUEST_TIMEOUT", "60")),
            max_tokens=int(kv.get("LLM_MAX_TOKENS", "2048")),
            temperature=float(kv.get("LLM_TEMPERATURE", "0.0")),
            allow_llm=True,
            redact=kv.get("LLM_REDACT", "true").lower() != "false",
            redaction_tags=kv.get("LLM_REDACT_TAGS", "").split(",") if kv.get("LLM_REDACT_TAGS") else [],
            redaction_attrs=kv.get("LLM_REDACT_ATTRS", "").split(",") if kv.get("LLM_REDACT_ATTRS") else [],
        )


# ---------------------------------------------------------------------------
# PII Redactor
# ---------------------------------------------------------------------------

_DEFAULT_REDACT_TAGS = ["ssn", "password", "secret", "credit_card", "card_number", "token", "api_key"]
_DEFAULT_REDACT_ATTRS = ["password", "secret", "token", "api_key", "ssn"]


class PIIRedactor:
    """Replace sensitive tag text content and attribute values with [REDACTED]."""

    def __init__(self, tags: list[str] | None = None, attrs: list[str] | None = None) -> None:
        self.tags = set((tags or []) + _DEFAULT_REDACT_TAGS)
        self.attrs = set((attrs or []) + _DEFAULT_REDACT_ATTRS)

    def redact(self, xml_str: str) -> str:
        """Return a redacted copy of the XML string."""
        # Redact tag content
        for tag in self.tags:
            xml_str = re.sub(
                rf"(<{re.escape(tag)}[^>]*>)[^<]*(</\s*{re.escape(tag)}\s*>)",
                r"\1[REDACTED]\2",
                xml_str,
                flags=re.IGNORECASE,
            )
        # Redact attribute values
        for attr in self.attrs:
            xml_str = re.sub(
                rf'({re.escape(attr)}\s*=\s*")[^"]*(")',
                r"\1[REDACTED]\2",
                xml_str,
                flags=re.IGNORECASE,
            )
        return xml_str


# ---------------------------------------------------------------------------
# Prompt builder
# ---------------------------------------------------------------------------

_PROMPT_TEMPLATE = """\
You are an assistant that corrects XML to meet a provided machine-readable schema. \
The schema (JSON) follows xmldev spec. Only respond with the corrected XML content, nothing else. \
If you cannot correct without adding data you do not know, return the original XML wrapped in a \
<xmldev_error> tag containing a short reason.

Schema JSON (trim to the minimum relevant subset):
{schema_json}

Broken XML:
{broken_xml}

Constraints:
- Output must be valid XML and use the root element: {root_name}.
- Do not invent extra top-level elements.
- Maintain existing values where uncertain. When changing values, include a comment attribute \
xmldev_confidence="low|medium|high".
- Strictly follow enum values when provided.
- Do not output any explanatory text outside XML.

Return only the XML bytes."""


def build_prompt(schema: Schema, broken_xml: str, deterministic_errors: list[str] | None = None) -> str:
    """Build the LLM repair prompt."""
    # Minimal schema snippet: just the root spec as JSON
    schema_snippet = _schema_to_json_snippet(schema)
    prompt = _PROMPT_TEMPLATE.format(
        schema_json=schema_snippet,
        broken_xml=broken_xml,
        root_name=schema.root.name,
    )
    if deterministic_errors:
        error_block = "\n".join(f"  - {e}" for e in deterministic_errors[:10])
        prompt += f"\n\nDeterministic repair attempts and failures:\n{error_block}"
    return prompt


def _schema_to_json_snippet(schema: Schema) -> str:
    """Convert schema root to a compact JSON snippet for the prompt."""
    def _elem_to_dict(spec: Any) -> dict:
        d: dict = {"name": spec.name}
        if spec.text_type != "string":
            d["text_type"] = spec.text_type
        if spec.enum_values:
            d["enum_values"] = spec.enum_values
        if spec.attrs:
            d["attrs"] = {k: {"name": v.name, "type": v.type, "required": v.required} for k, v in spec.attrs.items()}
        if spec.children:
            d["children"] = [{"spec": _elem_to_dict(cs.spec), "min_occurs": cs.min_occurs} for cs in spec.children]
        return d

    return json.dumps({"root": _elem_to_dict(schema.root)}, indent=2)


# ---------------------------------------------------------------------------
# LLM Client
# ---------------------------------------------------------------------------

class LLMError(RuntimeError):
    """Raised when LLM returns invalid XML after retry (XM05)."""
    code = "XM05"


class LLMClient:
    """
    OpenAI-compatible LLM client for XML repair.

    Parameters
    ----------
    config : LLMConfig
    """

    def __init__(self, config: LLMConfig) -> None:
        self.config = config
        self._redactor = PIIRedactor(
            tags=config.redaction_tags,
            attrs=config.redaction_attrs,
        )
        self._client = OpenAI(
            base_url=config.base_url or None,
            api_key=config.api_key or os.environ.get("OPENAI_API_KEY", ""),
            timeout=config.timeout,
        )

    def repair(
        self,
        schema: Schema,
        broken_xml: str,
        deterministic_errors: list[str] | None = None,
    ) -> str:
        """
        Call the LLM to repair broken_xml.

        Returns the corrected XML string.
        Raises LLMError (XM05) if the LLM response fails validation after one retry.
        """
        if self.config.redact:
            broken_xml_to_send = self._redactor.redact(broken_xml)
        else:
            broken_xml_to_send = broken_xml

        prompt = build_prompt(schema, broken_xml_to_send, deterministic_errors)
        logger.info("Sending LLM repair request (model=%s).", self.config.model_id)

        result = self._call_llm(prompt)
        logger.debug("LLM raw response (first 500 chars): %s", result[:500])

        if self._is_xmldev_error(result):
            raise LLMError(
                f"[XM05] LLM could not repair XML: {self._extract_error_reason(result)}"
            )

        # Validate the LLM output
        if self._looks_like_valid_xml(result):
            return result

        # One retry with a stricter prompt
        logger.warning("LLM returned invalid XML; retrying with stricter prompt.")
        strict_prompt = prompt + (
            "\n\nIMPORTANT: Your previous response was not valid XML. "
            "Return ONLY valid, well-formed XML, nothing else."
        )
        result = self._call_llm(strict_prompt)

        if self._looks_like_valid_xml(result) and not self._is_xmldev_error(result):
            return result

        raise LLMError(
            "[XM05] LLM returned invalid XML after retry. "
            f"Response preview: {result[:200]}"
        )

    def _call_llm(self, prompt: str) -> str:
        try:
            response = self._client.chat.completions.create(
                model=self.config.model_id,
                messages=[{"role": "user", "content": prompt}],
                temperature=self.config.temperature,
                max_tokens=self.config.max_tokens,
                top_p=1,
            )
            return (response.choices[0].message.content or "").strip()
        except OpenAIError as exc:
            raise LLMError(f"[XM05] LLM API call failed: {exc}") from exc

    @staticmethod
    def _looks_like_valid_xml(text: str) -> bool:
        try:
            from lxml import etree
            etree.fromstring(text.encode("utf-8"))
            return True
        except Exception:
            return False

    @staticmethod
    def _is_xmldev_error(text: str) -> bool:
        return "<xmldev_error>" in text

    @staticmethod
    def _extract_error_reason(text: str) -> str:
        m = re.search(r"<xmldev_error>(.*?)</xmldev_error>", text, re.DOTALL)
        return m.group(1).strip() if m else "unknown"
