import os
import pdb
import tempfile
from typing import Literal, TypedDict, List, Optional
from uuid import uuid4
import abc
from dataclasses import dataclass

from langchain.chat_models import init_chat_model
from langchain_core.messages import AIMessage, SystemMessage, HumanMessage
from langchain_core.runnables import RunnableConfig
from langgraph.graph import END, START, MessagesState, StateGraph
from langgraph.checkpoint.memory import MemorySaver 
from langchain_core.runnables.graph_png import PngDrawer
from langchain_core.language_models import BaseChatModel
from pydantic import BaseModel, Field
from langchain.output_parsers import PydanticOutputParser
from langchain_core.tools import StructuredTool
from langgraph.types import Command
from langgraph.prebuilt import ToolNode
from tqdm import tqdm

from .python_toolkit import (
        PythonRunnerToolkit, DockerConfig, PythonRunnerToolContext,
        PythonRunnerTool
)
from .utils import show_image

class SQLAgentPolicy(abc.ABC):
    """
    This policy class allows an application to customize the SQLAgent for their
    use case. The SQLAgent allows an LLM to interact with a database that is
    exposed as a SQLAlchemy Engine object. The LLM code that interacts with the
    SQLAlchemy Engine object runs inside a Docker container.

    The Docker container has the following "work" folders:

    - `/home/pythonuser`: This is the working folder for any code executed by the LLM.
    - `/home/pythonuser/src`: Hosts any source files needed to expose the Engine object to
      the LLM. Typically, it contains a file (say, `database.py`) from which an
      LLM can import the engine object using a line such as `from src.database
      import engine`. Any additional files needed by the application (such as
      an sqlite database), can be copied to this folder by overriding the
      `files_to_copy` method.
    - `/home/pythonuser/artifacts`: This folder inside the container contains artifacts
      generated by the LLM code (such as PNG plots).

    At a minimum, a policy should override `engine_export_code` to define the
    source code that will allow an agent to import an engine object. For
    example, the following code will expose the `Chinook.db` sqlite data base
    as an SQLAlchemy Engine object called `engine` to any code that needs it.

    ```python
    import os
    from sqlalchemy import create_engine
    script_dir = os.path.dirname(os.path.abspath(__file__))
    db = os.path.join(script_dir, 'Chinook.db')
    engine = create_engine(f"sqlite:///{db}")
    __all__ = [ 'engine' ]
    ```

    The above, of course, assumes that the application has enabled the copying
    of `Chinook.db` file to the Docker container by overriding the
    `files_to_copy` method as follows:

    ```python
    def files_to_copy(tmpdir: str):
        return [ '/local/path/to/Chinook.db' ]

    Applications can also override `database_hints` method to return a string
    describing the database table entries (for example, what's the unit of the
    currency column?). If provided, the agent will provide this information to
    the LLM for its reference.
    """
    
    def files_to_copy(self) -> List[str]:
        """
        Specify a list of files to copy into the working directory of the
        Docker Container.
        """
        return []

    @abc.abstractmethod
    def engine_export_code(self) -> str:
        """
        Return the code that will expose an SQLAlchemy Engine object under the
        name `engine`. See the description of SQLAgentPolicy for more details.
        """

    def database_hints(self) -> Optional[str]:
        """
        A string that provides any hints that an application may want to
        provide the LLM about specific tables and columns. Don't include schema
        since that will be already provided to the LLM using the SQLAlchemy
        engine object. 

        For example, you can include some text to indicate that an "Invoice"
        column is in the units of dollars, or a "Time" column is in the units
        of "milliseconds".
        """
        return None

class SQLAgentState(MessagesState):
    first_time: bool

SQLA_LIST_TABLES_PYTHON=r"""
from src.database import engine
from sqlalchemy import MetaData

# Create a metadata object
metadata = MetaData()
metadata.reflect(bind=engine)

# Get the table names
table_names = metadata.tables.keys()
print('Available tables are:', list(table_names))
"""

