from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from datetime import datetime
from typing import List, Dict, Any, Optional
import json
import re

app = FastAPI(title="AI Agent with Anthropic/OpenAI SDK Support", version="2.0.0")

# Enable CORS
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Data Models for SDK Compatibility
class ToolDefinition(BaseModel):
    name: str
    description: str
    parameters: Dict[str, Any]

class AnthropicRequest(BaseModel):
    messages: List[Dict[str, Any]]
    tools: Optional[List[ToolDefinition]] = None
    model: Optional[str] = "claude-3-sonnet-20240229"
    max_tokens: Optional[int] = 1024
    temperature: Optional[float] = 0.7
    reasoning_split: Optional[bool] = False

class OpenAIRequest(BaseModel):
    messages: List[Dict[str, Any]]
    tools: Optional[List[ToolDefinition]] = None
    model: Optional[str] = "gpt-4"
    max_tokens: Optional[int] = 1024
    temperature: Optional[float] = 0.7
    reasoning_split: Optional[bool] = False

class AgenticRequest(BaseModel):
    task: str
    context: Optional[str] = ""
    conversation_history: Optional[List[Dict[str, Any]]] = []
    tools: Optional[List[ToolDefinition]] = None
    reasoning_split: Optional[bool] = False

# Available Tools
AVAILABLE_TOOLS = {
    "search_web": {
        "name": "search_web",
        "description": "Search the web for information based on a query",
        "parameters": {
            "type": "object",
            "properties": {
                "query": {"type": "string", "description": "Search query to look up information"}
            },
            "required": ["query"]
        }
    },
    "analyze_code": {
        "name": "analyze_code",
        "description": "Analyze code for understanding, structure, and suggestions",
        "parameters": {
            "type": "object",
            "properties": {
                "code": {"type": "string", "description": "Code to analyze"},
                "task": {"type": "string", "description": "Analysis task (comprehensive_analysis, optimization, debugging)"}
            },
            "required": ["code"]
        }
    },
    "get_weather": {
        "name": "get_weather",
        "description": "Get current weather information for a location",
        "parameters": {
            "type": "object",
            "properties": {
                "location": {"type": "string", "description": "Location to get weather for"},
                "units": {"type": "string", "enum": ["celsius", "fahrenheit"], "default": "celsius"}
            },
            "required": ["location"]
        }
    },
    "math_calc": {
        "name": "math_calc",
        "description": "Perform mathematical calculations",
        "parameters": {
            "type": "object",
            "properties": {
                "expression": {"type": "string", "description": "Mathematical expression to calculate"},
                "precision": {"type": "integer", "default": 6, "description": "Number of decimal places"}
            },
            "required": ["expression"]
        }
    },
    "text_analyzer": {
        "name": "text_analyzer",
        "description": "Analyze text for sentiment, keywords, and insights",
        "parameters": {
            "type": "object",
            "properties": {
                "text": {"type": "string", "description": "Text to analyze"},
                "analysis_type": {"type": "string", "enum": ["sentiment", "keywords", "summary"], "default": "sentiment"}
            },
            "required": ["text"]
        }
    }
}

def determine_tools_needed(task: str) -> List[Dict[str, Any]]:
    """Determine which tools are needed based on task analysis"""
    task_lower = task.lower()
    selected_tools = []
    
    # Search tools
    if any(keyword in task_lower for keyword in ["search", "find", "look up", "research", "information"]):
        selected_tools.append(AVAILABLE_TOOLS["search_web"])
    
    # Code analysis tools
    if any(keyword in task_lower for keyword in ["code", "function", "program", "debug", "optimize"]):
        selected_tools.append(AVAILABLE_TOOLS["analyze_code"])
    
    # Weather tools
    if any(keyword in task_lower for keyword in ["weather", "temperature", "climate"]):
        selected_tools.append(AVAILABLE_TOOLS["get_weather"])
    
    # Math tools
    if any(keyword in task_lower for keyword in ["calculate", "math", "+", "-", "*", "/", "="]) or re.search(r'\d+\s*[+\-*/=]\s*\d+', task_lower):
        selected_tools.append(AVAILABLE_TOOLS["math_calc"])
    
    # Text analysis tools
    if any(keyword in task_lower for keyword in ["analyze", "sentiment", "summary", "text"]):
        selected_tools.append(AVAILABLE_TOOLS["text_analyzer"])
    
    return selected_tools

