"""Prompt templates for LIMEN-AI + LLM orchestration."""

from __future__ import annotations

import re
from textwrap import dedent
from typing import Optional, Sequence

from .schema import PredicateSchema, SchemaRegistry


def extract_candidate_tokens(text: str, limit: int = 12) -> list[str]:
    """Return literal entity strings that appear inside a text chunk."""

    tokens: list[str] = []
    seen: set[str] = set()

    for match in re.findall(r"`([^`]+)`", text):
        candidate = match.strip()
        if candidate and candidate not in seen:
            tokens.append(candidate)
            seen.add(candidate)
        if len(tokens) >= limit:
            return tokens

    for match in re.findall(r"\b[A-Z][A-Za-z0-9_.-]*\b", text):
        candidate = match.strip()
        if candidate and candidate not in seen:
            tokens.append(candidate)
            seen.add(candidate)
        if len(tokens) >= limit:
            break

    return tokens


def _select_example_schema(registry: SchemaRegistry) -> Optional[PredicateSchema]:
    """Choose a predicate schema to showcase in the extraction example."""

    # Prefer scenario predicates (introduced via heading fallback)
    for schema in registry.predicates.values():
        if "scenario" in schema.name.lower():
            return schema

    # Otherwise prefer the highest-arity predicate; fall back to first entry.
    sorted_schemas = sorted(
        registry.predicates.values(), key=lambda schema: schema.arity, reverse=True
    )
    if sorted_schemas:
        return sorted_schemas[0]
    return None


def build_extraction_prompt(
    chunk: str,
    registry: SchemaRegistry,
    *,
    candidate_tokens: Sequence[str] | None = None,
) -> str:
    """Prompt instructing the LLM to extract structured facts."""

    instruction = registry.as_instruction()
    tokens_block = ""
    scenario_note = ""
    example_block = ""
    example_schema = _select_example_schema(registry)
    if candidate_tokens:
        formatted = "\n".join(f"- {token}" for token in candidate_tokens[:12])
        tokens_block = (
            "\nConcrete entity tokens detected in the text (use these exact spellings in args):\n"
            f"{formatted}\n"
        )
        if example_schema:
            example_predicate = example_schema.name
            example_arity = max(1, example_schema.arity)
        else:
            example_predicate = "examplePredicate"
            example_arity = 2
        usable_args = list(candidate_tokens[: example_arity])
        if len(usable_args) == example_arity:
            args_literal = ", ".join(f'"{arg}"' for arg in usable_args)
            example_block = f"""
Example (each argument must be copied verbatim from the tokens listed above):
[
  {{
    "predicate": "{example_predicate}",
    "args": [{args_literal}],
    "confidence": 0.92
  }}
]
"""

    if any("scenario" in schema.name.lower() for schema in registry.predicates.values()):
        scenario_note = (
            "\nFor predicates whose names contain 'scenario', emit exactly one JSON object per scenario heading. "
            "Map **Initial actor** to argument #1, **Target asset** to argument #2, and the entire **Control evidence** line to argument #3."
        )

    prompt = f"""
You are an information extraction assistant. Analyse the following text chunk and emit
JSON with the format:
[
  {{
    "predicate": "predicate_name",
    "args": ["arg1", "arg2", ...],
    "confidence": 0.0-1.0,
    "provenance": "optional source snippet"
  }}
]
Only use predicates from the schema below, and make sure the number of arguments matches the definition.
If no facts are found, return an empty JSON list [].
Always fill the arguments with the concrete entity names that appear in the text.
Never echo template placeholders or inject example tokens that do not appear in the source.
{tokens_block}
{scenario_note}
{instruction}

Text chunk:
\"\"\"
{chunk}
\"\"\"

{example_block}
JSON:
"""
    return dedent(prompt).strip()


