"""
ContextAnalysisAgent: LangGraph workflow for context analysis.

This module implements a stateful workflow for analyzing code context using
LangGraph's StateGraph. The workflow includes file loading, dependency analysis,
and semantic relevance scoring.

Workflow Stages
---------------
1. **load_files**: Loads files using LangChain's TextLoader
2. **analyze_dependencies**: Extracts imports and dependencies using LLM
3. **chunk_documents**: Splits large files into overlapping chunks
4. **score_relevance**: Scores relevance of each file (0.0-1.0)
5. **summarize_context**: Creates high-level summary

Example Usage
-------------
::

    from cleveragents.agents import ContextAnalysisAgent, ContextAnalysisState

    # Create the agent
    agent = ContextAnalysisAgent(chunk_size=2000, chunk_overlap=200)

    # Prepare state
    state: ContextAnalysisState = {
        "file_paths": ["src/main.py", "src/utils.py"],
        "documents": [],
        "dependencies": {},
        "chunks": [],
        "relevance_scores": {},
        "summary": "",
        "error": None,
    }

    # Execute with thread isolation
    result = agent.invoke(state, config={"configurable": {"thread_id": "analysis-1"}})
"""

from collections.abc import AsyncIterator, Iterator
from pathlib import Path
from typing import Any, TypedDict, TypeVar, cast

from langchain_community.document_loaders import (
    TextLoader,  # type: ignore[import-untyped]
)
from langchain_core.documents import Document
from langchain_core.language_models import BaseLanguageModel
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import PromptTemplate  # type: ignore[attr-defined]
from langchain_core.runnables import Runnable, RunnableSequence
from langgraph.checkpoint.memory import MemorySaver  # type: ignore[import-untyped]
from langgraph.graph import END, StateGraph  # type: ignore[import-untyped]

ChainInput = TypeVar("ChainInput")
ChainOutput = TypeVar("ChainOutput")


class ContextAnalysisState(TypedDict):
    """State structure for context analysis workflow.

    Attributes:
        file_paths: List of file paths to analyze
        documents: Loaded documents from files
        dependencies: Extracted dependencies and imports
        summary: High-level summary of the context
        relevance_scores: Relevance scores for each file
        chunks: Chunked documents for large files
        error: Error message if any
    """

    file_paths: list[str]
    documents: list[Document]
    dependencies: dict[str, list[str]]
    summary: str
    relevance_scores: dict[str, float]
    chunks: list[Document]
    error: str | None


