"""LLM-based extraction of research-level knowledge from paper text.

Extracts IDEAS (findings, claims, methods, research problems) — not keywords.
Uses full-paper extraction (single LLM call) for coherent, non-fragmented graphs.
Falls back to per-chunk extraction + consolidation for very long papers.
"""

from __future__ import annotations

import json
import logging
from typing import Optional

from langchain_core.language_models import BaseChatModel
from langchain_core.messages import SystemMessage, HumanMessage

from res_sum.schemas import (
    DomainConfig,
    ExtractedEntity,
    ExtractedRelationship,
    ExtractionResult,
)

logger = logging.getLogger("res_sum")

# Max characters to send in a single LLM call for full-paper extraction.
# If paper exceeds this, fall back to per-chunk + consolidation.
_FULL_PAPER_CHAR_LIMIT = 30_000

# Appended to custom prompts so the LLM always returns parseable JSON
_JSON_OUTPUT_INSTRUCTIONS = """

## Required Output Format

Return your extraction as ONLY valid JSON (no markdown, no explanation):
{
  "entities": [
    {
      "name": "short label for this idea (5-10 words)",
      "node_type": "FINDING",
      "description": "Full sentence describing this idea",
      "evidence": "Statistical support if applicable",
      "confidence": "strong, moderate, or weak"
    }
  ],
  "relationships": [
    {
      "source": "name of source node (must match an entity name above)",
      "target": "name of target node (must match an entity name above)",
      "relation_type": "SUPPORTED_BY",
      "evidence": "Brief explanation"
    }
  ]
}

Valid node_type values: RESEARCH_PROBLEM, FINDING, CLAIM, METHOD, STUDY_SYSTEM, LIMITATION
Valid relation_type values: ADDRESSES, HAS_FINDING, HAS_CLAIM, USES_METHOD, STUDIED_IN, SUPPORTED_BY, PRODUCED_BY, HAS_LIMITATION

If no extractable content is found, return: {"entities": [], "relationships": []}"""


def _build_extraction_prompt(domain_config: DomainConfig) -> str:
    """Build the system prompt for research-level knowledge extraction."""

    domain_context = ""
    if domain_config.entity_types:
        domain_hints = "\n".join(
            f"  - {et.description} (e.g., {', '.join(et.examples[:2])})"
            for et in domain_config.entity_types
            if et.examples
        )
        if domain_hints:
            domain_context = f"\nDomain-specific guidance:\n{domain_hints}\n"

    return f"""You are a research analyst building a structured knowledge map of a scientific paper. Extract the INTELLECTUAL CONTENT — the ideas, findings, and claims — as a connected graph of nodes and relationships.

Think of it as notes a researcher would take during a literature review: what is the paper about, what did they do, what did they find, and what do they conclude?
{domain_context}
## Node Types to Extract

**RESEARCH_PROBLEM** — The central question this paper addresses. Usually one per paper.

**METHOD** — Specific research approaches used. Be precise:
  - Good: "MaxEnt SDM with 5 bioclimatic variables and 5-fold cross-validation"
  - Bad: "statistical analysis" or "modeling"

**STUDY_SYSTEM** — What was studied and where, as a single description:
  - Good: "Cane toad (Rhinella marina) invasion from South America to Australia, 1935-2008"
  - Bad: "Australia" or "cane toad"

**FINDING** — Specific empirical results. Include effect direction, magnitude, and statistics:
  - Good: "Native and invaded niches are not equivalent (Schoener's D = 0.28, P = 0.020)"
  - Bad: "There was a niche shift"

**CLAIM** — Interpretive conclusions the authors draw beyond the direct data:
  - Good: "Release from biotic interactions, not evolved tolerances, explains the niche shift"
  - Bad: "The results are important"

**LIMITATION** — Specific caveats acknowledged by the authors.

## Relationships

Connect nodes to show the paper's logical structure:
- USES_METHOD: Research Problem → Method
- STUDIED_IN: Research Problem → Study System
- PRODUCED_BY: Finding → Method (which method produced this finding)
- ADDRESSES: Finding → Research Problem (this finding answers the question)
- SUPPORTED_BY: Claim → Finding (this evidence supports this claim)
- HAS_LIMITATION: Research Problem or Method → Limitation

## Important Guidelines
- Extract the COMPLETE paper in one coherent graph — every finding should connect to a method and the research problem
- Use consistent naming — if you mention "niche shift" in one node, don't call it "range shift" in another
- Each node's "name" should be a short label (5-10 words), the "description" has the full detail
- Include ALL key findings, not just the main one
""" + _JSON_OUTPUT_INSTRUCTIONS