def generate_tool_calls(task: str, context: str = "") -> List[Dict[str, Any]]:
    """Generate appropriate tool calls based on task"""
    tool_calls = []
    current_time = datetime.now().strftime('%Y%m%d_%H%M%S')
    
    if "search" in task.lower() or "find" in task.lower():
        query = task.replace("search for", "").replace("find", "").replace("look up", "").strip()
        tool_calls.append({
            "id": f"search_{current_time}",
            "function": {
                "name": "search_web",
                "arguments": json.dumps({"query": query or task})
            }
        })
    
    elif "code" in task.lower():
        tool_calls.append({
            "id": f"code_analysis_{current_time}",
            "function": {
                "name": "analyze_code",
                "arguments": json.dumps({
                    "code": context or "Sample code for analysis",
                    "task": "comprehensive_analysis"
                })
            }
        })
    
    elif "weather" in task.lower():
        location = context or "current location"
        tool_calls.append({
            "id": f"weather_{current_time}",
            "function": {
                "name": "get_weather",
                "arguments": json.dumps({"location": location})
            }
        })
    
    elif any(char.isdigit() for char in task) and any(op in task for op in ["+", "-", "*", "/"]):
        tool_calls.append({
            "id": f"math_{current_time}",
            "function": {
                "name": "math_calc",
                "arguments": json.dumps({"expression": task})
            }
        })
    
    elif "analyze" in task.lower():
        tool_calls.append({
            "id": f"text_analysis_{current_time}",
            "function": {
                "name": "text_analyzer",
                "arguments": json.dumps({
                    "text": context or task,
                    "analysis_type": "sentiment"
                })
            }
        })
    
    return tool_calls

def create_anthropic_response(task: str, reasoning_split: bool = False) -> Dict[str, Any]:
    """Create Anthropic SDK compatible response"""
    tool_calls = generate_tool_calls(task)
    reasoning = f"Analyzing task: '{task}'. I need to determine the best approach. "
    
    if tool_calls:
        reasoning += f"I will use the {tool_calls[0]['function']['name']} tool to help with this task."
    else:
        reasoning += "This appears to be a general reasoning task that doesn't require specific tools."
    
    content_blocks = []
    
    if reasoning_split:
        content_blocks.append({
            "type": "thinking",
            "text": reasoning
        })
    
    if tool_calls:
        content_blocks.append({
            "type": "tool_use",
            "id": tool_calls[0]["id"],
            "name": tool_calls[0]["function"]["name"],
            "input": json.loads(tool_calls[0]["function"]["arguments"])
        })
    
    content_blocks.append({
        "type": "text",
        "text": f"I've analyzed your task: '{task}'. I'm ready to proceed with the best approach."
    })
    
    return {
        "id": f"msg_{datetime.now().strftime('%Y%m%d_%H%M%S')}",
        "type": "message",
        "role": "assistant",
        "content": content_blocks,
        "model": "claude-3-sonnet-20240229",
        "stop_reason": "tool_use" if tool_calls else "end_turn"
    }

def create_openai_response(task: str, reasoning_split: bool = False) -> Dict[str, Any]:
    """Create OpenAI SDK compatible response"""
    tool_calls = generate_tool_calls(task)
    
    reasoning = f"I'm analyzing the task: '{task}'. "
    if tool_calls:
        reasoning += f"I need to use the {tool_calls[0]['function']['name']} tool to help complete this."
    else:
        reasoning += "This is a reasoning task that requires my cognitive abilities."
    
    return {
        "id": f"chatcmpl-{datetime.now().strftime('%Y%m%d_%H%M%S')}",
        "object": "chat.completion",
        "created": int(datetime.now().timestamp()),
        "model": "gpt-4",
        "choices": [
            {
                "index": 0,
                "message": {
                    "role": "assistant",
                    "content": f"Working on your task: '{task}'. Let me proceed with the analysis." if not reasoning_split else None,
                    "function_call": None,
                    "tool_calls": tool_calls if tool_calls else None
                },
                "finish_reason": "tool_calls" if tool_calls else "stop"
            }
        ],
        "usage": {
            "prompt_tokens": 10,
            "completion_tokens": 50,
            "total_tokens": 60
        },
        "reasoning_details": reasoning if reasoning_split else None
    }

# API Endpoints
@app.get("/")
async def root():
    return {
        "message": "AI Agent with Anthropic/OpenAI SDK Support",
        "version": "2.0.0",
        "status": "running",
        "timestamp": datetime.now().isoformat(),
        "endpoints": [
            "/",
            "/health", 
            "/agentic",
            "/anthropic/compatible",
            "/openai/compatible"
        ]
    }

@app.get("/health")
async def health():
    return {"status": "healthy", "timestamp": datetime.now().isoformat()}