SQLA_GET_SCHEMA_PYTHON=r"""
from src.database import engine
from sqlalchemy import MetaData
from sqlalchemy.schema import CreateTable

# Load metadata
metadata = MetaData()
metadata.reflect(bind=engine)

# Generate CREATE TABLE statements
for table_name, table in metadata.tables.items():
    create_statement = str(CreateTable(table).compile(engine))
    print(f"Schema for '{table_name} table':\n{create_statement}\n")

"""
def _run_python_tool(code: str, config:RunnableConfig):

    # create an artificial tool call and an AI message.
    tool_call = {
        "name": "run_python_tool",
        "args": {
            'python_code': code,
            'artifacts_abs_paths': [],
        },
        "id": uuid4().hex,
        "type": "tool_call",
    }
    tool_call_message = AIMessage(content="", tool_calls=[tool_call])

    # run the python tool on behalf of AI.
    config_idx: int = config['configurable'].get('config_idx', 0) # pyright: ignore[reportTypedDictNotRequiredAccess]
    python_runner = agent_globals_registry.get_globals(config_idx).agent.python_runner_tool
    tool_message = python_runner.invoke(tool_call, config=config)
    
    # return the request and response.
    return {"messages": [tool_call_message, tool_message]}

def list_tables(state: SQLAgentState, config: RunnableConfig):
    """
    Create a pre-determined tool call to list database table entries.
    """
    return _run_python_tool(SQLA_LIST_TABLES_PYTHON, config)

def call_get_schema(state: SQLAgentState, config: RunnableConfig):
    """
    Force the model to query the schema of the SQL database.
    """
    return _run_python_tool(SQLA_GET_SCHEMA_PYTHON, config)

class AnswerSchema(BaseModel):
    message: str = Field(description="""
    The answer to the user's question in markdown format. Include links to artifacts as appropriate.
    """)
    user_artifacts_abs_paths: List[str] = Field(description="""
    Local paths of artifacts (if any) produced as a part of the user's answer.
    """)

REACT_SYSTEM_PROMPT ="""
You are an agent designed to interact with a SQL database. You'll do so by
writing and executing python code using a python runner tool
("run_python_tool") that is equipped with Python 3.12+.

The database is exposed to you using an SQLAlchemy Engine object which you can
import as follows:

```
from src.database import engine
```

Always include the above line when referring to the databse. You can now use
the engine object imported above to list tables, schema etc., and also to run
queries.

The python tool returns the stdout of the program and hence you can use print
statements to get the information you want. 

Given an input question, create a syntactically correct code to run, then look
at the results of the run, and return the answer. You can also use the tool to
create artifiacts such as plots to show the user.

Here is an example of code that lists all the tables in the database:

```
{example_code}
```

You also have a available a `get_database_hints` tool that returns some hints
about the contents os the database.

DO NOT make any DML statements (INSERT, UPDATE, DELETE, DROP etc.) to the database.

{output_instructions}
"""

def react_node(state: SQLAgentState, config: RunnableConfig):
    """
    The main node in which we invoke the reasoning ability of the agent.
    """
    # get agent instance.
    config_idx: int = config['configurable'].get('config_idx', 0)  # pyright: ignore[reportTypedDictNotRequiredAccess]
    agent: SQLAgent = agent_globals_registry.get_globals(config_idx).agent
    # create a system prompt.
    system_message = SystemMessage(agent.system_prompt)
    # bind the python runner tool the the LLM.
    llm_with_tools = agent.llm.bind_tools([agent.python_runner_tool, get_database_hints])
    # invoke the LLM.
    response = llm_with_tools.invoke([system_message] + state["messages"])
    return {"messages": [response]}

def route_tool(state: SQLAgentState) -> Literal[END, 'run_python']:
    messages = state["messages"]
    last_message = messages[-1]
    assert isinstance(last_message, AIMessage)
    tool_calls = last_message.tool_calls
    if len(tool_calls) == 0:
        return END
    else:
        return 'run_python'

def should_get_table_meta(state: SQLAgentState, config: RunnableConfig) -> Command[Literal["react_node", "list_tables"]]:
    
    # if this is not the first time, then go directly to the react node and
    # don't update any state.
    if 'first_time' in state and state['first_time'] == 0:
        return Command(update={}, goto='react_node')
    # if this is the first time, go to list_tables and clear first_time.
    # - also, add runtime configuration to the state, so it available to all
    #   the tools.
    else:
        return Command(
                update = {
                    'first_time': 0,
                    }, 
                goto="list_tables"
        )


