import os
import uuid
import tempfile
from typing import Annotated, List, TypedDict, Dict, Any

import streamlit as st

from langchain_core.documents import Document
from langchain_core.messages import SystemMessage, HumanMessage, AIMessage, ToolMessage
from langchain_core.tools import tool
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter

try:
    from langchain_chroma import Chroma
except Exception:
    from langchain_community.vectorstores import Chroma

try:
    from langchain_community.tools.tavily_search import TavilySearchResults
except Exception:
    TavilySearchResults = None

from langgraph.graph import StateGraph, END
from langgraph.graph.message import add_messages
from langgraph.prebuilt import ToolNode
from langgraph.checkpoint.memory import MemorySaver


st.set_page_config(
    page_title="Legal Agentic AI Project",
    page_icon="⚖️",
    layout="wide",
)


SAMPLE_CASE = '''
Case No: 2024/CR/12345
Court: District Court, US

Facts:
The defendant, Mr. Alex Carter, was charged with breach of contract dated January 15, 2023.
The plaintiff, ABC Corporation, claims that the defendant failed to deliver 1000 units of
electronic components as per the signed agreement. The contract specified delivery by March 31, 2023.

The defendant claims that due to unforeseen supply chain disruptions caused by the pandemic,
delivery was impossible. However, no force majeure clause was explicitly invoked in writing
within the stipulated 30-day notice period as mentioned in Section 12 of the contract.

The plaintiff seeks damages of Rs. 50,00,000 for business losses incurred due to delayed delivery.
The defendant argues that the damages claimed are excessive and not proven with adequate documentation.

Key Issues:
1. Was there a valid force majeure event?
2. Did the defendant provide proper notice?
3. Are the damages claimed reasonable and substantiated?
4. Was there any contributory negligence by the plaintiff?
'''


LEGAL_KNOWLEDGE_BASE = [
    '''
    Precedent Case: Smith v. Johnson (2022)
    Court ruled that force majeure must be invoked in writing within the contractual notice period.
    Failure to do so waives the right to claim force majeure defense.
    Citation: [2022] 3 SCC 145
    ''',
    '''
    Precedent Case: ABC Corp v. XYZ Ltd (2021)
    In breach of contract cases, damages must be proven with documentary evidence.
    Speculative damages without supporting documentation are not admissible.
    The burden of proof lies with the plaintiff.
    Citation: [2021] 2 BomHC 78
    ''',
    '''
    Legal Principle: Contributory Negligence
    If the plaintiff's actions contributed to the breach or losses, damages may be reduced proportionally.
    Courts have recognized this principle in commercial contract disputes.
    Reference: US Contract Act, 1872, Section 73
    ''',
    '''
    Force Majeure During Pandemic
    Courts have taken a nuanced view of COVID-19 as force majeure.
    The invoking party must demonstrate: (1) impossibility, not mere difficulty,
    (2) proper notice, and (3) mitigation efforts.
    Citation: Energy Watchdog v. CERC (2020)
    ''',
    '''
    Force Majeure Invocation - Notice Requirement
    Relevant Precedent: Smith v. Johnson (2022)
    Citation: [2022] 3 SCC 145
    Legal Principle:
    Force majeure must be invoked in writing within the contractual notice period.
    Failure to comply with the notice requirement waives the right to claim force majeure defense.

    Application to Current Case:
    The contract specified a 30-day notice period under Section 12.
    Mr. Alex Carter failed to invoke force majeure in writing within this stipulated period.
    This precedent strongly supports the plaintiff's position.
    ''',
    '''
    Burden of Proof for Damages
    Relevant Precedent: ABC Corp v. XYZ Ltd (2021)
    Citation: [2021] 2 BomHC 78
    Legal Principle:
    In breach of contract cases, damages must be proven with documentary evidence.
    Speculative damages without supporting documentation are not admissible.
    The burden of proof lies with the plaintiff.

    Application to Current Case:
    ABC Corporation claims Rs. 50,00,000 in damages.
    The plaintiff must provide concrete evidence such as lost orders, financial statements,
    alternative procurement costs, and documented business loss.
    ''',
    '''
    Relevant Precedent: Energy Watchdog v. CERC (2020)
    Legal Principle:
    Courts apply a nuanced, three-part test for COVID-19 related force majeure claims:
    1. Impossibility, not mere difficulty or inconvenience
    2. Proper notice as per contractual terms
    3. Mitigation efforts undertaken by the affected party

    Application to Current Case:
    The defendant must demonstrate absolute impossibility of performance, not merely increased cost or delay.
    The defendant must also show good faith mitigation efforts, such as seeking alternative suppliers.
    ''',
]