class ContextAnalysisAgent:
    """LangGraph workflow for analyzing code context.

    This workflow orchestrates the context analysis process through multiple nodes:
    1. load_files: Loads files and creates Document objects
    2. analyze_dependencies: Extracts imports and dependencies
    3. chunk_documents: Splits large files into chunks
    4. score_relevance: Scores relevance of each file
    5. summarize_context: Creates high-level summary

    The workflow includes error handling and checkpointing for resumable execution.
    """

    def __init__(
        self,
        llm: BaseLanguageModel | None = None,
        chunk_size: int = 2000,
        chunk_overlap: int = 200,
        retry_attempts: int = 3,
    ):
        """Initialize the context analysis agent.

        Args:
            llm: Language model to use (defaults to FakeListLLM for testing)
            chunk_size: Maximum size of each chunk in characters
            chunk_overlap: Overlap between chunks in characters
            retry_attempts: Number of times to retry transient LLM failures
        """
        self.chunk_size = chunk_size
        self.chunk_overlap = chunk_overlap
        self.retry_attempts = max(1, retry_attempts)

        # Initialize LLM - an LLM must be provided explicitly
        if llm is None:
            raise ValueError(
                "No LLM provider configured. "
                "Set provider credentials or enable core.mock_providers for testing."
            )
        self.llm: BaseLanguageModel = llm

        # Create prompts
        self._create_prompts()

        # Build the graph
        self.graph = self._build_graph()

        # Compile with checkpointing
        self.checkpointer = MemorySaver()
        self.app = self.graph.compile(checkpointer=self.checkpointer)

    def _create_prompts(self) -> None:
        """Create prompt templates for each workflow node."""
        self.dependency_prompt: Any = PromptTemplate(
            template=(
                "Analyze the following code and extract all imports and dependencies.\n"
                "List them in a structured format.\n\n"
                "Code:\n"
                "{code}\n\n"
                "Extracted Dependencies:"
            ),
            input_variables=["code"],
        )

        self.relevance_prompt: Any = PromptTemplate(
            template=(
                "Score the relevance of this code file for the given task.\n"
                "Provide a score from 0.0 to 1.0 and a brief explanation.\n\n"
                "Task: {task}\n"
                "File: {file_name}\n"
                "Code Preview:\n"
                "{code_preview}\n\n"
                "Relevance Score and Explanation:"
            ),
            input_variables=["task", "file_name", "code_preview"],
        )

        self.summary_prompt: Any = PromptTemplate(
            template=(
                "Provide a high-level summary of the following codebase context.\n\n"
                "Files analyzed: {file_count}\n"
                "Total size: {total_size} characters\n"
                "Dependencies found: {dependency_count}\n\n"
                "Key files:\n"
                "{key_files}\n\n"
                "Summary:"
            ),
            input_variables=[
                "file_count",
                "total_size",
                "dependency_count",
                "key_files",
            ],
        )

    def _with_retry(
        self, chain: Runnable[ChainInput, ChainOutput]
    ) -> Runnable[ChainInput, ChainOutput]:
        """Wrap a runnable with LangChain's retry support."""

        attempts = max(1, self.retry_attempts)
        if hasattr(chain, "with_retry"):
            return chain.with_retry(stop_after_attempt=attempts)
        return chain

    def _build_graph(self) -> StateGraph:
        """Build the LangGraph workflow."""
        workflow = StateGraph(ContextAnalysisState)

        # Add nodes
        workflow.add_node("load_files", self._load_files)
        workflow.add_node("analyze_dependencies", self._analyze_dependencies)
        workflow.add_node("chunk_documents", self._chunk_documents)
        workflow.add_node("score_relevance", self._score_relevance)
        workflow.add_node("summarize_context", self._summarize_context)

        # Define edges
        workflow.set_entry_point("load_files")
        workflow.add_edge("load_files", "analyze_dependencies")
        workflow.add_edge("analyze_dependencies", "chunk_documents")
        workflow.add_edge("chunk_documents", "score_relevance")
        workflow.add_edge("score_relevance", "summarize_context")
        workflow.add_edge("summarize_context", END)

        return workflow

    def _load_files(self, state: ContextAnalysisState) -> dict[str, Any]:
        """Load files and create Document objects.

        Args:
            state: Current workflow state

        Returns:
            Updated state with loaded documents
        """
        preloaded_docs = state["documents"]
        if preloaded_docs:
            return {
                "documents": preloaded_docs,
                "error": state["error"],
            }

        documents: list[Document] = []
        errors: list[str] = []

        for file_path in state["file_paths"]:
            try:
                path = Path(file_path)
                if not path.exists():
                    errors.append(f"File not found: {file_path}")
                    continue

                if not path.is_file():
                    errors.append(f"Not a file: {file_path}")
                    continue

                loader = TextLoader(str(path))
                loaded_docs: list[Document] = loader.load()
                documents.extend(loaded_docs)
            except Exception as exc:  # pragma: no cover - defensive
                errors.append(f"Error loading {file_path}: {exc!s}")

        error_msg: str | None = "; ".join(errors) if errors else None

        return {
            "documents": documents,
            "error": error_msg,
        }

    def _analyze_dependencies(self, state: ContextAnalysisState) -> dict[str, Any]:
        """Analyze dependencies and imports in the loaded files."""
        dependencies: dict[str, list[str]] = {}
        base_chain = cast(
            RunnableSequence[dict[str, Any], str],
            self.dependency_prompt | self.llm | StrOutputParser(),
        )
        chain = self._with_retry(base_chain)

        for doc_value in state["documents"]:
            doc = cast(Any, doc_value)
            metadata = cast(dict[str, Any], getattr(doc, "metadata", {}))
            file_path = str(metadata.get("source", "unknown"))

            try:
                snippet = doc.page_content[:1000]
                result = chain.invoke({"code": snippet})
                dependencies[file_path] = self._parse_dependencies(result)
            except Exception as exc:  # pragma: no cover - defensive
                dependencies[file_path] = []
                new_error = f"Dependency analysis error for {file_path}: {exc!s}"
                current_error = state.get("error")
                combined_error = (
                    f"{current_error}; {new_error}"
                    if isinstance(current_error, str) and current_error
                    else new_error
                )
                return {"dependencies": dependencies, "error": combined_error}

        return {"dependencies": dependencies}

    def _parse_dependencies(self, llm_output: str) -> list[str]:
        """Parse dependencies from LLM output."""
        deps: list[str] = []
        for line in llm_output.split("\n"):
            line = line.strip()
            if line and not line.startswith("#"):
                parts = (
                    line.replace("[", "")
                    .replace("]", "")
                    .replace("'", "")
                    .replace('"', "")
                    .split(",")
                )
                deps.extend(part.strip() for part in parts if part.strip())
        return deps[:10]

    def _chunk_documents(self, state: ContextAnalysisState) -> dict[str, Any]:
        """Chunk large documents for processing."""
        chunks: list[Document] = []

        for doc in state["documents"]:
            content = doc.page_content
            metadata_source = cast(
                dict[str, Any] | None, getattr(doc, "metadata", None)
            )
            metadata: dict[str, Any] = metadata_source or {}

            if len(content) <= self.chunk_size:
                chunks.append(doc)
                continue

            step = max(1, self.chunk_size - self.chunk_overlap)
            for index, start in enumerate(range(0, len(content), step)):
                chunk_content = content[start : start + self.chunk_size]
                chunk_metadata: dict[str, Any] = {
                    **metadata,
                    "chunk_index": index,
                }
                chunks.append(
                    Document(page_content=chunk_content, metadata=chunk_metadata)
                )

        return {"chunks": chunks}

    def _score_relevance(self, state: ContextAnalysisState) -> dict[str, Any]:
        """Score relevance of each file for the analysis task."""
        relevance_scores: dict[str, float] = {}
        files_seen: set[str] = set()
        base_chain = cast(
            RunnableSequence[dict[str, Any], str],
            self.relevance_prompt | self.llm | StrOutputParser(),
        )
        chain = self._with_retry(base_chain)

        for chunk_value in state["chunks"]:
            chunk = cast(Any, chunk_value)
            metadata = cast(dict[str, Any], getattr(chunk, "metadata", {}))
            file_path = str(metadata.get("source", "unknown"))
            if file_path in files_seen:
                continue
            files_seen.add(file_path)

            try:
                result = chain.invoke(
                    {
                        "task": "code analysis",
                        "file_name": file_path,
                        "code_preview": str(chunk.page_content)[:500],
                    }
                )
                relevance_scores[file_path] = self._parse_relevance_score(result)
            except Exception as exc:  # pragma: no cover - defensive
                relevance_scores[file_path] = 0.5
                current_error = state.get("error")
                new_error = f"Relevance scoring error for {file_path}: {exc!s}"
                combined_error = (
                    f"{current_error}; {new_error}"
                    if isinstance(current_error, str) and current_error
                    else new_error
                )
                return {
                    "relevance_scores": relevance_scores,
                    "error": combined_error,
                }

        return {"relevance_scores": relevance_scores}

    def _parse_relevance_score(self, llm_output: str) -> float:
        """Parse relevance score from LLM output."""
        import re

        matches = re.findall(r"0?\.\d+|1\.0|[01]", llm_output.lower())
        if matches:
            try:
                score = float(matches[0])
                return max(0.0, min(1.0, score))
            except ValueError:  # pragma: no cover - defensive
                pass

        if "high" in llm_output.lower():
            return 0.8
        if "low" in llm_output.lower():
            return 0.3
        return 0.5

    def _summarize_context(self, state: ContextAnalysisState) -> dict[str, Any]:
        """Create a high-level summary of the analyzed context."""
        try:
            file_count = len(state["documents"])
            total_size = sum(len(doc.page_content) for doc in state["documents"])
            dependency_count = sum(len(deps) for deps in state["dependencies"].values())

            sorted_files: list[tuple[str, float]] = sorted(
                state["relevance_scores"].items(),
                key=lambda item: item[1],
                reverse=True,
            )[:3]
            key_files = "\n".join(
                f"- {path} (score: {score:.2f})" for path, score in sorted_files
            )

            base_chain = cast(
                RunnableSequence[dict[str, Any], str],
                self.summary_prompt | self.llm | StrOutputParser(),
            )
            chain = self._with_retry(base_chain)
            summary = chain.invoke(
                {
                    "file_count": file_count,
                    "total_size": total_size,
                    "dependency_count": dependency_count,
                    "key_files": key_files,
                }
            )

            return {"summary": summary}
        except Exception as exc:  # pragma: no cover - defensive
            current_error = state.get("error")
            new_error = f"Summarization error: {exc!s}"
            combined_error = (
                f"{current_error}; {new_error}"
                if isinstance(current_error, str) and current_error
                else new_error
            )
            return {"summary": "Context analysis failed", "error": combined_error}

    def invoke(
        self, input_state: ContextAnalysisState, config: dict[str, Any] | None = None
    ) -> ContextAnalysisState:
        """Synchronously execute the context analysis workflow."""
        config = config or {}
        result = self.app.invoke(cast(dict[str, Any], input_state), config)
        return cast(ContextAnalysisState, result)

    async def ainvoke(
        self, input_state: ContextAnalysisState, config: dict[str, Any] | None = None
    ) -> ContextAnalysisState:
        """Asynchronously execute the context analysis workflow."""
        config = config or {}
        result = await self.app.ainvoke(cast(dict[str, Any], input_state), config)
        return cast(ContextAnalysisState, result)

    def stream(
        self, input_state: ContextAnalysisState, config: dict[str, Any] | None = None
    ) -> Iterator[dict[str, Any]]:
        """Stream the context analysis workflow execution."""
        config = config or {}
        app_obj = cast(Any, self.app)
        stream_iter = cast(
            Iterator[dict[str, Any]],
            app_obj.stream(cast(dict[str, Any], input_state), config),
        )
        yield from stream_iter

    async def astream(
        self, input_state: ContextAnalysisState, config: dict[str, Any] | None = None
    ) -> AsyncIterator[dict[str, Any]]:
        """Asynchronously stream the context analysis workflow execution."""
        config = config or {}
        app_obj = cast(Any, self.app)
        astream_iter = cast(
            AsyncIterator[dict[str, Any]],
            app_obj.astream(cast(dict[str, Any], input_state), config),
        )
        async for event in astream_iter:
            yield event