SQLITE_DATABASE_LOADER_CODE="""
import os
from sqlalchemy import create_engine
script_dir = os.path.dirname(os.path.abspath(__file__))
db = os.path.join(script_dir, '{db_filename}')
engine = create_engine("sqlite:///" + db)
__all__ = [ 'engine' ]
"""

class SQLiteAgentPolicy(SQLAgentPolicy):
    """
    A helper for SQLite databases that includes some default implementation
    code that pertains to SQLite databases.
    """

    def __init__(self, db_file: str):
        self.db_file = db_file

    def files_to_copy(self):
        return [ self.db_file ]

    def engine_export_code(self):
        return SQLITE_DATABASE_LOADER_CODE.format(
                db_filename = os.path.basename(self.db_file)
        )

def _get_database_hints(config: RunnableConfig) -> str:
    config_idx: int = config['configurable'].get('config_idx', 0)  # pyright: ignore[reportTypedDictNotRequiredAccess]
    agent_policy: SQLAgentPolicy = agent_globals_registry.get_globals(config_idx).agent.agent_policy
    hints = agent_policy.database_hints() 
    if hints is not None:
        return hints
    else:
        return "No hints are available."

GET_DATABASE_HINTS_DESC="""
Returns some hints about tables and columns in the database.

Returns:
    str: Hints about the tables and columns in a database.
"""
get_database_hints = StructuredTool.from_function(
        func=_get_database_hints,
        name="get_database_tool",
        description=GET_DATABASE_HINTS_DESC,
        return_direct=True
)

def database_hints_node(state: SQLAgentState, config: RunnableConfig):
    """
    Create a pre-determined tool call to provide hints about the database.
    """
    tool_call = {
        "name": "get_database_hints",
        "args": {},
        "id": uuid4().hex,
        "type": "tool_call",
    }
    tool_call_message = AIMessage(content="", tool_calls=[tool_call])
    tool_message = get_database_hints.invoke(tool_call, config=config)
    return {"messages": [tool_call_message, tool_message]}

class SQLAgentConfigSchema(TypedDict):
    config_idx: int  