class AgentState(TypedDict):
    messages: Annotated[list, add_messages]
    case_text: str
    final_analysis: str


def secrets_ready() -> bool:
    return bool(os.getenv("OPENAI_API_KEY"))


@st.cache_resource(show_spinner=False)
def create_vectorstore(openai_key: str):
    docs = [Document(page_content=doc) for doc in LEGAL_KNOWLEDGE_BASE]

    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=1000,
        chunk_overlap=200,
    )
    splits = text_splitter.split_documents(docs)

    embedding_model = OpenAIEmbeddings(
        model="text-embedding-3-small",
        api_key=openai_key,
    )

    persist_dir = tempfile.mkdtemp(prefix="legal_agent_chroma_")

    vectorstore = Chroma.from_documents(
        documents=splits,
        embedding=embedding_model,
        collection_name=f"legal_cases_{uuid.uuid4().hex[:8]}",
        persist_directory=persist_dir,
    )

    return vectorstore, len(splits)


def build_legal_agent(model_name: str, temperature: float = 0.0):
    openai_key = os.getenv("OPENAI_API_KEY")
    tavily_key = os.getenv("TAVILY_API_KEY")

    llm = ChatOpenAI(
        model=model_name,
        temperature=temperature,
        api_key=openai_key,
    )

    vectorstore, chunk_count = create_vectorstore(openai_key)

    @tool
    def rag_search_tool(query: str) -> str:
        """Search legal case documents using RAG for similar precedents and relevant information."""
        docs = vectorstore.similarity_search(query, k=3)

        if not docs:
            return "No relevant legal precedents found in the knowledge base."

        results = []
        for i, doc in enumerate(docs, 1):
            results.append(f"Precedent {i}:\n{doc.page_content}")

        return "\n\n".join(results)

    @tool
    def tavily_web_search_tool(query: str) -> str:
        """Search the web using Tavily for recent legal cases and legal developments."""
        if not tavily_key:
            return (
                "Tavily web search is not configured. "
                "Add TAVILY_API_KEY in Hugging Face Space secrets to enable live web search."
            )

        if TavilySearchResults is None:
            return "Tavily package is not installed. Add tavily-python to requirements.txt."

        tavily_tool = TavilySearchResults(max_results=3)
        results = tavily_tool.invoke({"query": query})

        if not results:
            return "No web search results found."

        formatted_results = []
        for i, result in enumerate(results, 1):
            content = result.get("content", "No content available")
            url = result.get("url", "No URL")
            formatted_results.append(f"Result {i}:\n{content}\nSource: {url}")

        return "\n\n".join(formatted_results)

    @tool
    def analyze_loopholes_tool(case_summary: str) -> str:
        """Analyze legal case to identify potential loopholes using LLM reasoning."""
        analysis_prompt = f'''
You are a senior legal analyst. Analyze the following case and identify potential loopholes.

Case Summary:
{case_summary}

Identify:
1. Procedural loopholes
2. Evidentiary gaps
3. Jurisdictional issues
4. Statute of limitations concerns
5. Constitutional challenges

Provide a detailed analysis with reasoning.
'''
        response = llm.invoke([HumanMessage(content=analysis_prompt)])
        return response.content

    tools = [rag_search_tool, tavily_web_search_tool, analyze_loopholes_tool]

    def agent_node(state: AgentState) -> Dict[str, Any]:
        system_prompt = SystemMessage(
            content=f'''
You are an expert legal analysis AI agent specializing in identifying loopholes and weaknesses in legal cases.

Case to Analyze:
{state["case_text"]}

Your Task:
Conduct a comprehensive legal analysis by:
1. Searching for similar precedent cases using rag_search_tool
2. Finding recent legal developments using tavily_web_search_tool
3. Analyzing loopholes using analyze_loopholes_tool

Analysis Framework:
- Force majeure defenses and procedural requirements
- Notice requirements and timelines
- Burden of proof and evidence standards
- Contributory negligence considerations
- Damages calculation and substantiation

After gathering all information, provide a comprehensive final analysis covering:
- Case summary
- Relevant precedents found
- Current legal developments, if available
- Procedural loopholes
- Evidentiary weaknesses
- Risk assessment
- Strategic recommendations for plaintiff and defendant

Use the tools before concluding.
'''
        )

        llm_with_tools = llm.bind_tools(tools)
        response = llm_with_tools.invoke([system_prompt] + state["messages"])
        return {"messages": [response]}

    def should_continue(state: AgentState) -> str:
        last_message = state["messages"][-1]

        if isinstance(last_message, AIMessage) and getattr(last_message, "tool_calls", None):
            return "tools"

        return "end"

    workflow = StateGraph(AgentState)
    workflow.add_node("agent", agent_node)
    workflow.add_node("tools", ToolNode(tools))

    workflow.set_entry_point("agent")
    workflow.add_conditional_edges(
        "agent",
        should_continue,
        {
            "tools": "tools",
            "end": END,
        },
    )
    workflow.add_edge("tools", "agent")

    memory = MemorySaver()
    graph = workflow.compile(checkpointer=memory)

    return graph, chunk_count


