"""AI-powered edge proposals using LLM providers."""

from __future__ import annotations

import json

from orxaq.kernel.world_dag import WorldDAG

# Domain knowledge: known causal relationships in credit risk.
CREDIT_PRIORS: dict[str, list[str]] = {
    "gdp_growth": ["unemployment_rate", "default_rate", "origination_volume"],
    "unemployment_rate": ["default_rate", "delinquency_rate", "income"],
    "fed_funds_rate": ["credit_spread", "origination_volume", "hpi"],
    "hpi": ["ltv", "recovery_rate", "lgd"],
    "credit_score": ["default_rate", "credit_spread"],
    "ltv": ["lgd", "default_rate"],
    "dti": ["default_rate", "delinquency_rate"],
    "delinquency_rate": ["charge_off_rate", "default_rate"],
    "charge_off_rate": ["net_charge_off", "provision_expense"],
    "recovery_rate": ["lgd", "net_charge_off"],
    "default_rate": ["lgd", "charge_off_rate", "cecl_allowance"],
    "lgd": ["net_charge_off", "cecl_allowance"],
    "net_charge_off": ["cecl_allowance", "provision_expense"],
    "cecl_allowance": ["coverage_ratio", "provision_expense"],
    "inflation_rate": ["fed_funds_rate", "hpi"],
    "prepayment_rate": ["origination_volume"],
    "origination_volume": ["delinquency_rate"],
}


class EdgeProposer:
    """Proposes causal edges for a World DAG.

    Strategy:
    1. Use domain priors (CREDIT_PRIORS) for instant suggestions.
    2. If an LLM provider is available, enhance proposals with AI.
    """

    def __init__(self, dag: WorldDAG, provider=None) -> None:
        self.dag = dag
        self.provider = provider

    def propose_for_node(self, node_id: str, max_proposals: int = 5) -> list[dict]:
        """Propose edges involving a specific node."""
        proposals: list[dict] = []
        proposals.extend(self._prior_proposals(node_id))
        if self.provider is not None:
            proposals.extend(self._llm_proposals(node_id))

        seen: set[str] = set()
        unique: list[dict] = []
        for p in proposals:
            key = f"{p['source']}->{p['target']}"
            if key not in seen:
                seen.add(key)
                unique.append(p)

        unique.sort(key=lambda x: x.get("confidence", 0), reverse=True)
        return unique[:max_proposals]

    def propose_all(self, max_proposals: int = 10) -> list[dict]:
        """Propose edges for the entire graph."""
        proposals: list[dict] = []
        for node_id in self.dag._entities:
            proposals.extend(self._prior_proposals(node_id))

        seen: set[str] = set()
        unique: list[dict] = []
        for p in proposals:
            key = f"{p['source']}->{p['target']}"
            if key not in seen:
                seen.add(key)
                unique.append(p)

        unique.sort(key=lambda x: x.get("confidence", 0), reverse=True)
        return unique[:max_proposals]

    def _prior_proposals(self, node_id: str) -> list[dict]:
        """Generate proposals from domain knowledge priors."""
        proposals: list[dict] = []
        node_lower = node_id.lower()

        targets = CREDIT_PRIORS.get(node_lower, [])
        for target in targets:
            if target in self.dag._entities and node_id in self.dag._entities:
                edge_id = f"{node_id}\u2192{target}"
                if edge_id not in self.dag._edges:
                    proposals.append(
                        {
                            "source": node_id,
                            "target": target,
                            "confidence": 0.7,
                            "reason": "domain prior (credit risk ontology)",
                            "method": "prior",
                        }
                    )

        for source, tgts in CREDIT_PRIORS.items():
            if node_lower in tgts and source in self.dag._entities:
                edge_id = f"{source}\u2192{node_id}"
                if edge_id not in self.dag._edges:
                    proposals.append(
                        {
                            "source": source,
                            "target": node_id,
                            "confidence": 0.7,
                            "reason": "domain prior (credit risk ontology)",
                            "method": "prior",
                        }
                    )

        return proposals

    def _llm_proposals(self, node_id: str) -> list[dict]:
        """Generate proposals using an LLM provider."""
        existing_nodes = sorted(self.dag._entities.keys())
        existing_edges = [f"{e.source} -> {e.target}" for e in self.dag._edges.values()]

        prompt = (
            f"Given a causal graph for credit risk analysis with these variables:\n"
            f"{', '.join(existing_nodes)}\n\n"
            f"Existing edges:\n{chr(10).join(existing_edges) if existing_edges else '(none)'}\n\n"
            f"Suggest up to 3 new causal edges involving '{node_id}'. "
            f"Return JSON array of objects with: source, target, confidence (0-1), reason.\n"
            f"Only suggest edges where a plausible causal mechanism exists."
        )

        response = self.provider.generate(prompt, max_tokens=512, temperature=0.3)

        if response.metadata.get("error"):
            return []

        try:
            text = response.text.strip()
            if text.startswith("```"):
                text = text.split("\n", 1)[1].rsplit("```", 1)[0]
            proposals_raw = json.loads(text)
            if not isinstance(proposals_raw, list):
                return []

            proposals: list[dict] = []
            for p in proposals_raw:
                if isinstance(p, dict) and "source" in p and "target" in p:
                    proposals.append(
                        {
                            "source": p["source"],
                            "target": p["target"],
                            "confidence": float(p.get("confidence", 0.5)),
                            "reason": p.get("reason", "LLM suggestion"),
                            "method": "llm",
                        }
                    )
            return proposals
        except (json.JSONDecodeError, ValueError):
            return []
