import os
from typing import List, Dict, Any, Optional, Callable
from dataclasses import dataclass, field
import json
from datetime import datetime
import inspect
import typing
from typing import Union
from swarms import Agent
from swarm_models import OpenAIChat


@dataclass
class ToolDefinition:
    name: str
    description: str
    parameters: Dict[str, Any]
    required_params: List[str]
    callable: Optional[Callable] = None


def extract_type_hints(func: Callable) -> Dict[str, Any]:
    """Extract parameter types from function type hints."""
    return typing.get_type_hints(func)


def extract_tool_info(func: Callable) -> ToolDefinition:
    """Extract tool information from a callable function."""
    # Get function name
    name = func.__name__

    # Get docstring
    description = inspect.getdoc(func) or "No description available"

    # Get parameters and their types
    signature = inspect.signature(func)
    type_hints = extract_type_hints(func)

    parameters = {}
    required_params = []

    for param_name, param in signature.parameters.items():
        # Skip self parameter for methods
        if param_name == "self":
            continue

        param_type = type_hints.get(param_name, Any)

        # Handle optional parameters
        is_optional = (
            param.default != inspect.Parameter.empty
            or getattr(param_type, "__origin__", None) is Union
            and type(None) in param_type.__args__
        )

        if not is_optional:
            required_params.append(param_name)

        parameters[param_name] = {
            "type": str(param_type),
            "default": (
                None
                if param.default is inspect.Parameter.empty
                else param.default
            ),
            "required": not is_optional,
        }

    return ToolDefinition(
        name=name,
        description=description,
        parameters=parameters,
        required_params=required_params,
        callable=func,
    )


@dataclass
class FunctionSpec:
    """Specification for a callable tool function."""

    name: str
    description: str
    parameters: Dict[
        str, dict
    ]  # Contains type and description for each parameter
    return_type: str
    return_description: str


@dataclass
class ExecutionStep:
    """Represents a single step in the execution plan."""

    step_id: int
    function_name: str
    parameters: Dict[str, Any]
    expected_output: str
    completed: bool = False
    result: Any = None


@dataclass
class ExecutionContext:
    """Maintains state during execution."""

    task: str
    steps: List[ExecutionStep] = field(default_factory=list)
    results: Dict[int, Any] = field(default_factory=dict)
    current_step: int = 0
    history: List[Dict[str, Any]] = field(default_factory=list)