def run_legal_analysis(case_text: str, model_name: str, temperature: float):
    graph, chunk_count = build_legal_agent(model_name, temperature)

    thread_id = f"legal_case_{uuid.uuid4().hex[:8]}"
    config = {"configurable": {"thread_id": thread_id}}

    initial_state = {
        "messages": [HumanMessage(content="Analyze this legal case for loopholes and weaknesses.")],
        "case_text": case_text,
        "final_analysis": "",
    }

    trace = []
    final_answer = ""

    for step_output in graph.stream(initial_state, config, stream_mode="values"):
        last_message = step_output["messages"][-1]

        if isinstance(last_message, AIMessage):
            tool_calls = getattr(last_message, "tool_calls", None)

            if tool_calls:
                for call in tool_calls:
                    trace.append(
                        {
                            "type": "tool_call",
                            "name": call.get("name", "Unknown Tool"),
                            "args": call.get("args", {}),
                        }
                    )
            elif last_message.content:
                trace.append(
                    {
                        "type": "agent_message",
                        "content": last_message.content,
                    }
                )
                final_answer = last_message.content

        elif isinstance(last_message, ToolMessage):
            trace.append(
                {
                    "type": "tool_result",
                    "name": getattr(last_message, "name", "Tool"),
                    "content": last_message.content,
                }
            )

    final_state = graph.get_state(config)
    final_message = final_state.values["messages"][-1]
    if isinstance(final_message, AIMessage) and final_message.content:
        final_answer = final_message.content

    return {
        "final_answer": final_answer,
        "trace": trace,
        "chunk_count": chunk_count,
        "thread_id": thread_id,
    }


def display_trace(trace: List[Dict[str, Any]]):
    if not trace:
        st.info("No trace available yet.")
        return

    for i, item in enumerate(trace, 1):
        if item["type"] == "tool_call":
            with st.expander(f"Step {i}: Agent called tool → {item['name']}", expanded=True):
                st.json(item["args"])

        elif item["type"] == "tool_result":
            with st.expander(f"Step {i}: Tool result ← {item['name']}", expanded=False):
                st.write(item["content"])

        elif item["type"] == "agent_message":
            with st.expander(f"Step {i}: Agent response", expanded=False):
                st.write(item["content"])


st.title("⚖️ Legal Agentic AI Project")
st.caption("LangGraph ReAct Agent + RAG + Web Search + LLM Reasoning | Hugging Face Spaces Deployment")

with st.sidebar:
    st.header("Configuration")

    model_name = st.selectbox(
        "OpenAI Model",
        ["gpt-4o-mini", "gpt-4o", "gpt-4.1-mini"],
        index=0,
    )

    temperature = st.slider(
        "Temperature",
        min_value=0.0,
        max_value=1.0,
        value=0.0,
        step=0.1,
    )

    st.divider()

    st.subheader("Secrets Status")
    if os.getenv("OPENAI_API_KEY"):
        st.success("OPENAI_API_KEY found")
    else:
        st.error("OPENAI_API_KEY missing")

    if os.getenv("TAVILY_API_KEY"):
        st.success("TAVILY_API_KEY found")
    else:
        st.warning("TAVILY_API_KEY missing - web search will be skipped")

    st.divider()

    st.markdown(
        """
        **Agent Tools**
        1. RAG Search Tool  
        2. Tavily Web Search Tool  
        3. Loophole Analysis Tool  
        """
    )


tab_demo, tab_architecture, tab_code, tab_deployment = st.tabs(
    ["🚀 Live Demo", "🧠 Architecture", "🧩 Code Walkthrough", "☁️ HF Deployment"]
)