def build_schema_induction_prompt(
    text: str,
    max_predicates: int = 6,
    existing_predicates: list[str] | None = None,
    candidate_tokens: Sequence[str] | None = None,
) -> str:
    """Prompt that asks the LLM to propose predicate schemas from free-form text."""

    if existing_predicates:
        sorted_names = ", ".join(sorted(existing_predicates))
        existing_note = (
            "\n- The current schema already contains the following predicates: "
            f"{sorted_names}. Only propose additional predicate definitions that capture NEW concepts present in the narrative."
        )
    else:
        existing_note = ""

    token_block = ""
    if candidate_tokens:
        formatted = "\n".join(f"- {token}" for token in candidate_tokens[:12])
        token_block = (
            "\nConcrete entities mentioned in the narrative (align argument slots to these tokens when possible):\n"
            f"{formatted}\n"
        )

    prompt = f"""
You are designing predicate schemas for the LIMEN-AI probabilistic logic engine.
Read the narrative below and derive up to {max_predicates} predicate definitions that
would help represent its facts in first-order logic form.

Respond strictly in JSON:
{{
  "predicates": [
    {{
      "name": "camelCasePredicateName",
      "arity": 2,
      "arg_names": ["argOne", "argTwo"],
      "description": "Concise natural-language description"
    }}
  ]
}}

Guidelines:
- Use camelCase predicate names without spaces or punctuation.
- Set the arity equal to the length of arg_names.
- Each arg_names entry must be a descriptive camelCase identifier.
- Only include predicates that are clearly supported by the text.
- Avoid duplicate or overlapping predicates; prefer concise coverage.
- Under no circumstance repeat template placeholders such as "camelCasePredicateName",
  "argOne", "argTwo", or "Concise natural-language description".
- Every predicate must be tied to a concrete concept present in the narrative (e.g., riskCategory,
  mitigationControl, kpiMeasurement).{existing_note}
{token_block}

Bad example (do NOT output):
{{
  "predicates": [
    {{
      "name": "camelCasePredicateName",
      "arity": 2,
      "arg_names": ["argOne", "argTwo"],
      "description": "Concise natural-language description"
    }}
  ]
}}

Good example (structure only, invent actual concepts from the text instead of copying this):
{{
  "predicates": [
    {{
      "name": "multiFactorRequirement",
      "arity": 2,
      "arg_names": ["policyControl", "userRole"],
      "description": "Links MFA controls to the roles they protect"
    }},
    {{
      "name": "riskMeasurement",
      "arity": 3,
      "arg_names": ["riskType", "metricName", "kpiValue"],
      "description": "Captures KPIs associated with each ISO risk type"
    }}
  ]
}}

Narrative:
\"\"\"
{text}
\"\"\"

JSON:
"""
    return dedent(prompt).strip()


def build_query_prompt(
    question: str, registry: SchemaRegistry, known_constants: list[str] | None = None
) -> str:
    """Prompt that translates a user question into predicate calls."""

    instruction = registry.as_instruction()
    constants_block = ""
    constants_instructions = ""
    example_block = ""
    if known_constants:
        formatted = "\n".join(f"- {name}" for name in sorted(known_constants))
        constants_block = f"\nKnown constants already present in the KB:\n{formatted}\n"
        constants_instructions = (
            "Only use constants from the list above; do not invent descriptors or prefix them with labels such as "
            "'targetAsset:' or 'controlEvidence:'."
        )

    scenario_schema = next(
        (schema for schema in registry.predicates.values() if "scenario" in schema.name.lower()),
        None,
    )
    if scenario_schema and known_constants:
        usable_args = []
        for token in known_constants:
            if token not in usable_args:
                usable_args.append(token)
            if len(usable_args) == scenario_schema.arity:
                break
        if len(usable_args) == scenario_schema.arity:
            args_literal = ", ".join(f'"{arg}"' for arg in usable_args)
            example_block = f"""
Example (showing how to ask about a specific scenario predicate):
{{
  "queries": [
    {{"predicate": "{scenario_schema.name}", "args": [{args_literal}]}}
  ]
}}
"""

    prompt = f"""
Translate the user's question into a list of predicate calls from the schema below.
Respond in JSON with the format:
{{
  "queries": [
    {{"predicate": "name", "args": ["arg1", "arg2", ...]}}
  ]
}}
If the question cannot be mapped, respond with {{"queries": []}}.

{instruction}
{constants_block}
{constants_instructions}

Ensure every predicate uses exactly the number of arguments defined in the schema,
and every argument string must match an existing KB constant literally (same spelling,
capitalisation, and punctuation).

{example_block}

Question: \"{question}\"
JSON:
"""
    return dedent(prompt).strip()


def build_response_prompt(structured_answer: str, user_question: str) -> str:
    """Prompt instructing the LLM to turn structured data into natural language."""

    prompt = f"""
You are an analyst assistant. Convert the following structured LIMEN-AI answer into
a short, user-friendly explanation. Ensure you mention the key predicates involved and
highlight any provenance.

Question: {user_question}

Structured answer:
{structured_answer}

Response:
"""
    return dedent(prompt).strip()