class ToolAgent:
    def __init__(
        self,
        functions: List[Callable],
        openai_api_key: str,
        model_name: str = "gpt-4",
        temperature: float = 0.1,
    ):
        self.functions = {func.__name__: func for func in functions}
        self.function_specs = self._analyze_functions(functions)

        self.model = OpenAIChat(
            openai_api_key=openai_api_key,
            model_name=model_name,
            temperature=temperature,
        )

        self.system_prompt = self._create_system_prompt()
        self.agent = Agent(
            agent_name="Tool-Agent",
            system_prompt=self.system_prompt,
            llm=self.model,
            max_loops=1,
            verbose=True,
        )

    def _analyze_functions(
        self, functions: List[Callable]
    ) -> Dict[str, FunctionSpec]:
        """Analyze functions to create detailed specifications."""
        specs = {}
        for func in functions:
            hints = get_type_hints(func)
            sig = inspect.signature(func)
            doc = inspect.getdoc(func) or ""

            # Parse docstring for parameter descriptions
            param_descriptions = {}
            current_param = None
            for line in doc.split("\n"):
                if ":param" in line:
                    param_name = (
                        line.split(":param")[1].split(":")[0].strip()
                    )
                    desc = line.split(":", 2)[-1].strip()
                    param_descriptions[param_name] = desc
                elif ":return:" in line:
                    return_desc = line.split(":return:")[1].strip()

            # Build parameter specifications
            parameters = {}
            for name, param in sig.parameters.items():
                param_type = hints.get(name, Any)
                parameters[name] = {
                    "type": str(param_type),
                    "type_class": param_type,
                    "description": param_descriptions.get(name, ""),
                    "required": param.default == param.empty,
                }

            specs[func.__name__] = FunctionSpec(
                name=func.__name__,
                description=doc.split("\n")[0],
                parameters=parameters,
                return_type=str(hints.get("return", Any)),
                return_description=(
                    return_desc if "return_desc" in locals() else ""
                ),
            )

        return specs

    def _create_system_prompt(self) -> str:
        """Create system prompt with detailed function specifications."""
        functions_desc = []
        for spec in self.function_specs.values():
            params_desc = []
            for name, details in spec.parameters.items():
                params_desc.append(
                    f"    - {name}: {details['type']} - {details['description']}"
                )

            functions_desc.append(
                f"""
Function: {spec.name}
Description: {spec.description}
Parameters:
{chr(10).join(params_desc)}
Returns: {spec.return_type} - {spec.return_description}
            """
            )

        return f"""You are an AI agent that creates and executes plans using available functions.

Available Functions:
{chr(10).join(functions_desc)}

You must respond in two formats depending on the phase:

1. Planning Phase:
{{
    "phase": "planning",
    "plan": {{
        "description": "Overall plan description",
        "steps": [
            {{
                "step_id": 1,
                "function": "function_name",
                "parameters": {{
                    "param1": "value1",
                    "param2": "value2"
                }},
                "purpose": "Why this step is needed"
            }}
        ]
    }}
}}

2. Execution Phase:
{{
    "phase": "execution",
    "analysis": "Analysis of current result",
    "next_action": {{
        "type": "continue|request_input|complete",
        "reason": "Why this action was chosen",
        "needed_input": {{}} # If requesting input
    }}
}}

Always:
- Use exact function names
- Ensure parameter types match specifications
- Provide clear reasoning for each decision
"""

    def _execute_function(
        self, spec: FunctionSpec, parameters: Dict[str, Any]
    ) -> Any:
        """Execute a function with type checking."""
        converted_params = {}
        for name, value in parameters.items():
            param_spec = spec.parameters[name]
            try:
                # Convert value to required type
                param_type = param_spec["type_class"]
                if param_type in (int, float, str, bool):
                    converted_params[name] = param_type(value)
                else:
                    converted_params[name] = value
            except (ValueError, TypeError) as e:
                raise ValueError(
                    f"Parameter '{name}' conversion failed: {str(e)}"
                )

        return self.functions[spec.name](**converted_params)

    def run(self, task: str) -> Dict[str, Any]:
        """Execute task with planning and step-by-step execution."""
        context = ExecutionContext(task=task)
        execution_log = {
            "task": task,
            "start_time": datetime.utcnow().isoformat(),
            "steps": [],
            "final_result": None,
        }

        try:
            # Planning phase
            plan_prompt = f"Create a plan to: {task}"
            plan_response = self.agent.run(plan_prompt)
            plan_data = json.loads(
                plan_response.replace("System:", "").strip()
            )

            # Convert plan to execution steps
            for step in plan_data["plan"]["steps"]:
                context.steps.append(
                    ExecutionStep(
                        step_id=step["step_id"],
                        function_name=step["function"],
                        parameters=step["parameters"],
                        expected_output=step["purpose"],
                    )
                )

            # Execution phase
            while context.current_step < len(context.steps):
                step = context.steps[context.current_step]
                print(
                    f"\nExecuting step {step.step_id}: {step.function_name}"
                )

                try:
                    # Execute function
                    spec = self.function_specs[step.function_name]
                    result = self._execute_function(
                        spec, step.parameters
                    )
                    context.results[step.step_id] = result
                    step.completed = True
                    step.result = result

                    # Get agent's analysis
                    analysis_prompt = f"""
                    Step {step.step_id} completed:
                    Function: {step.function_name}
                    Result: {json.dumps(result)}
                    Remaining steps: {len(context.steps) - context.current_step - 1}
                    
                    Analyze the result and decide next action.
                    """

                    analysis_response = self.agent.run(
                        analysis_prompt
                    )
                    analysis_data = json.loads(
                        analysis_response.replace(
                            "System:", ""
                        ).strip()
                    )

                    execution_log["steps"].append(
                        {
                            "step_id": step.step_id,
                            "function": step.function_name,
                            "parameters": step.parameters,
                            "result": result,
                            "analysis": analysis_data,
                        }
                    )

                    if (
                        analysis_data["next_action"]["type"]
                        == "complete"
                    ):
                        if (
                            context.current_step
                            < len(context.steps) - 1
                        ):
                            continue
                        break

                    context.current_step += 1

                except Exception as e:
                    print(f"Error in step {step.step_id}: {str(e)}")
                    execution_log["steps"].append(
                        {
                            "step_id": step.step_id,
                            "function": step.function_name,
                            "parameters": step.parameters,
                            "error": str(e),
                        }
                    )
                    raise

            # Final analysis
            final_prompt = f"""
            Task completed. Results:
            {json.dumps(context.results, indent=2)}
            
            Provide final analysis and recommendations.
            """

            final_analysis = self.agent.run(final_prompt)
            execution_log["final_result"] = {
                "success": True,
                "results": context.results,
                "analysis": json.loads(
                    final_analysis.replace("System:", "").strip()
                ),
            }

        except Exception as e:
            execution_log["final_result"] = {
                "success": False,
                "error": str(e),
            }

        execution_log["end_time"] = datetime.utcnow().isoformat()
        return execution_log


def calculate_investment_return(
    principal: float, rate: float, years: int
) -> float:
    """Calculate investment return with compound interest.

    :param principal: Initial investment amount in dollars
    :param rate: Annual interest rate as decimal (e.g., 0.07 for 7%)
    :param years: Number of years to invest
    :return: Final investment value
    """
    return principal * (1 + rate) ** years


agent = ToolAgent(
    functions=[calculate_investment_return],
    openai_api_key=os.getenv("OPENAI_API_KEY"),
)

result = agent.run(
    "Calculate returns for $10000 invested at 7% for 10 years"
)