#!/usr/bin/env python3

import inspect
import json
from pathlib import Path
from typing import Any, Dict, List, Optional, Type

from pydantic import BaseModel

from crewai_tools import tools
from crewai.tools.base_tool import BaseTool, EnvVar

from pydantic.json_schema import GenerateJsonSchema
from pydantic_core import PydanticOmit


class SchemaGenerator(GenerateJsonSchema):
    def handle_invalid_for_json_schema(self, schema, error_info):
        raise PydanticOmit


class ToolSpecExtractor:
    def __init__(self) -> None:
        self.tools_spec: List[Dict[str, Any]] = []
        self.processed_tools: set[str] = set()

    def extract_all_tools(self) -> List[Dict[str, Any]]:
        for name in dir(tools):
            if name.endswith("Tool") and name not in self.processed_tools:
                obj = getattr(tools, name, None)
                if inspect.isclass(obj):
                    self.extract_tool_info(obj)
                    self.processed_tools.add(name)
        return self.tools_spec

    def extract_tool_info(self, tool_class: BaseTool) -> None:
        try:
            core_schema = tool_class.__pydantic_core_schema__
            if not core_schema:
                return

            schema = self._unwrap_schema(core_schema)
            fields = schema.get("schema", {}).get("fields", {})

            tool_info = {
                "name": tool_class.__name__,
                "humanized_name": self._extract_field_default(
                    fields.get("name"), fallback=tool_class.__name__
                ),
                "description": self._extract_field_default(
                    fields.get("description")
                ).strip(),
                "run_params_schema": self._extract_params(fields.get("args_schema")),
                "init_params_schema": self._extract_init_params(tool_class),
                "env_vars": self._extract_env_vars(fields.get("env_vars")),
                "package_dependencies": self._extract_field_default(
                    fields.get("package_dependencies"), fallback=[]
                ),
            }

            self.tools_spec.append(tool_info)

        except Exception as e:
            print(f"Error extracting {tool_class.__name__}: {e}")

    def _unwrap_schema(self, schema: Dict) -> Dict:
        while (
            schema.get("type") in {"function-after", "default"} and "schema" in schema
        ):
            schema = schema["schema"]
        return schema

    def _extract_field_default(self, field: Optional[Dict], fallback: str = "") -> str:
        if not field:
            return fallback

        schema = field.get("schema", {})
        default = schema.get("default")
        return default if isinstance(default, (list, str, int)) else fallback

    def _extract_params(
        self, args_schema_field: Optional[Dict]
    ) -> List[Dict[str, str]]:
        if not args_schema_field:
            return {}

        args_schema_class = args_schema_field.get("schema", {}).get("default")
        if not (
            inspect.isclass(args_schema_class)
            and hasattr(args_schema_class, "__pydantic_core_schema__")
        ):
            return {}

        try:
            return args_schema_class.model_json_schema(
                schema_generator=SchemaGenerator, mode="validation"
            )
        except Exception as e:
            print(f"Error extracting params from {args_schema_class}: {e}")
            return {}

    def _extract_env_vars(self, env_vars_field: Optional[Dict]) -> List[Dict[str, str]]:
        if not env_vars_field:
            return []

        env_vars = []
        for env_var in env_vars_field.get("schema", {}).get("default", []):
            if isinstance(env_var, EnvVar):
                env_vars.append(
                    {
                        "name": env_var.name,
                        "description": env_var.description,
                        "required": env_var.required,
                        "default": env_var.default,
                    }
                )
        return env_vars

    def _extract_init_params(self, tool_class: BaseTool) -> dict:
        ignored_init_params = [
            "name",
            "description",
            "env_vars",
            "args_schema",
            "description_updated",
            "cache_function",
            "result_as_answer",
            "max_usage_count",
            "current_usage_count",
            "package_dependencies",
        ]

        json_schema = tool_class.model_json_schema(
            schema_generator=SchemaGenerator, mode="serialization"
        )

        properties = {}
        for key, value in json_schema["properties"].items():
            if key not in ignored_init_params:
                properties[key] = value

        json_schema["properties"] = properties
        return json_schema

    def save_to_json(self, output_path: str) -> None:
        with open(output_path, "w", encoding="utf-8") as f:
            json.dump({"tools": self.tools_spec}, f, indent=2, sort_keys=True)
        print(f"Saved tool specs to {output_path}")


if __name__ == "__main__":
    output_file = Path(__file__).parent / "tool.specs.json"
    extractor = ToolSpecExtractor()

    specs = extractor.extract_all_tools()
    extractor.save_to_json(str(output_file))

    print(f"Extracted {len(specs)} tool classes.")