import os
import dspy
from dotenv import load_dotenv
from modaic import PrecompiledProgram, PrecompiledConfig
from src.neo4j import Neo4j

load_dotenv()

# set up Neo4j using NEO4J_URI
neo4j = Neo4j(
    uri=os.getenv("NEO4J_URI"),
    user=os.getenv("NEO4J_USER"),
    password=os.getenv("NEO4J_PASSWORD"),
)


class CypherFromText(dspy.Signature):
    """Task: Generate Cypher statement to query a graph database. 
    Instructions: Use only the provided relationship types and properties in the schema. 
    Do not use any other relationship types or properties that are not provided in the schema. 
    Do not include any explanations or apologies in your responses. 
    Do not respond to any questions that might ask anything else than for you to construct a Cypher statement. 
    Do not include any text except the generated Cypher statement.
    """
    question = dspy.InputField(
        desc="Question to model using a cypher statement."
    )
    neo4j_schema = dspy.InputField(
        desc="Current graph schema in Neo4j as a list of NODES and RELATIONSHIPS."
    )
    statement = dspy.OutputField(
        desc="Cypher statement to merge nodes and relationships found in the text."
    )


class GenerateCypherConfig(PrecompiledConfig):
    model: str = "openai/gpt-4o"
    max_tokens: int = 1024


class GenerateCypher(PrecompiledProgram):
    config: GenerateCypherConfig

    def __init__(self, config: GenerateCypherConfig, **kwargs):
        super().__init__(config=config, **kwargs)
        self.lm = dspy.LM(
            model=config.model,
            max_tokens=config.max_tokens,
        )
        self.generate_cypher = dspy.ChainOfThought(CypherFromText)
        self.generate_cypher.set_lm(self.lm)

    def forward(self, text: str, neo4j_schema: list[str]):
        return self.generate_cypher(text=text, neo4j_schema=neo4j_schema)


generate_cypher = GenerateCypher(GenerateCypherConfig())

if __name__ == "__main__":
    """
    from pathlib import Path
    import json

    examples_path = Path(__file__).parent / "examples" / "wikipedia-abstracts-v0_0_1.ndjson"
    with open(examples_path, "r") as f:
         for line in f:
             data = json.loads(line)
             text = data["text"]
             print("TEXT TO PROCESS:\n", text[:50])
             cypher = generate_cypher(text=text, neo4j_schema=neo4j.fmt_schema())
             neo4j.query(cypher.statement.replace('```', ''))
             print("CYPHER STATEMENT:\n", cypher.statement)
    
    schema = neo4j.fmt_schema()   
    print("SCHEMA:\n", schema)
    """
    generate_cypher.push_to_hub(
        "farouk1/text-to-cypher",
        with_code=True,
        tag="v0.0.8",
        commit_message="Update README.md",
    )