with tab_demo:
    st.subheader("Run the Legal Agent")

    st.markdown(
        """
        This demo takes a legal case, searches an internal legal knowledge base,
        optionally searches the web using Tavily, analyzes loopholes, and produces a final legal analysis.
        """
    )

    case_text = st.text_area(
        "Legal Case Input",
        value=SAMPLE_CASE,
        height=360,
    )

    col1, col2 = st.columns([1, 2])

    with col1:
        run_button = st.button("Run Legal Agent", type="primary", use_container_width=True)

    with col2:
        st.info("For learners: watch the trace below to understand Observe → Reason → Tool Use → Final Answer.")

    if run_button:
        if not secrets_ready():
            st.error("OPENAI_API_KEY is missing. Add it in Hugging Face Space → Settings → Repository secrets.")
        elif not case_text.strip():
            st.error("Please enter a legal case.")
        else:
            with st.spinner("Agent is analyzing the case..."):
                try:
                    result = run_legal_analysis(case_text, model_name, temperature)
                    st.session_state["last_result"] = result
                except Exception as e:
                    st.exception(e)

    if "last_result" in st.session_state:
        result = st.session_state["last_result"]

        st.success(f"Analysis complete. RAG initialized with {result['chunk_count']} document chunks.")

        st.markdown("## Final Legal Analysis")
        st.write(result["final_answer"])

        st.markdown("## Agent Execution Trace")
        display_trace(result["trace"])


with tab_architecture:
    st.subheader("How the Agent Works")

    st.markdown(
        """
        The notebook project has been converted into a deployable Streamlit app.

        The core design is a **ReAct Agent** built using **LangGraph**.

        **Flow:**

        ```text
        User Case
            ↓
        Agent Node
            ↓
        Decide: Need tools?
            ↓
        Tool Node
            ├── RAG Search Tool
            ├── Tavily Web Search Tool
            └── Loophole Analysis Tool
            ↓
        Agent Node again
            ↓
        Final Legal Analysis
        ```
        """
    )

    st.markdown("### What each tool does")

    st.table(
        [
            {
                "Tool": "rag_search_tool",
                "Purpose": "Searches internal legal knowledge base for precedents.",
                "Example": "Force majeure notice requirement",
            },
            {
                "Tool": "tavily_web_search_tool",
                "Purpose": "Searches current web information using Tavily.",
                "Example": "Recent force majeure legal developments",
            },
            {
                "Tool": "analyze_loopholes_tool",
                "Purpose": "Uses LLM reasoning to identify loopholes and weaknesses.",
                "Example": "Procedural gaps, evidentiary issues, damages weakness",
            },
        ]
    )

    st.markdown("### Teaching point")
    st.info(
        "This is not just a chatbot. It is an agent because it can reason, decide which tool to use, call tools, observe results, and then continue reasoning."
    )


with tab_code:
    st.subheader("Code Walkthrough for Learners")

    st.markdown("### 1. State")
    st.code(
        '''
class AgentState(TypedDict):
    messages: Annotated[list, add_messages]
    case_text: str
    final_analysis: str
''',
        language="python",
    )

    st.markdown("### 2. Tools")
    st.code(
        '''
tools = [
    rag_search_tool,
    tavily_web_search_tool,
    analyze_loopholes_tool
]
''',
        language="python",
    )

    st.markdown("### 3. Agent Node")
    st.code(
        '''
def agent_node(state):
    system_prompt = SystemMessage(content="Legal analysis instructions...")
    llm_with_tools = llm.bind_tools(tools)
    response = llm_with_tools.invoke([system_prompt] + state["messages"])
    return {"messages": [response]}
''',
        language="python",
    )

    st.markdown("### 4. Conditional Routing")
    st.code(
        '''
def should_continue(state):
    last_message = state["messages"][-1]

    if last_message.tool_calls:
        return "tools"

    return "end"
''',
        language="python",
    )

    st.markdown("### 5. LangGraph Workflow")
    st.code(
        '''
workflow = StateGraph(AgentState)
workflow.add_node("agent", agent_node)
workflow.add_node("tools", ToolNode(tools))

workflow.set_entry_point("agent")
workflow.add_conditional_edges("agent", should_continue, {
    "tools": "tools",
    "end": END
})
workflow.add_edge("tools", "agent")

graph = workflow.compile(checkpointer=MemorySaver())
''',
        language="python",
    )


with tab_deployment:
    st.subheader("Hugging Face Spaces Deployment")

    st.markdown(
        """
        Your Space should have this structure:

        ```text
        Hands-On-Agentic-Project/
        │
        ├── app.py
        ├── Dockerfile
        ├── requirements.txt
        ├── README.md
        │
        └── src/
            └── streamlit_app.py
        ```

        **Secrets required:**

        ```text
        OPENAI_API_KEY
        TAVILY_API_KEY
        ```

        `OPENAI_API_KEY` is required.  
        `TAVILY_API_KEY` is optional but recommended for the web-search tool.

        After uploading this file inside the `src` folder, run:

        **Factory rebuild** in Hugging Face Spaces.
        """
    )

    st.warning(
        "This app is for educational demonstration. It is not legal advice and should not be used as a substitute for a qualified legal professional."
    )