from langgraph.graph import StateGraph, END
from langchain_core.messages import HumanMessage, AIMessage
from langchain_core.tools import tool
from langchain_ollama import ChatOllama
import datetime
import re
from typing import TypedDict, List
import sys
import os

# -------------------------
# Configuration
# -------------------------

OLLAMA_MODEL = "qwen2.5:3b"

print("🔧 Initializing FlawHunt Local Agent with Ollama")
print(f"🤖 Model: {OLLAMA_MODEL}")
print("📝 Make sure Ollama is running: ollama serve")
print()

# -------------------------
# Basic Tools (that work with mistral)
# -------------------------

@tool
def time_tool() -> str:
    """Get current local time"""
    return datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")

@tool
def calc_tool(expression: str) -> str:
    """Evaluate a math expression"""
    if not re.match(r"^[0-9+\-*/(). ]+$", expression):
        return "Invalid expression"
    try:
        return str(eval(expression))
    except Exception as e:
        return f"Error: {str(e)}"

@tool  
def list_files_tool() -> str:
    """List all files in the current directory"""
    import subprocess
    try:
        result = subprocess.run(['ls', '-la'], capture_output=True, text=True, timeout=5)
        return result.stdout if result.stdout else "No files found"
    except Exception as e:
        return f"Error: {str(e)}"

@tool
def read_file_tool(filename: str) -> str:
    """Read contents of a file"""
    try:
        with open(filename, 'r') as f:
            content = f.read()
            # Limit to first 1000 chars
            if len(content) > 1000:
                return content[:1000] + "\n... (truncated)"
            return content
    except Exception as e:
        return f"Error reading {filename}: {str(e)}"

@tool
def search_web_tool(query: str) -> str:
    """Search the web using DuckDuckGo"""
    try:
        # Import here to avoid issues if not available
        sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
        from ai_terminal.tools import WebResearchTool
        
        web_tool = WebResearchTool()
        result = web_tool._run(query=query, mode='search', max_results=3)
        return result
    except Exception as e:
        return f"Web search error: {str(e)}"

tools = [time_tool, calc_tool, list_files_tool, read_file_tool, search_web_tool]

print("✅ Loaded tools:")
for tool in tools:
    print(f"   - {tool.name}: {tool.description}")
print()

# -------------------------
# LLM Setup - Dual Model System
# -------------------------

# Fast model for tool calling
TOOL_MODEL = "qwen2.5:3b"
# Better model for summarization
SUMMARY_MODEL = "mistral:7b"

print(f"🔧 Tool Model: {TOOL_MODEL}")
print(f"📝 Summary Model: {SUMMARY_MODEL}")
print()

try:
    # Main LLM for tool calling (fast)
    llm = ChatOllama(
        model=TOOL_MODEL,
        temperature=0.1,
        num_predict=256
    ).bind_tools(tools)
    
    # Separate LLM for summarization (better quality)
    summary_llm = ChatOllama(
        model=SUMMARY_MODEL,
        temperature=0.3,
        num_predict=512
    )
    
    print("✅ Both models initialized successfully!")
    print()
except Exception as e:
    print(f"❌ Failed to initialize Ollama: {e}")
    exit(1)

# -------------------------
# State
# -------------------------

class AgentState(TypedDict):
    messages: List

# -------------------------
# Nodes (EXACT same as working calc example)
# -------------------------

def llm_node(state: AgentState):
    response = llm.invoke(state["messages"])
    return {"messages": state["messages"] + [response]}

def tool_node(state: AgentState):
    last_msg = state["messages"][-1]
    tool_call = last_msg.tool_calls[0]
    
    tool_name = tool_call["name"]
    tool_args = tool_call.get("args", {})
    
    # Find and execute the tool
    tool_instance = next((t for t in tools if t.name == tool_name), None)
    
    if tool_instance:
        print(f"🔧 Executing: {tool_name}")
        try:
            result = tool_instance.invoke(tool_args)
        except Exception as e:
            result = f"❌ Error: {str(e)}"
    else:
        result = f"❌ Unknown tool: {tool_name}"
    
    return {
        "messages": state["messages"] + [
            AIMessage(content=str(result))
        ]
    }

def should_use_tool(state: AgentState):
    last_msg = state["messages"][-1]
    if getattr(last_msg, "tool_calls", None):
        return "tool"
    return END

# -------------------------
# Graph - Loop back to LLM after tool execution
# -------------------------

graph = StateGraph(AgentState)

graph.add_node("llm", llm_node)
graph.add_node("tool", tool_node)

