from sqlalchemy import create_engine, MetaData, Table, Column, Integer, String, Text, text
from sqlalchemy.exc import SQLAlchemyError

# Try to import real API utilities from the project
try:
    from pageindex_.utils import ChatGPT_API
except ImportError:
    # If not found, we will just log it or define a dummy if needed,
    # but the requirement is "No need HAS_REAL_API", so we won't rely on it for logic switching.
    print("Warning: pageindex.utils not found. ChatGPT_API will fail if called.")
    ChatGPT_API = None

class MetadataDB:
    """
    Handles Database operations using SQLAlchemy for abstraction.
    """
    def __init__(self, connection_string='sqlite:///:memory:'):
        """
        Initialize database connection.
        connection_string: SQLAlchemy connection string (e.g., sqlite:///example.db, postgresql://...)
        """
        self.engine = create_engine(connection_string, echo=False)
        self.metadata = MetaData()
        self.connection = self.engine.connect()
        self.tables = {}

    def create_table(self, table_config):
        """
        Creates a table based on the provided configuration using SQLAlchemy.
        
        table_config: dict containing 'table_name' and 'columns' list.
        """
        table_name = table_config['table_name']
        columns_config = table_config['columns']
        
        cols = []
        for col in columns_config:
            col_name = col['name']
            col_type_str = col['type'].upper()
            
            # Map string types to SQLAlchemy types
            if "INTEGER" in col_type_str:
                col_type = Integer
            elif "TEXT" in col_type_str:
                col_type = Text
            elif "STRING" in col_type_str or "VARCHAR" in col_type_str:
                col_type = String
            else:
                col_type = Text # Default fallback
            
            # Check for primary key
            is_pk = "PRIMARY KEY" in col_type_str
            
            cols.append(Column(col_name, col_type, primary_key=is_pk))
        
        # Define table
        table = Table(table_name, self.metadata, *cols)
        self.tables[table_name] = table
        
        # Create table in DB
        self.metadata.create_all(self.engine)
        print(f"✅ Table '{table_name}' created/verified via SQLAlchemy.")

    def insert_data(self, table_name, columns, data):
        """
        Inserts data into the specified table.
        
        columns: list of column names.
        data: list of tuples/lists matching the columns.
        """
        if table_name not in self.tables:
            print(f"❌ Table '{table_name}' not found in metadata.")
            return

        table = self.tables[table_name]
        
        # Convert list of tuples to list of dicts for SQLAlchemy
        records = []
        for row in data:
            record = {col: val for col, val in zip(columns, row)}
            records.append(record)
            
        try:
            self.connection.execute(table.insert(), records)
            self.connection.commit()
            print(f"✅ Inserted {len(data)} rows into '{table_name}'.")
        except SQLAlchemyError as e:
            print(f"❌ Insert Error: {e}")

    def execute_query(self, sql_query):
        """
        Executes a raw SQL query and returns results.
        """
        try:
            result = self.connection.execute(text(sql_query))
            return result.fetchall()
        except SQLAlchemyError as e:
            print(f"❌ SQL Execution Error: {e}")
            return None

    def close(self):
        self.connection.close()


class SQLGenerator:
    """
    Generates SQL queries from natural language using an LLM.
    """
    def __init__(self, table_config, model="qwen3-max"):
        self.table_name = table_config['table_name']
        self.schema_desc = self._generate_schema_description(table_config)
        self.model = model

    def _generate_schema_description(self, table_config):
        lines = [f"Table: {table_config['table_name']}", "Columns:"]
        for col in table_config['columns']:
            desc = col.get('desc', '')
            lines.append(f"- {col['name']} ({col['type']}): {desc}")
        return "\n".join(lines)

    def generate_sql(self, user_query):
        prompt = f"""
        You are a SQL expert. Convert the following natural language user query into a SQL query.
        
        Database Schema:
        {self.schema_desc}

        User Query: "{user_query}"

        Rules:
        1. Return ONLY the SQL query. No markdown, no explanations.
        2. The query should select relevant columns based on the request.
        3. Use standard SQL compatible with the underlying database.
        
        SQL Query:
        """
        
        print(f"\n🤖 Generating SQL for: '{user_query}'")
        
        if ChatGPT_API:
            try:
                sql_query = ChatGPT_API(model=self.model, prompt=prompt)
                # Clean up potential markdown formatting
                sql_query = sql_query.replace("```sql", "").replace("```", "").strip()
                return sql_query
            except Exception as e:
                print(f"Error calling LLM: {e}")
                return None
        else:
            print("❌ LLM API not available. Cannot generate SQL.")
            return None

