import re

import duckdb
import pandas as pd
import requests
from observability import span
from resolver import block
from sqlalchemy import create_engine, text
from trino.sqlalchemy import URL

from blocks.base import BaseBlock
from patterns import Secret

from ..patterns.data_operator import ExtendedDataFrame, ExtendedDataFrameList
from ..patterns.derived_str import JsonString
from ..patterns.http_header_private import HttpHeadersPrivate


@block(name="Extract_Regex_All", kind="tools")
class ExtractRegex(BaseBlock):
    """
    Extract all occurrences of a regex pattern in a text.

    Example:

    ```
    node = ExtractRegex(r"\d+")
    print(node("I have 3 apples and 5 bananas."))
    ```

    Will output:

    ["3", "5"]
    """

    def __init__(self, pattern: str, flags: int = 0):
        self.pattern = pattern
        self.flags = flags

    @span(name="ExtractRegex")
    def __call__(self, text: str) -> list:
        return re.findall(self.pattern, text, flags=self.flags)


@block(name="Trino_Query", kind="tools")
class TrinoQuery(BaseBlock):
    """Execute a Trino query."""

    def __init__(self, host: str, port: int = 8080):
        self.engine = create_engine(URL(host=host, port=port, user="admin"))

    @span(name="TrinoQuery")
    def __call__(self, query: str) -> ExtendedDataFrame:
        try:
            conn = self.engine.connect()
            df = pd.read_sql_query(text(query), conn)
            return ExtendedDataFrame(df)
        finally:
            conn.close()


@block(name="Batch_Trino_Query", kind="tools")
class BatchTrinoQuery(BaseBlock):
    """Execute a list of Trino queries in batch."""

    def __init__(self, host: str, port: int = 8080):
        self.trino = TrinoQuery(host, port)

    @span(name="BatchTrinoQueries")
    def __call__(self, queries: list) -> ExtendedDataFrameList:
        return ExtendedDataFrameList([self.trino(query) for query in queries])


@block(name="HTTP_API_Private", kind="tools")
class HttpApiPrivate(BaseBlock):
    """
    Executes an HTTP API request with the specified method and headers.
    Filter is a duckdb SQL statement and the table name is always `result`!

    Args:
        method (str): The HTTP method to use for the request (e.g. "get", "post", "put", "delete").
        headers (HttpHeaders): The headers to include in the request.

    Returns:
        ExtendedDataFrame: The response data from the HTTP API request.

    Raises:
        Exception: If the request fails with a non-200 status code.
    """

    def __init__(self, method: str, headers: HttpHeadersPrivate):
        assert method.lower() in ["get", "post", "put", "delete"]
        self.method = method.lower()
        self.headers = headers.headers

    def __call__(
        self, url: str, body: JsonString = "", filter: str = ""
    ) -> ExtendedDataFrame:
        if self.method in ["get", "delete"]:
            response = requests.request(self.method, url, headers=self.headers)
        else:
            response = requests.request(
                self.method, url, headers=self.headers, json=body.parsed()
            )
        if response.status_code != 200:
            raise Exception(
                f"Request failed with status code {response.status_code}, reason: {response.reason}"
            )
        result = pd.json_normalize(response.json())
        if len(filter) > 0:
            try:
                return ExtendedDataFrame(duckdb.sql(filter).df())
            except duckdb.BinderException as e:
                raise Exception(f"DuckDB encountered an error: {str(e)}")
        else:
            return ExtendedDataFrame(result)