"""Example of using LLMs with PyOMOP

pip install pyomop[llm]

"""

from pyomop.cdm54 import Cohort, Base
from pyomop import CdmEngineFactory, CdmLLMQuery, CDMDatabase, CdmVector
import re
import os
from sqlalchemy import text
import datetime
import asyncio

# Import any LLMs that llama_index supports and you have access to
# Require OpenAI API key to use OpenAI LLMs
from llama_index.llms.google_genai import GoogleGenAI


async def main():
    # Create a sqllite database by default
    # You can also connect to an existing CDM database using the CdmEngineFactory
    cdm = CdmEngineFactory()
    # Postgres example below (db='mysql' also supported)
    # cdm = CdmEngineFactory(db='pgsql', host='', port=5432,
    #                       user='', pw='',
    #                       name='', schema='cdm6')

    engine = cdm.engine
    # Create Tables if required
    await cdm.init_models(Base.metadata)

    async with cdm.session() as session:  # type: ignore
        async with session.begin():

            # Adding  a cohort just for the example (not required if you have a OMOP CDM database)
            session.add(
                Cohort(
                    cohort_definition_id=2,
                    subject_id=100,
                    cohort_end_date=datetime.datetime.now(),
                    cohort_start_date=datetime.datetime.now(),
                )
            )
            await session.commit()

            # Use any LLM that llama_index supports
            llm = GoogleGenAI(
                model="gemini-2.0-flash",
                api_key=os.getenv("GOOGLE_GENAI_API_KEY"),
            )
            # Include tables that you want to query
            sql_database = CDMDatabase(
                engine,  # type: ignore
                include_tables=[  # type: ignore
                    "cohort",
                ],
                version="cdm54",
            )  # version can be 'cdm54' or 'cdm6'
            query_engine = CdmLLMQuery(sql_database, llm=llm).query_engine
            # Try any complex query.
            response = query_engine.query(
                "Show each in table cohort with a subject id of 100?"
            )
            print(f"Response: {response}\n")
            print(f"Response Metadata: {response.metadata}\n")
            # You may need to adjust the splitting based on the LLM output format. TODO: improve this
            sqls = response.metadata["sql_query"].split("\n")  # type: ignore

        # Execute the SQLs generated by the LLM using CdmVector for easy conversion to dataframe
        vec = CdmVector()
        for sql in sqls:
            try:
                if re.search(r"SELECT", sql):
                    result = await vec.execute(cdm, query=sql)
                    print("Executing custom query:")
                    df = vec.result_to_df(result)
                    print("DataFrame from result:")
                    print(df.head())
            except Exception as e:
                print(f"Error executing SQL: {sql}")
                print(e)

        # Or you can execute the SQLs directly using the session
        async with session.begin():
            # run each sql query
            for sql in sqls:
                try:
                    if re.search(r"SELECT", sql):
                        result = await session.execute(text(sql))
                        # print(result.all())
                        print(f"SQL: {sql}")
                        # Print the result in a table format
                        print("\n".join([str(row) for row in result.all()]))
                except Exception as e:
                    print(f"Error executing SQL: {sql}")
                    print(e)
    """

| cohort_id | subject_id | cohort_name |
|---|---|---|
| 1 | 100 | Math |

    """
    # Close session
    await session.close()
    await engine.dispose()  # type: ignore


# Run the main function
asyncio.run(main())