def extract_from_full_paper(
    full_text: str,
    llm: BaseChatModel,
    domain_config: DomainConfig,
    paper_filename: str = "",
    custom_prompt: Optional[str] = None,
) -> ExtractionResult:
    """Extract knowledge from the full paper text in a single LLM call.

    This produces a coherent, non-fragmented graph because the LLM sees
    the entire paper and can create consistent node names and relationships.

    Args:
        full_text: The complete paper text (intro through conclusion).
        llm: The LLM to use.
        domain_config: Domain configuration.
        paper_filename: Paper filename for logging.
        custom_prompt: Optional user-provided prompt (overrides default).

    Returns:
        ExtractionResult with the paper's knowledge graph.
    """
    if custom_prompt:
        system_prompt = custom_prompt + _JSON_OUTPUT_INSTRUCTIONS
    else:
        system_prompt = _build_extraction_prompt(domain_config)

    messages = [
        SystemMessage(content=system_prompt),
        HumanMessage(content=f"Build a knowledge graph from this research paper:\n\n{full_text}"),
    ]

    try:
        response = llm.invoke(messages)
        return _parse_extraction_response(response.content, chunk_id=paper_filename)
    except Exception as e:
        logger.warning(f"Full-paper extraction failed for {paper_filename}: {e}")
        return ExtractionResult()


def extract_from_chunks_with_consolidation(
    chunks: list[str],
    chunk_ids: list[str],
    llm: BaseChatModel,
    domain_config: DomainConfig,
    custom_prompt: Optional[str] = None,
) -> ExtractionResult:
    """Extract per-chunk, then consolidate into a coherent graph.

    Used as fallback for very long papers that exceed the context limit.

    Args:
        chunks: Text chunks.
        chunk_ids: Corresponding IDs.
        llm: The LLM.
        domain_config: Domain config.
        custom_prompt: Optional custom prompt.

    Returns:
        Consolidated ExtractionResult.
    """
    # Step 1: Extract from each chunk
    all_entities: list[ExtractedEntity] = []
    all_relationships: list[ExtractedRelationship] = []

    if custom_prompt:
        system_prompt = custom_prompt + _JSON_OUTPUT_INSTRUCTIONS
    else:
        system_prompt = _build_extraction_prompt(domain_config)

    for i, (chunk, chunk_id) in enumerate(zip(chunks, chunk_ids)):
        logger.debug(f"  Extracting from chunk {i + 1}/{len(chunks)}")
        messages = [
            SystemMessage(content=system_prompt),
            HumanMessage(content=f"Extract research ideas from this section:\n\n{chunk}"),
        ]
        try:
            response = llm.invoke(messages)
            result = _parse_extraction_response(response.content, chunk_id=chunk_id)
            all_entities.extend(result.entities)
            all_relationships.extend(result.relationships)
        except Exception as e:
            logger.warning(f"Chunk extraction failed for {chunk_id}: {e}")

    if not all_entities:
        return ExtractionResult()

    # Step 2: Consolidate — ask LLM to merge duplicates
    consolidated = _consolidate_entities(all_entities, all_relationships, llm)
    return consolidated


