import dataclasses
import logging
from decimal import Decimal

import clients
import requests_cache
import secret_loader
from langchain_core.tools import tool

try:
    import pycountry.db
except ImportError:
    pycountry = None

logger = logging.getLogger(__name__)


@tool
def search(query: str):
    """Call to surf the web."""
    if 'sf' in query.lower() or 'san francisco' in query.lower():
        return "It's 60 degrees and foggy."
    return "It's 90 degrees and sunny."


def dataclasses_to_json(data):
    if pycountry and isinstance(data, pycountry.db.Country):
        return data.alpha_2
    if isinstance(data, list | tuple):
        return [dataclasses_to_json(d) for d in data]
    if isinstance(data, dict):
        return {k: dataclasses_to_json(v) for k, v in data.items() if v}
    if isinstance(data, Decimal):
        return float(data)
    if dataclasses.is_dataclass(data):
        result = {}
        for field in dataclasses.fields(data):
            result[field.name] = getattr(data, field.name)
        return dataclasses_to_json(result)
    return data


assert dataclasses_to_json([1, 2, 3]) == [1, 2, 3]


RETURN_FORMATS = {
    'raw_python': lambda x: x,
    'json': dataclasses_to_json,
}

RETURN_FORMAT = 'json'


def wrap_method(class_, method):
    logger.info('Wrapping %s.%s', class_.__name__, method.__name__)
    is_iterator = str(method.__annotations__.get('return', '')).startswith(
        'collections.abc.Iterator',
    )

    def wrapper(input_value):
        if isinstance(input_value, dict):
            logger.warning('Silently converting from dict to plain value!')
            input_value = next(input_value.values())
        logger.info(
            'AI called %s.%s(%s)', class_.__name__, method.__name__, repr(input_value),
        )
        try:
            result = method(input_value)
            if is_iterator:
                result = list(result)
            return RETURN_FORMATS[RETURN_FORMAT](result)
        except:
            logger.exception(
                'AI invocation of %s.%s(%s) failed!',
                class_.__name__,
                method.__name__,
                repr(input_value),
            )
            raise

    wrapper.__name__ = f'{class_.__name__}.{method.__name__}'
    wrapper.__doc__ = method.__doc__
    wrapper.__annotations__ = method.__annotations__
    return tool(wrapper)


def wrap_all_methods_on_client(obj):
    for field_name in dir(obj):
        if field_name.startswith('_'):
            continue
        method = getattr(obj, field_name)
        if not callable(method):
            continue
        if method.__doc__ is None:
            continue
        yield wrap_method(obj.__class__, method)


def get_tools():
    session = requests_cache.CachedSession('output/test.sqlite')
    session.headers['User-Agent'] = 'Test Test'

    secrets = secret_loader.SecretLoader()

    tools = []
    for client in clients.all_clients(session, secrets):
        tools.extend(wrap_all_methods_on_client(client))
    return tools