graph.set_entry_point("llm")
graph.add_conditional_edges("llm", should_use_tool)
graph.add_edge("tool", "llm")  # Loop back to LLM to summarize tool results

agent = graph.compile()

# -------------------------
# Run loop
# -------------------------

if __name__ == "__main__":
    print("=" * 70)
    print("⚡ FlawHunt Local Agent Ready (Ollama)")
    print("=" * 70)
    print()
    print("🔧 Available Tools:")
    print("   • time_tool - Get current time")
    print("   • calc_tool - Calculate math expressions")
    print("   • list_files_tool - List files in current directory")
    print("   • read_file_tool - Read a file")
    print("   • search_web_tool - Search the web")
    print()
    print("💬 Examples:")
    print("   - 'what time is it?'")
    print("   - 'what is 2+2?'")
    print("   - 'list all files'")
    print("   - 'read README.md'")
    print("   - 'search for CVE 2024'")
    print()
    print("Type 'exit' to quit")
    print("=" * 70)
    print()

    while True:
        try:
            user_input = input("You: ")
            if user_input.lower() in ["exit", "quit", "q"]:
                print("👋 Goodbye!")
                break

            result = agent.invoke(
                {"messages": [HumanMessage(content=user_input)]},
                {"recursion_limit": 10}  # Prevent infinite loops
            )

            # Get all messages
            messages = result["messages"]
            
            # Find tool results (AIMessages that came from tools)
            tool_results = []
            llm_responses = []
            
            for i, msg in enumerate(messages):
                if isinstance(msg, AIMessage):
                    # Check if this is a tool result (comes after a tool call)
                    if i > 0 and hasattr(messages[i-1], 'tool_calls') and messages[i-1].tool_calls:
                        # This is a tool result - show it directly
                        if msg.content.strip():
                            tool_results.append(msg.content)
                    elif not hasattr(msg, 'tool_calls'):
                        # This is an LLM response
                        if msg.content.strip():
                            llm_responses.append(msg.content)
            
            # Display results
            if tool_results:
                print("\n" + "="*70)
                for result in tool_results:
                    print(result)
                print("="*70)
                
                # Generate summary using mistral:7b
                print("\n💭 Generating summary...")
                try:
                    # Determine the type of query for better summarization
                    user_query = user_input.lower()
                    
                    if any(word in user_query for word in ['cve', 'vulnerability', 'exploit', 'security']):
                        summary_prompt = f"""You are a cybersecurity analyst. Analyze the following search results and provide a concise summary focusing on:
1. Key vulnerabilities or threats mentioned
2. CVSS scores and severity levels
3. Affected systems or software
4. Recommended actions

Search Results:
{tool_results[0][:1500]}

Provide a 2-3 sentence technical summary highlighting the most critical information."""
                    
                    elif any(word in user_query for word in ['news', 'latest', 'recent']):
                        summary_prompt = f"""You are a cybersecurity news analyst. The user asked for latest news. Based on these search results, provide:
1. What are the actual news stories or topics (NOT just the website names)
2. Key takeaways or trends
3. Any urgent threats or important updates

Search Results:
{tool_results[0][:1500]}

Focus on the CONTENT and NEWS, not the sources. Be specific about what's happening in cybersecurity right now."""
                    
                    elif 'file' in user_query or 'list' in user_query:
                        summary_prompt = f"""Summarize the following file listing:
1. How many files/directories are there?
2. What types of files are present?
3. Any notable or important files?

{tool_results[0][:1500]}

Keep it brief and actionable."""
                    
                    else:
                        summary_prompt = f"""Provide a concise, actionable summary of the following information in 2-3 sentences. Focus on key facts and insights, not descriptions of sources:

{tool_results[0][:1500]}"""
                    
                    summary_response = summary_llm.invoke(summary_prompt)
                    summary = summary_response.content if hasattr(summary_response, 'content') else str(summary_response)
                    
                    if summary.strip():
                        print(f"\n📝 Summary: {summary}")
                except Exception as e:
                    print(f"⚠️  Summary generation failed: {e}")
            
            # Show LLM response if available (from qwen)
            elif llm_responses:
                print(f"\nAgent: {llm_responses[-1]}")
            else:
                # No tool results and no LLM response
                last_msg = messages[-1]
                if hasattr(last_msg, 'content') and last_msg.content.strip():
                    print(f"Agent: {last_msg.content}")
                else:
                    print("Agent: [No response generated]")
            
            print()
        except KeyboardInterrupt:
            print("\n👋 Goodbye!")
            break
        except Exception as e:
            print(f"❌ Error: {str(e)}")
            print()