@app.get("/models")
async def get_models():
    return {
        "models": [
            {
                "name": "anthropic-claude",
                "status": "loaded",
                "description": "Anthropic SDK compatible agentic model"
            },
            {
                "name": "openai-gpt4", 
                "status": "loaded",
                "description": "OpenAI SDK compatible agentic model"
            }
        ]
    }

@app.post("/agentic")
async def agentic_endpoint(request: AgenticRequest):
    """General agentic endpoint with interleaved thinking"""
    tools = request.tools or list(AVAILABLE_TOOLS.values())
    tool_calls = generate_tool_calls(request.task, request.context)
    
    reasoning = f"I need to analyze the task: '{request.task}'. "
    if tool_calls:
        reasoning += f"This task requires using the {tool_calls[0]['function']['name']} tool. "
    
    reasoning += "I'm reasoning through each step to determine the best approach for tool use."
    
    return {
        "reasoning_details": reasoning if request.reasoning_split else None,
        "content": f"Working on your task: '{request.task}'. I have the tools I need to assist.",
        "tool_calls": tool_calls,
        "status": "success",
        "thinking_process": "Interleaved thinking enabled - reasoning performed between tool interactions",
        "tools_available": len(tools),
        "conversation_length": len(request.conversation_history) if request.conversation_history else 0
    }

@app.post("/anthropic/compatible")
async def anthropic_endpoint(request: AnthropicRequest):
    """Anthropic SDK compatible endpoint with interleaved thinking"""
    # Get the latest user message
    latest_message = request.messages[-1] if request.messages else {"content": "Hello"}
    user_content = latest_message.get("content", "")
    
    if isinstance(user_content, list):
        user_text = " ".join([item.get("text", "") for item in user_content if item.get("type") == "text"])
    else:
        user_text = str(user_content)
    
    response = create_anthropic_response(user_text, request.reasoning_split)
    
    # Add available tools if specified
    if request.tools:
        available_tools = [tool for tool in AVAILABLE_TOOLS.values() if tool["name"] in [t.name for t in request.tools]]
    else:
        available_tools = list(AVAILABLE_TOOLS.values())
    
    return {
        **response,
        "available_tools": available_tools,
        "reasoning_split": request.reasoning_split
    }

@app.post("/openai/compatible")
async def openai_endpoint(request: OpenAIRequest):
    """OpenAI SDK compatible endpoint with interleaved thinking"""
    # Get the latest user message
    latest_message = request.messages[-1] if request.messages else {"content": "Hello"}
    user_content = latest_message.get("content", "")
    
    if isinstance(user_content, list):
        user_text = " ".join([item.get("text", "") for item in user_content if item.get("type") == "text"])
    else:
        user_text = str(user_content)
    
    response = create_openai_response(user_text, request.reasoning_split)
    
    # Add available tools if specified
    if request.tools:
        available_tools = [tool for tool in AVAILABLE_TOOLS.values() if tool["name"] in [t.name for t in request.tools]]
    else:
        available_tools = list(AVAILABLE_TOOLS.values())
    
    return {
        **response,
        "available_tools": available_tools,
        "reasoning_split": request.reasoning_split
    }

@app.get("/tools")
async def list_tools():
    """List all available tools"""
    return {
        "tools": list(AVAILABLE_TOOLS.values()),
        "count": len(AVAILABLE_TOOLS)
    }

@app.post("/tools/search")
async def simulate_tool_use(tool_name: str, arguments: Dict[str, Any]):
    """Simulate tool use for demonstration"""
    if tool_name not in AVAILABLE_TOOLS:
        raise HTTPException(status_code=400, detail=f"Tool '{tool_name}' not found")
    
    # Simulate tool responses
    if tool_name == "search_web":
        return {
            "tool": tool_name,
            "result": f"Search results for '{arguments.get('query', 'unknown')}': Found relevant information.",
            "status": "success"
        }
    elif tool_name == "analyze_code":
        return {
            "tool": tool_name,
            "result": "Code analysis complete: Found 2 functions, 1 class, good structure overall.",
            "status": "success"
        }
    elif tool_name == "get_weather":
        return {
            "tool": tool_name,
            "result": f"Weather for {arguments.get('location', 'unknown')}: 22°C, partly cloudy",
            "status": "success"
        }
    elif tool_name == "math_calc":
        try:
            expression = arguments.get("expression", "0")
            result = eval(expression)
            return {
                "tool": tool_name,
                "result": f"Calculation result: {result}",
                "status": "success"
            }
        except:
            return {
                "tool": tool_name,
                "result": "Error: Invalid mathematical expression",
                "status": "error"
            }
    elif tool_name == "text_analyzer":
        return {
            "tool": tool_name,
            "result": "Text analysis complete: Positive sentiment, found 3 key topics",
            "status": "success"
        }

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)