import json
from pathlib import Path
from typing import AsyncGenerator
from langflow.api.v1.flows import get_session

from langflow.graph.graph.base import Graph
import pytest
from fastapi.testclient import TestClient
from httpx import AsyncClient
from sqlmodel import SQLModel, Session, create_engine
from sqlmodel.pool import StaticPool


def pytest_configure():
    pytest.BASIC_EXAMPLE_PATH = (
        Path(__file__).parent.absolute() / "data" / "basic_example.json"
    )
    pytest.COMPLEX_EXAMPLE_PATH = (
        Path(__file__).parent.absolute() / "data" / "complex_example.json"
    )
    pytest.OPENAPI_EXAMPLE_PATH = (
        Path(__file__).parent.absolute() / "data" / "Openapi.json"
    )

    pytest.CODE_WITH_SYNTAX_ERROR = """
def get_text():
    retun "Hello World"
    """


@pytest.fixture()
async def async_client() -> AsyncGenerator:
    from langflow.main import create_app

    app = create_app()
    async with AsyncClient(app=app, base_url="http://testserver") as client:
        yield client


# Create client fixture for FastAPI
@pytest.fixture(scope="module")
def client():
    from langflow.main import create_app

    app = create_app()

    with TestClient(app) as client:
        yield client


def get_graph(_type="basic"):
    """Get a graph from a json file"""

    if _type == "basic":
        path = pytest.BASIC_EXAMPLE_PATH
    elif _type == "complex":
        path = pytest.COMPLEX_EXAMPLE_PATH
    elif _type == "openapi":
        path = pytest.OPENAPI_EXAMPLE_PATH

    with open(path, "r") as f:
        flow_graph = json.load(f)
    data_graph = flow_graph["data"]
    nodes = data_graph["nodes"]
    edges = data_graph["edges"]
    return Graph(nodes, edges)


@pytest.fixture
def basic_graph_data():
    with open(pytest.BASIC_EXAMPLE_PATH, "r") as f:
        return json.load(f)


@pytest.fixture
def basic_graph():
    return get_graph()


@pytest.fixture
def complex_graph():
    return get_graph("complex")


@pytest.fixture
def openapi_graph():
    return get_graph("openapi")


@pytest.fixture
def json_flow():
    with open(pytest.BASIC_EXAMPLE_PATH, "r") as f:
        return f.read()


@pytest.fixture(name="session")  #
def session_fixture():  #
    engine = create_engine(
        "sqlite://", connect_args={"check_same_thread": False}, poolclass=StaticPool
    )
    SQLModel.metadata.create_all(engine)
    with Session(engine) as session:
        yield session


@pytest.fixture(name="client")  #
def client_fixture(session: Session):  #
    def get_session_override():  #
        return session

    from langflow.main import create_app

    app = create_app()

    app.dependency_overrides[get_session] = get_session_override  #

    yield TestClient(app)
    app.dependency_overrides.clear()  #


