"""Stage 2: LLM-assisted entity and relationship extraction.

Uses an OpenAI-compatible LLM (e.g., Qwen on Rascal) to extract
entities and relationships from document sections that deterministic
extraction cannot handle.

Operates on document SECTIONS (not chunks) — each section is sent
as a structured extraction prompt with the entity type registry.
"""

from __future__ import annotations

import json
import logging
import os
import urllib.request

from axiom.graph.schema import Edge, Entity, EntityTypeRegistry

log = logging.getLogger(__name__)

_DEFAULT_PROMPT = """\
Extract entities and relationships from the following technical document section.

Entity types: {entity_types}
Relationship types: {relationship_types}

Return a JSON object with:
{{
  "entities": [
    {{"label": "EntityType", "name": "entity name", "properties": {{}}, "confidence": 0.0-1.0}}
  ],
  "relationships": [
    {{"type": "REL_TYPE", "from_name": "...", "from_label": "...", "to_name": "...", "to_label": "...", "confidence": 0.0-1.0}}
  ]
}}

Only extract what is explicitly stated or strongly implied.
Do not infer relationships that require domain expertise beyond the text.
If nothing can be extracted, return {{"entities": [], "relationships": []}}.

Document section:
---
{text}
---
"""


def extract_from_section(
    text: str,
    source_path: str,
    registry: EntityTypeRegistry | None = None,
    llm_url: str | None = None,
    llm_model: str | None = None,
    llm_api_key: str | None = None,
) -> tuple[list[Entity], list[Edge]]:
    """Extract entities and relationships from a document section using LLM.

    Args:
        text: Section text (NOT a chunk — full section)
        source_path: Source document path (for provenance)
        registry: Entity type registry (for prompt construction)
        llm_url: OpenAI-compatible API base URL
        llm_model: Model name
        llm_api_key: API key

    Returns:
        (entities, edges) tuple
    """
    if not text.strip():
        return [], []

    # Resolve LLM endpoint
    url = llm_url or os.environ.get("AXIOM_LLM_URL", "")
    model = llm_model or os.environ.get("AXIOM_LLM_MODEL", "")
    api_key = llm_api_key or os.environ.get("AXIOM_LLM_KEY", "")

    if not url:
        log.debug("No LLM URL configured — skipping LLM extraction")
        return [], []

    if registry is None:
        registry = EntityTypeRegistry()

    # Build prompt
    entity_types = ", ".join(registry.entity_labels())
    rel_types = ", ".join(registry.relationship_types())
    prompt = _DEFAULT_PROMPT.format(
        entity_types=entity_types,
        relationship_types=rel_types,
        text=text[:4000],  # Truncate to avoid token limits
    )

    # Call LLM
    try:
        response = _call_llm(url, model, api_key, prompt)
    except Exception as e:
        log.warning("LLM extraction failed for %s: %s", source_path, e)
        return [], []

    # Parse response
    return _parse_extraction_response(response, source_path)


def _call_llm(url: str, model: str, api_key: str, prompt: str) -> str:
    """Call an OpenAI-compatible chat completions endpoint."""
    endpoint = f"{url.rstrip('/')}/chat/completions"

    payload = {
        "model": model,
        "messages": [
            {
                "role": "system",
                "content": "You are an entity extraction assistant. Return only valid JSON.",
            },
            {"role": "user", "content": prompt},
        ],
        "temperature": 0.1,
        "max_tokens": 2000,
    }

    headers = {"Content-Type": "application/json"}
    if api_key:
        headers["Authorization"] = f"Bearer {api_key}"

    data = json.dumps(payload).encode("utf-8")
    req = urllib.request.Request(endpoint, data=data, headers=headers, method="POST")

    with urllib.request.urlopen(req, timeout=60) as resp:
        body = json.loads(resp.read())

    return body["choices"][0]["message"]["content"]


def _parse_extraction_response(
    response: str,
    source_path: str,
) -> tuple[list[Entity], list[Edge]]:
    """Parse LLM extraction response into Entity and Edge objects."""
    entities = []
    edges = []

    # Extract JSON from response (may be wrapped in markdown code fence)
    json_str = response.strip()
    if json_str.startswith("```"):
        lines = json_str.split("\n")
        json_str = "\n".join(line for line in lines if not line.strip().startswith("```"))

    try:
        data = json.loads(json_str)
    except json.JSONDecodeError:
        log.warning("Could not parse LLM extraction response as JSON")
        return [], []

    for e in data.get("entities", []):
        try:
            entities.append(
                Entity(
                    label=e.get("label", "Concept"),
                    name=e.get("name", ""),
                    properties=e.get("properties", {}),
                    confidence=float(e.get("confidence", 0.5)),
                    provenance="llm_extracted",
                    source_path=source_path,
                )
            )
        except Exception:
            continue

    for r in data.get("relationships", []):
        try:
            edges.append(
                Edge(
                    rel_type=r.get("type", "REFERENCES"),
                    from_name=r.get("from_name", ""),
                    from_label=r.get("from_label", "Concept"),
                    to_name=r.get("to_name", ""),
                    to_label=r.get("to_label", "Concept"),
                    confidence=float(r.get("confidence", 0.5)),
                    provenance="llm_extracted",
                    source_chunk_id=None,
                )
            )
        except Exception:
            continue

    return entities, edges