def _consolidate_entities(
    entities: list[ExtractedEntity],
    relationships: list[ExtractedRelationship],
    llm: BaseChatModel,
) -> ExtractionResult:
    """Ask the LLM to merge duplicate/overlapping entities from chunk-based extraction."""

    entity_list = "\n".join(
        f"  - [{e.node_type}] {e.name}: {e.description[:100]}"
        for e in entities
    )

    messages = [
        SystemMessage(content="""You are merging extracted knowledge nodes from different sections of the same paper. Some nodes may be duplicates or near-duplicates referring to the same concept.

Review the list below and return a CLEAN, DEDUPLICATED list. For duplicates:
- Keep the more descriptive version
- Merge their evidence and descriptions

Return ONLY valid JSON:
{
  "entities": [{"name": "...", "node_type": "...", "description": "...", "evidence": "...", "confidence": "..."}],
  "merge_map": {"old_name": "kept_name", ...}
}"""),
        HumanMessage(content=f"Deduplicate these extracted nodes:\n\n{entity_list}"),
    ]

    try:
        response = llm.invoke(messages)
        content = response.content.strip()

        # Strip markdown fences
        if content.startswith("```"):
            content = content.split("\n", 1)[1] if "\n" in content else content[3:]
        if content.endswith("```"):
            content = content[:-3]
        content = content.strip()
        if content.startswith("json"):
            content = content[4:].strip()

        data = json.loads(content)

        # Build merged entities
        merged_entities = [
            ExtractedEntity(
                name=e.get("name", ""),
                node_type=e.get("node_type", "FINDING"),
                entity_type=e.get("node_type", "FINDING"),
                description=e.get("description", ""),
                evidence=e.get("evidence", ""),
                confidence=e.get("confidence", ""),
            )
            for e in data.get("entities", [])
            if e.get("name") and e.get("description")
        ]

        # Apply merge map to relationships
        merge_map = data.get("merge_map", {})
        merged_relationships = []
        for rel in relationships:
            source = merge_map.get(rel.source, rel.source)
            target = merge_map.get(rel.target, rel.target)
            merged_relationships.append(ExtractedRelationship(
                source=source,
                target=target,
                relation_type=rel.relation_type,
                evidence=rel.evidence,
            ))

        return ExtractionResult(entities=merged_entities, relationships=merged_relationships)

    except Exception as e:
        logger.warning(f"Consolidation failed: {e}. Using raw entities.")
        return ExtractionResult(entities=entities, relationships=relationships)


def _parse_extraction_response(content: str, chunk_id: str = "") -> ExtractionResult:
    """Parse LLM response into ExtractionResult."""
    content = content.strip()

    # Strip markdown code fences
    if content.startswith("```"):
        content = content.split("\n", 1)[1] if "\n" in content else content[3:]
    if content.endswith("```"):
        content = content[:-3]
    content = content.strip()
    if content.startswith("json"):
        content = content[4:].strip()

    try:
        data = json.loads(content)
    except json.JSONDecodeError as e:
        logger.warning(f"Failed to parse extraction JSON for {chunk_id}: {e}")
        return ExtractionResult()

    entities = [
        ExtractedEntity(
            name=e.get("name", ""),
            node_type=e.get("node_type", "FINDING"),
            entity_type=e.get("node_type", "FINDING"),
            description=e.get("description", ""),
            evidence=e.get("evidence", ""),
            confidence=e.get("confidence", ""),
            aliases=e.get("aliases", []),
            source_chunk_id=chunk_id,
        )
        for e in data.get("entities", [])
        if e.get("name") and e.get("description")
    ]

    relationships = [
        ExtractedRelationship(
            source=r.get("source", ""),
            target=r.get("target", ""),
            relation_type=r.get("relation_type", "ADDRESSES"),
            evidence=r.get("evidence", ""),
            source_chunk_id=chunk_id,
        )
        for r in data.get("relationships", [])
        if r.get("source") and r.get("target")
    ]

    return ExtractionResult(entities=entities, relationships=relationships)


# Keep batch_extract for backward compatibility
def batch_extract(
    chunks: list[str],
    chunk_ids: list[str],
    llm: BaseChatModel,
    domain_config: DomainConfig,
    custom_prompt: Optional[str] = None,
) -> list[ExtractionResult]:
    """Legacy per-chunk extraction. Prefer extract_from_full_paper instead."""
    results = []
    if custom_prompt:
        system_prompt = custom_prompt + _JSON_OUTPUT_INSTRUCTIONS
    else:
        system_prompt = _build_extraction_prompt(domain_config)

    for i, (chunk, chunk_id) in enumerate(zip(chunks, chunk_ids)):
        logger.debug(f"  Extracting from chunk {i + 1}/{len(chunks)}")
        messages = [
            SystemMessage(content=system_prompt),
            HumanMessage(content=f"Analyze this text:\n\n{chunk}"),
        ]
        try:
            response = llm.invoke(messages)
            result = _parse_extraction_response(response.content, chunk_id=chunk_id)
            results.append(result)
        except Exception as e:
            logger.warning(f"Extraction failed for chunk {chunk_id}: {e}")
            results.append(ExtractionResult())

    return results