@pytest.fixture
def custom_chain():
    return '''
from __future__ import annotations
from typing import Any, Dict, List, Optional

from pydantic import Extra

from langchain.schema import BaseLanguageModel, Document
from langchain.callbacks.manager import (
    AsyncCallbackManagerForChainRun,
    CallbackManagerForChainRun,
)
from langchain.chains.base import Chain
from langchain.prompts import StringPromptTemplate
from langflow.interface.custom.base import CustomComponent

class MyCustomChain(Chain):
    """
    An example of a custom chain.
    """

from typing import Any, Dict, List, Optional

from pydantic import Extra

from langchain.schema import BaseLanguageModel, Document
from langchain.callbacks.manager import (
    AsyncCallbackManagerForChainRun,
    CallbackManagerForChainRun,
)
from langchain.chains.base import Chain
from langchain.prompts import StringPromptTemplate
from langflow.interface.custom.base import CustomComponent

class MyCustomChain(Chain):
    """
    An example of a custom chain.
    """

    prompt: StringPromptTemplate
    """Prompt object to use."""
    llm: BaseLanguageModel
    output_key: str = "text"  #: :meta private:

    class Config:
        """Configuration for this pydantic object."""

        extra = Extra.forbid
        arbitrary_types_allowed = True

    @property
    def input_keys(self) -> List[str]:
        """Will be whatever keys the prompt expects.

        :meta private:
        """
        return self.prompt.input_variables

    @property
    def output_keys(self) -> List[str]:
        """Will always return text key.

        :meta private:
        """
        return [self.output_key]

    def _call(
        self,
        inputs: Dict[str, Any],
        run_manager: Optional[CallbackManagerForChainRun] = None,
    ) -> Dict[str, str]:
        # Your custom chain logic goes here
        # This is just an example that mimics LLMChain
        prompt_value = self.prompt.format_prompt(**inputs)

        # Whenever you call a language model, or another chain, you should pass
        # a callback manager to it. This allows the inner run to be tracked by
        # any callbacks that are registered on the outer run.
        # You can always obtain a callback manager for this by calling
        # `run_manager.get_child()` as shown below.
        response = self.llm.generate_prompt(
            [prompt_value],
            callbacks=run_manager.get_child() if run_manager else None,
        )

        # If you want to log something about this run, you can do so by calling
        # methods on the `run_manager`, as shown below. This will trigger any
        # callbacks that are registered for that event.
        if run_manager:
            run_manager.on_text("Log something about this run")

        return {self.output_key: response.generations[0][0].text}

    async def _acall(
        self,
        inputs: Dict[str, Any],
        run_manager: Optional[AsyncCallbackManagerForChainRun] = None,
    ) -> Dict[str, str]:
        # Your custom chain logic goes here
        # This is just an example that mimics LLMChain
        prompt_value = self.prompt.format_prompt(**inputs)

        # Whenever you call a language model, or another chain, you should pass
        # a callback manager to it. This allows the inner run to be tracked by
        # any callbacks that are registered on the outer run.
        # You can always obtain a callback manager for this by calling
        # `run_manager.get_child()` as shown below.
        response = await self.llm.agenerate_prompt(
            [prompt_value],
            callbacks=run_manager.get_child() if run_manager else None,
        )

        # If you want to log something about this run, you can do so by calling
        # methods on the `run_manager`, as shown below. This will trigger any
        # callbacks that are registered for that event.
        if run_manager:
            await run_manager.on_text("Log something about this run")

        return {self.output_key: response.generations[0][0].text}

    @property
    def _chain_type(self) -> str:
        return "my_custom_chain"

class CustomChain(CustomComponent):
    display_name: str = "Custom Chain"
    field_config = {
        "prompt": {"field_type": "prompt"},
        "llm": {"field_type": "BaseLanguageModel"},
    }

    def build(self, prompt, llm, input: str) -> Document:
        chain = MyCustomChain(prompt=prompt, llm=llm)
        return chain(input)'''


@pytest.fixture
def data_processing():
    return """
import pandas as pd
from langchain.schema import Document
from langflow.interface.custom.base import CustomComponent

class CSVLoaderComponent(CustomComponent):
    display_name: str = "CSV Loader"
    field_config = {
        "filename": {"field_type": "str", "required": True},
        "column_name": {"field_type": "str", "required": True},
    }

    def build(self, filename: str, column_name: str) -> Document:
        # Load the CSV file
        df = pd.read_csv(filename)

        # Verify the column exists
        if column_name not in df.columns:
            raise ValueError(f"Column '{column_name}' not found in the CSV file")

        # Convert each row of the specified column to a document object
        documents = []
        for content in df[column_name]:
            metadata = {"filename": filename}
            documents.append(Document(page_content=str(content), metadata=metadata))

        return documents
"""


@pytest.fixture
def filter_docs():
    return """
from langchain.schema import Document
from langflow.interface.custom.base import CustomComponent
from typing import List

class DocumentFilterByLengthComponent(CustomComponent):
    display_name: str = "Document Filter By Length"
    field_config = {
        "documents": {"field_type": "Document", "required": True},
        "max_length": {"field_type": "int", "required": True},
    }

    def build(self, documents: List[Document], max_length: int) -> List[Document]:
        # Filter the documents by length
        filtered_documents = [doc for doc in documents if len(doc.page_content) <= max_length]

        return filtered_documents
"""


@pytest.fixture
def get_request():
    return """import requests
from typing import Dict, Union
from langchain.schema import Document
from langflow.interface.custom.base import CustomComponent

class GetRequestComponent(CustomComponent):
    display_name: str = "GET Request"
    field_config = {
        "url": {"field_type": "str", "required": True},
    }

    def build(self, url: str) -> Document:
        # Send a GET request to the URL
        response = requests.get(url)

        # Raise an exception if the request was not successful
        if response.status_code != 200:
            raise ValueError(f"GET request failed: {response.status_code} status code")

        # Create a document with the response text and the URL as metadata
        document = Document(page_content=response.text, metadata={"url": url})

        return document"""


@pytest.fixture
def post_request():
    return """import requests
from typing import Dict, Union
from langchain.schema import Document
from langflow.interface.custom.base import CustomComponent

class PostRequestComponent(CustomComponent):
    display_name: str = "POST Request"
    field_config = {
        "url": {"field_type": "str", "required": True},
        "data": {"field_type": "dict", "required": True},
    }

    def build(self, url: str, data: Dict[str, Union[str, int]]) -> Document:
        # Send a POST request to the URL
        response = requests.post(url, data=data)

        # Raise an exception if the request was not successful
        if response.status_code != 200:
            raise ValueError(f"POST request failed: {response.status_code} status code")

        # Create a document with the response text and the URL and data as metadata
        document = Document(page_content=response.text, metadata={"url": url, "data": data})

        return document"""