class SQLAgent:

    def __init__(self, 
                 agent_policy:SQLAgentPolicy,
                 docker_config: DockerConfig,
                 tmpdir: str,
                 model: str | BaseChatModel = "openai:gpt-4.1",
                 show_artifacts: bool =False,
                 recursion_limit: int =25
    ):
        """
        Creates a new SQL Agent.

        Args:
            agent_policy: App specific customization for the agent.
            docker_config: Docker configuration for Python Tool.
            tmpdir: A temporary directory to use for the temporary files
                created during the agent run (recommend using
                tempfile.TemporaryDirectory() context manager for this).
            model: The model to use.
            recursion_limit: The max number of graph state transitions to
                execute.
        """

        # latch parameters.
        self.recursion_limit = recursion_limit

        # initialize LLM.
        if isinstance(model, str):
            self.llm = init_chat_model(model)
        else:
            self.llm = model

        # configure the python runner.
        self.python_tool = PythonRunnerToolContext(
                docker_config,
                tmpdir,
                ignore_dependencies=['src'],
                ignore_unsafe_functions=['compile']
        )
        self.python_toolkit = PythonRunnerToolkit.from_context(
                config=self.python_tool
        )
        # python toolkit only exposes a single tool
        self.python_runner_tool: PythonRunnerTool = next(
                tool 
                for tool in self.python_toolkit.get_tools() 
                if tool.name == "run_python_tool"
        ) # pyright: ignore[reportAttributeAccessIssue]

        # latch agent policy.
        self.agent_policy: SQLAgentPolicy = agent_policy

        # initialize the database configuration.
        # - copy any files the application needs.
        for file in self.agent_policy.files_to_copy():
            self.python_tool.copy_file_to_container(file, '/home/pythonuser/src')
        # -- copy the database loader code to container.
        self.python_tool.copy_code_to_container(
                self.agent_policy.engine_export_code(),
                '/home/pythonuser/src/database.py'
        )

        # graph building.
        self._build_graph()

        # initialize runtime configuration.
        self._init_config()

        # generate system prompt for react agent.
        self.show_artifacts = show_artifacts
        if show_artifacts:
            output_instructions = PydanticOutputParser(pydantic_object=AnswerSchema).get_format_instructions()
        else:
            output_instructions = "Return your answer with Markdown formatting."
        self.system_prompt = REACT_SYSTEM_PROMPT.format(
            example_code=SQLA_LIST_TABLES_PYTHON,
            output_instructions=output_instructions
        )

    def _set_globals(self) -> int:
        
        return agent_globals_registry.register_globals(self)

    def _build_graph(self):

        self.config_idx = self._set_globals()
        builder = StateGraph(SQLAgentState, config_schema=SQLAgentConfigSchema)

        # are hints available?
        hints_avail = self.agent_policy.database_hints() is not None

        # add nodes.
        builder.add_node(should_get_table_meta)
        builder.add_node(list_tables)
        builder.add_node(call_get_schema)
        builder.add_node(react_node)
        builder.add_node(ToolNode([self.python_runner_tool], name="run_python"))
        if hints_avail:
            builder.add_node(database_hints_node)

        # connect them.
        builder.add_edge(START, "should_get_table_meta")
        builder.add_edge("list_tables", "call_get_schema")
        if hints_avail:
            builder.add_edge("call_get_schema", "database_hints_node")
            builder.add_edge("database_hints_node", "react_node")
        else:
            builder.add_edge("call_get_schema", "react_node")
        builder.add_conditional_edges(
            "react_node",
            route_tool,
        )
        builder.add_edge("run_python", "react_node")

        self.memory = MemorySaver()
        self.agent = builder.compile(checkpointer=self.memory)

    def print_ascii(self):
        self.agent.get_graph().print_ascii()

    def display_graph(self):

        with tempfile.TemporaryDirectory() as tmpdir:
            drawer = PngDrawer()
            png_path = os.path.join(tmpdir, '_lang_graph.png')
            drawer.draw(self.agent.get_graph(), png_path)
            show_image(png_path)

    def display_artifacts(self, artifacts: List[str]):
        """
        Displays artifacts generated by the Model.
        """
        for artifact in artifacts:
            if not os.path.isfile(artifact):
                raise RuntimeError(f'Artifact {artifact} is not found!')
            _, ext = os.path.splitext(artifact)
            match ext:
                case '.png': show_image(artifact)

    def artifacts(self) -> List[str]:
        """
        Returns the list of artifacts from the last AI message.
        """
        final_state = self.agent.get_state(self.config) # pyright: ignore[reportArgumentType]
        final_message = final_state.values['messages'][-1]
        answer: AnswerSchema = AnswerSchema.model_validate_json(final_message.content)
        return answer.user_artifacts_abs_paths

    def _init_config(self):

        self.config = {
                'recursion_limit': self.recursion_limit,
                'configurable': {
                    'thread_id': 1,
                    'config_idx': self.config_idx
                }
        }

    def chat(self, question:str, quiet=True):

        gen = self.agent.stream(
            {"messages": [HumanMessage(question)]}, 
            stream_mode="values", 
            config=self.config # type: ignore
        )
        gen = tqdm(gen, desc="Thinking ...", total=self.recursion_limit) if quiet else gen
        for step in gen:
            if not quiet:
                step["messages"][-1].pretty_print()

        print('================================== Final Answer ================================== ')
        final_state = self.agent.get_state(self.config) # pyright: ignore[reportArgumentType]
        final_message = final_state.values['messages'][-1]
        if self.show_artifacts:
            answer: AnswerSchema = AnswerSchema.model_validate_json(final_message.content)
            print(answer.message)
            self.display_artifacts(answer.user_artifacts_abs_paths)
        else:
            print(final_message.content)

@dataclass
class SQLAgentGlobals:
    """
    A global variable for tools to access non-trivial agent related objects.
    The config object will just contain an index of the global object
    associated with the agent.
    """
    agent: SQLAgent

class SQLAgentGlobalsRegistry:
    """
    A registry to keep track of every SQLAgentGlobal object created.
    """
    def __init__(self):
        self.globals: List[SQLAgentGlobals] = []

    def register_globals(
            self,
            agent: SQLAgent
            ):
        idx = len(self.globals)
        self.globals.append(SQLAgentGlobals(agent))
        return idx

    def get_globals(self, index) -> SQLAgentGlobals:
        return self.globals[index]

# a registry to keep track of non-trivial objects needed by the nodes and
# tools.
agent_globals_registry = SQLAgentGlobalsRegistry()

