from mcp.server.fastmcp import FastMCP
import os
import psycopg2
from psycopg2.extras import RealDictCursor


financial_model_schema = {
  "general_ledger_master": {
    "description": "Represents metadata for each general ledger account.",
    "columns": {
      "id": "UUID - Primary key.",
      "general_ledger_code": "Numeric - Unique code representing the GL account.",
      "account_name": "String - Name of the account.",
      "category": "String - Category of the account (e.g., revenue, expense, asset).",
      "tags": "Array of String - Custom tags used for classification (e.g., 'cogs', 'cash', 'revenue').",
      "created_at": "DateTime - Timestamp when the entry was created.",
      "updated_at": "DateTime - Timestamp when the entry was last updated."
    }
  },
  "general_ledger": {
    "description": "Stores actual transactional entries for the general ledger.",
    "columns": {
      "id": "UUID - Primary key.",
      "date": "DateTime - Date of the transaction.",
      "general_ledger_code": "Numeric - Code referencing a GL master record.",
      "account_name": "String - Name of the account used in this entry.",
      "category": "String - Category type (used as fallback if GL Master mapping is missing).",
      "description": "String - Optional description or narration.",
      "debit": "Numeric - Debit amount.",
      "credit": "Numeric - Credit amount.",
      "is_active": "Boolean - Status flag to indicate if the record is active.",
      "entry_by": "String - User who made the entry.",
      "entry_system_source": "String - System or tool used to generate this entry.",
      "created_at": "DateTime - Timestamp when the record was inserted.",
      "updated_at": "DateTime - Timestamp of last update.",
      "integration_id": "UUID - Foreign key to IntegrationHistory table.",
      "integration_history": "Relationship - Backref to IntegrationHistory."
    }
  },
  
}



def test_run_query(query: str) -> list[dict]:
    """
    Executes a SQL query and returns the results as a list of JSON-like dicts.
    Only intended for read-only (SELECT) queries.
    """
    dbname = "postgres"
    user = "adminfpna"
    password =  "Maven%404321" 
    host = "fpna-dev-postgresql.postgres.database.azure.com"
    port = 5432

    conn = psycopg2.connect(
        dbname=dbname,
        user=user,
        password=password,
        host=host,
        port=port
    )
    try:
        with conn.cursor(cursor_factory=RealDictCursor) as cur:
            cur.execute(query)
            if cur.description:
                return cur.fetchall()
            return []
    finally:
        conn.close()


query = """
SELECT SUM(gl.credit) AS total_revenue
FROM general_ledger gl
JOIN general_ledger_master glm
  ON gl.general_ledger_code = glm.general_ledger_code
WHERE (glm.category = 'revenue' OR glm.tags @> ARRAY['revenue'])
  AND gl.is_active = TRUE
"""
print("--------------------------------------------------")
print("Result: ",test_run_query(query))
print("---------------------------------------------------")


mcp = FastMCP(
    name="Demo-Server",
    stateless_http=True,
    host="0.0.0.0",  # Accept connections from any network interface
    port=4000        # Specify the desired port
)





@mcp.tool(description="Run a read-only SQL query on PostgreSQL")
def run_query(query: str) -> list[dict]:
    """
    Executes a SQL query and returns the results as a list of JSON-like dicts.
    Only intended for read-only (SELECT) queries.
    """
    dbname = "postgres"
    user = "adminfpna"
    password =  "Maven%404321" 
    host = "fpna-dev-postgresql.postgres.database.azure.com"
    port = 5432

    conn = psycopg2.connect(
        dbname=dbname,
        user=user,
        password=password,
        host=host,
        port=port
    )
    try:
        with conn.cursor(cursor_factory=RealDictCursor) as cur:
            cur.execute(query)
            if cur.description:
                return cur.fetchall()
            return []
    finally:
        conn.close()

@mcp.prompt(description="query decompose")
def decompose_query(company_name: str) -> str:
    
    return f"""
        You are an expert in interpreting and simplifying ambiguous or complex queries, focusing on query decomposition and asking clarifying questions when necessary. Your goal is to break down user queries into clearly understandable components, identify missing context, and request clarification when the query is ambiguous or lacks necessary details.

        Instructions:

        1. Rewriting User Query:
        - Simplify and clarify the user’s query for better readability and understanding.
        - Retain the original structure and intent of the query.
        - Avoid unnecessary additions like “you want” or “you need.”

        2. Query Decomposition:
        - Break down the query into its core components.
        - Ensure that each part is concise and directly derived from the original query.

        3. Identify Missing Information:
        - If any part of the query is ambiguous, ask for clarification.
        - If essential details (e.g., time frame) are missing, ask the user for that context before proceeding.

        4. You must only answer questions that are related to the **allowed topics** for the **{company_name}**.
        -  Topics = ["financial statements", "revenue trends", "profit and loss", "balance sheet", "cash flow"]

        5. Response Format:
        Respond in the following JSON structure:

        {{
        "simplified_query": "",
        "query_components": {{
            "query_1": "",
            "query_2": ""
        }},
        "clarification_needed": "Yes" | "No",
        "clarification": [
            "Question 1",
            "Question 2"
        ]
        "allowed_topics : "Yes" | "No",
        "topics" : [
        "topic 1",
        "topic 2",
        .....
        ]
        }}
        """

