"""
llm_agents.logic_agent
=======================
Hybrid LLM + propositional KB agent for logical reasoning (AIMA Ch 7-9).

The agent can:
  1. Accept new facts / rules (TELL)
  2. Answer logical queries (ASK) by:
     a. First checking the local KB with DPLL (fast, exact)
     b. Falling back to LLM reasoning if the KB cannot answer
  3. Translate natural language into propositional sentences (via LLM)

Usage
-----
    from ml_ai_library.llm_agents.logic_agent import LogicAgent

    agent = LogicAgent(model="gpt-4o-mini")
    agent.tell("It is raining")
    agent.tell("If it is raining then the ground is wet")
    print(agent.ask("Is the ground wet?"))   # True

    # Natural-language tell
    agent.nl_tell("Birds can fly, but penguins cannot fly and penguins are birds.")
    print(agent.nl_ask("Can a penguin fly?"))   # "No, penguins cannot fly."
"""

from __future__ import annotations
from typing import List

from ml_ai_library.llm_agents.base_agent import BaseAgent
from ml_ai_library.norvig.logic import KB, parse_expr

_TRANSLATE_SYSTEM = """\
You are a propositional logic translator.

Given a natural-language statement, convert it to one or more propositional
logic sentences using this syntax:
  - Atoms: single CamelCase words (e.g. IsRaining, GroundIsWet)
  - NOT: ~Atom
  - AND: Atom1 & Atom2
  - OR:  Atom1 | Atom2
  - IMPLIES: Atom1 => Atom2
  - IFF: Atom1 <=> Atom2

Return ONLY the list of sentences, one per line. No explanation.
"""

_REASON_SYSTEM = """\
You are a logical reasoning assistant. You have access to the following
known facts and rules. Use them to answer the question logically.

Answer with "Yes", "No", or "Unknown", followed by a brief explanation.
"""


class LogicAgent:
    """
    Hybrid propositional logic agent.

    Parameters
    ----------
    model : litellm model string (used for NL→logic translation and fallback)
    """

    def __init__(self, model: str = "gpt-4o-mini", temperature: float = 0.0):
        self.kb      = KB()
        self.model   = model
        self._nl_facts: List[str] = []          # natural-language memory
        self._translator = BaseAgent(model=model, system=_TRANSLATE_SYSTEM,
                                     temperature=temperature)
        self._reasoner   = BaseAgent(model=model, system=_REASON_SYSTEM,
                                     temperature=temperature)

    # ------------------------------------------------------------------
    # Formal interface
    # ------------------------------------------------------------------

    def tell(self, sentence: str):
        """Add a propositional sentence to the KB (AIMA TELL)."""
        self.kb.tell(sentence)
        self._nl_facts.append(sentence)

    def ask(self, query: str) -> bool:
        """Query the KB with DPLL (AIMA ASK). Falls back to LLM if uncertain."""
        try:
            return self.kb.ask(query)
        except Exception:
            return self._llm_ask(query)

    # ------------------------------------------------------------------
    # Natural language interface
    # ------------------------------------------------------------------

    def nl_tell(self, natural_language: str):
        """
        Convert a natural-language statement to propositional sentences
        and add them to the KB.
        """
        self._translator.reset()
        raw = self._translator.chat(natural_language)
        for line in raw.strip().splitlines():
            line = line.strip()
            if line:
                try:
                    self.kb.tell(line)
                except Exception:
                    pass       # skip unparseable lines
        self._nl_facts.append(natural_language)

    def nl_ask(self, question: str) -> str:
        """
        Answer a natural-language yes/no question.
        First attempts DPLL on a propositional translation; falls back to LLM.
        """
        # Translate question to atom and try KB
        self._translator.reset()
        q_atom = self._translator.chat(question).strip().splitlines()[0]
        try:
            result = self.kb.ask(q_atom)
            return "Yes." if result else "No."
        except Exception:
            pass

        # LLM fallback with full fact list
        facts_str = "\n".join(f"- {f}" for f in self._nl_facts)
        self._reasoner.reset()
        prompt = f"Known facts:\n{facts_str}\n\nQuestion: {question}"
        return self._reasoner.chat(prompt)

    # ------------------------------------------------------------------
    # Internal LLM fallback
    # ------------------------------------------------------------------

    def _llm_ask(self, query: str) -> bool:
        facts_str = "\n".join(str(c) for c in self.kb.clauses)
        self._reasoner.reset()
        reply = self._reasoner.chat(
            f"KB:\n{facts_str}\n\nQuery: Does '{query}' follow? Answer Yes or No."
        )
        return reply.strip().lower().startswith("yes")