if __name__ == "__main__":
    # ==========================================
    # Configuration
    # ==========================================

    DOCUMENTS_SCHEMA = {
        "table_name": "documents",
        "columns": [
            {"name": "id", "type": "INTEGER PRIMARY KEY AUTOINCREMENT", "desc": "Primary Key"},
            {"name": "doc_id", "type": "TEXT NOT NULL", "desc": "Unique identifier for the document in PageIndex"},
            {"name": "title", "type": "TEXT NOT NULL", "desc": "Document title"},
            {"name": "category", "type": "TEXT", "desc": "e.g., 'Finance', 'Legal', 'HR', 'Healthcare', 'ESG'"},
            {"name": "year", "type": "INTEGER", "desc": "e.g., 2022, 2023, 2024"},
            {"name": "region", "type": "TEXT", "desc": "e.g., 'North America', 'Europe', 'Asia', 'Global'"},
            {"name": "summary", "type": "TEXT", "desc": "Brief description of content"}
        ]
    }

    MOCK_DATA = [
        ("doc_001", "Q1 2024 Financial Report - TechCorp", "Finance", 2024, "North America", "Quarterly financial results showing 15% growth in cloud revenue."),
        ("doc_002", "2023 Annual Sustainability Report", "ESG", 2023, "Global", "Overview of carbon footprint reduction initiatives and social responsibility goals."),
        ("doc_003", "Legal Framework for AI in EU", "Legal", 2024, "Europe", "Analysis of the new EU AI Act and its compliance requirements for tech companies."),
        ("doc_004", "Healthcare Market Trends 2023", "Healthcare", 2023, "Asia", "Survey of emerging medical device markets in Southeast Asia."),
        ("doc_005", "TechCorp Employee Handbook", "HR", 2022, "North America", "Internal policies regarding remote work, benefits, and code of conduct."),
        ("doc_006", "Q4 2023 Financial Report - TechCorp", "Finance", 2023, "North America", "End of year financial summary with strong performance in hardware sector.")
    ]

    DATA_COLUMNS = ["doc_id", "title", "category", "year", "region", "summary"]

    # ==========================================
    # Main Execution
    # ==========================================

    def run_demo():
        # 1. Initialize DB (Can easily switch to 'postgresql://user:password@localhost/dbname')
        db = MetadataDB(connection_string='sqlite:///:memory:')
        
        # 2. Configure DB Structure
        db.create_table(DOCUMENTS_SCHEMA)
        
        # 3. Configure Data
        # Note: 'id' is auto-increment, so we don't include it in insert columns if the DB handles it.
        # We pass the columns that match our data.
        db.insert_data(DOCUMENTS_SCHEMA['table_name'], DATA_COLUMNS, MOCK_DATA)
        
        # 4. Initialize SQL Generator
        sql_gen = SQLGenerator(DOCUMENTS_SCHEMA)

        # 5. Run Test Queries
        test_queries = [
            "Find all financial reports from 2023.",
            "Show me documents related to Europe.",
            "What HR policies do we have?"
        ]

        for q in test_queries:
            # Note: Since we removed the simulation fallback, this will only work if LLM API is available.
            # If not available, it prints an error, which is consistent with "No need HAS_REAL_API" (don't fake it).
            sql = sql_gen.generate_sql(q)
            if sql:
                print(f"   -> Generated SQL: {sql}")
                results = db.execute_query(sql)
                
                if results is not None:
                    print(f"   -> Found {len(results)} documents:")
                    for row in results:
                        print(f"      - {row}")
            else:
                print("   -> Skipping execution (No SQL generated).")
            print("-" * 50)

        db.close()


    run_demo()