@mcp.prompt(description="natural lang to sql query for general ledger")
def general_ledger(user_questions: dict) -> str:
    
    return f"""
    You are a PostgreSQL SQL query generator that uses **Chain-of-Thought (CoT)** reasoning to generate accurate, valid SQL queries from a natural language question.

    Your job is to deeply understand the questions, break it down, reason step by step, and generate SQL queries by using ONLY the provided schema and the metadata or Master_data.

    ---

    ### Important Rules (MUST follow strictly):
    1. **NO hallucinations**: Never invent table names, column names, or relationships not present in the schema.
    2. **Strict Schema Adherence**: Use only the tables, fields, and relationships explicitly defined in the provided schema.
    3. **Context Preservation**: Stay completely within the scope of the question and schema. Do not assume anything beyond what is provided.
    4. **SQL Syntax Validation**: Ensure the final SQL queries are free of syntax errors.
    5. **Simulate the Query (Mentally)**: Before outputting the final queries, simulate running them on the schema to confirm they produce correct results.
    6. **Format Clearly**: All queries must be combined in a single `query` string output, separated by semicolons.

    ---

    ###Step-by-Step Chain of Thought Instructions:

    **Step 1: Understand the user questions.**  
    - Break it into sub-questions if needed (e.g., Query 1, Query 2).

    **Step 2: Reason through each sub-question:**  
    - Identify: What data is needed? What are the metrics?
    - Find: Which tables contain the required data?
    - Join: Decide if any table joins are needed.
    - Filter: Apply conditions like year = 2023 if required.
    - Aggregate: Use SUM, COUNT, MAX etc. as needed.
    - Group: Group by month, product, department etc.
    - Sort and Limit: Use ORDER BY and LIMIT for top results.

    **Step 3: Write SQL queries:**  
    - Use correct and complete PostgreSQL syntax.
    - Validate fields, joins, and filters.
    - Recheck each query for errors.


    ---
    Schema:

    "general_ledger": {
    "description": "Stores  transactional entries for the general ledger.",
    "columns": {
      "id": "UUID - Primary key.",
      "date": "DateTime - Date of the transaction.",
      "general_ledger_code": "Numeric - Code referencing a GL master record.",
      "account_name": "String - Name of the account used in this entry.",
      "category": "String - Category type (used as fallback if GL Master mapping is missing).",
      "description": "String - Optional description or narration.",
      "debit": "Numeric - Debit amount.",
      "credit": "Numeric - Credit amount.",
      "is_active": "Boolean - Status flag to indicate if the record is active.",
      "entry_by": "String - User who made the entry.",
      "entry_system_source": "String - System or tool used to generate this entry.",
      "created_at": "DateTime - Timestamp when the record was inserted.",
      "updated_at": "DateTime - Timestamp of last update.",
      "integration_id": "UUID - Foreign key to IntegrationHistory table.",
      "integration_history": "Relationship - Backref to IntegrationHistory."
    }
    }
    
    metadata/ Master_data: 
    {[
    {
      "category": "asset",
      "general_ledger_code": 1100,
      "account_name": "accountsReceivable",
    },
    {
      "category": "asset",     
      "general_ledger_code": 1200,
      "account_name": "inventory",
    },
    {
      "category": "liability",
      "general_ledger_code": 2000,
      "account_name": "accountsPayable",
    },
    {
      "category": "liability",
      "general_ledger_code": 2100,
      "account_name": "accruedLiabilities",
    },
    {
      "category": "equity",
      "general_ledger_code": 3000,
      "account_name": "commonStock",
    },
    {
      "category": "equity",
      "general_ledger_code": 3100,
      "account_name": "retainedEarnings",
    },
    {
      "category": "revenue",
      "general_ledger_code": 4000,
      "account_name": "productRevenue",
    },
    {
      "category": "revenue",
      "general_ledger_code": 4010,
      "account_name": "serviceRevenue",
    },
    {
      "category": "expense",
      "general_ledger_code": 5000,
      "account_name": "cogsProduct",
    },
    {
      "category": "expense",
      "general_ledger_code": 5010,
      "account_name": "cogsService",
    },
    {
      "category": "expense",
      "general_ledger_code": 6000,
      "account_name": "salaries",
    },
    {
      "category": "expense",
      "general_ledger_code": 6100,
      "account_name": "marketing",
    },
    {
      "category": "expense",
      "general_ledger_code": 6200,
      "account_name": "rent",
    },
    {
      "category": "expense",
      "general_ledger_code": 6400,
      "account_name": "softwareSubscriptions",
    },
    {
      "category": "expense",
      "general_ledger_code": 6500,
      "account_name": "professionalServices",
    },
    {
      "category": "expense",
      "general_ledger_code": 6300,
      "account_name": "utilities",
    },
    {
      "category": "asset",
      "general_ledger_code": 1000,
      "account_name": "cash",
    }
  ],
    }

    ---

    question:{
        user_questions
    }

    **Step 4: Final Output:**  
    - Return the result as a single JSON object:
    ```json
    {{
    "query_1": ["sql_query": "", "question":" "],
    "query_2": ["sql_query": "", "question":" "]
    }}

    """


if __name__ == "__main__":
    mcp.run(
        transport="streamable-http")
    


