from pyspark.sql import SparkSession
from pyspark.sql.types import StructType, StructField, StringType, IntegerType

spark = SparkSession.builder.appName("SparkSQL").getOrCreate()

# Sample data
data = [
    ("Alice", 30, "Engineering", 95000),
    ("Bob",   25, "Marketing",   72000),
    ("Carol", 35, "Engineering", 110000),
    ("Dave",  28, "Marketing",   68000),
    ("Eve",   40, "HR",           85000),
]

schema = StructType([
    StructField("name",       StringType()),
    StructField("age",        IntegerType()),
    StructField("department", StringType()),
    StructField("salary",     IntegerType()),
])

df = spark.createDataFrame(data, schema)

# Register as temporary view
df.createOrReplaceTempView("employees")

# Now query it with SQL
result = spark.sql("SELECT * FROM employees WHERE age > 28")
result.show()

# This will FAIL if "employees" view already exists
# df.createTempView("employees")  # ← AnalysisException!

# This is SAFE — replaces existing view silently
df.createOrReplaceTempView("employees")

# Check what views exist in the current session
spark.catalog.listTables()  # returns list of registered views

# Drop a temporary view
spark.catalog.dropTempView("employees")

# Check if a view exists
exists = spark.catalog.tableExists("employees")
print(exists)  # False after dropping

# Register as a GLOBAL temporary view
df.createOrReplaceGlobalTempView("employees_global")

# MUST prefix with 'global_temp.' when querying
result = spark.sql("SELECT * FROM global_temp.employees_global")
result.show()

# From a DIFFERENT SparkSession in the same application:
spark2 = SparkSession.builder.appName("SparkSQL").getOrCreate()
# This works! global_temp view is accessible across sessions
result2 = spark2.sql("SELECT * FROM global_temp.employees_global")
result2.show()

# Drop global temp view
spark.catalog.dropGlobalTempView("employees_global")

# List all tables (temp + global) in the catalog
for t in spark.catalog.listTables():
    print(t.name, t.tableType, t.isTemporary)

# Setup — reuse data from section 1
data = [
    ("Alice", 30, "Engineering", 95000),
    ("Bob",   25, "Marketing",   72000),
    ("Carol", 35, "Engineering", 110000),
    ("Dave",  28, "Marketing",   68000),
    ("Eve",   40, "HR",           85000),
]
df = spark.createDataFrame(data, ["name", "age", "department", "salary"])
df.createOrReplaceTempView("employees")

# Basic spark.sql() call
result = spark.sql("SELECT name, salary FROM employees")

# This is LAZY — nothing runs yet
# This triggers execution:
result.show()

# Multi-line SQL using triple quotes — clean and readable
result = spark.sql("""
    SELECT
        department,
        COUNT(*)        AS headcount,
        AVG(salary)     AS avg_salary,
        MAX(salary)     AS max_salary
    FROM employees
    GROUP BY department
    ORDER BY avg_salary DESC
""")

result.show()

from pyspark.sql.functions import col, round

# SQL result → chain DataFrame operations
result = (
    spark.sql("""
        SELECT department, AVG(salary) AS avg_salary
        FROM employees
        GROUP BY department
    """)
    .withColumn("avg_salary", round(col("avg_salary"), 0))  # DataFrame op
    .orderBy("avg_salary", ascending=False)             # DataFrame op
)

result.show()

# Dynamic SQL using Python f-strings
min_age = 28
dept = "Engineering"

query = f"""
    SELECT name, age, salary
    FROM employees
    WHERE age >= {min_age}
    AND department = '{dept}'
"""

result = spark.sql(query)
result.show()

# Show all tables in current catalog
spark.sql("SHOW TABLES").show()

# Describe the schema of a view
spark.sql("DESCRIBE employees").show()

# Create a table (Hive-style, requires metastore)
spark.sql("""
    CREATE TABLE IF NOT EXISTS prod_employees
    USING DELTA
    AS SELECT * FROM employees
""")

# DROP TABLE
spark.sql("DROP TABLE IF EXISTS prod_employees")

# Select specific columns
spark.sql("SELECT name, department FROM employees").show()

# Select with expressions and aliases
spark.sql("""
    SELECT
        name,
        salary,
        salary * 1.10            AS salary_with_raise,
        UPPER(department)        AS dept_upper,
        CONCAT(name, ' @ ', department) AS label
    FROM employees
""").show()

# Select ALL columns
spark.sql("SELECT * FROM employees").show()

# Select with DISTINCT
spark.sql("SELECT DISTINCT department FROM employees").show()

# Simple condition
spark.sql("SELECT * FROM employees WHERE salary > 80000").show()

# AND / OR conditions
spark.sql("""
    SELECT * FROM employees
    WHERE department = 'Engineering'
    AND salary > 90000
""").show()

# IN operator
spark.sql("""
    SELECT * FROM employees
    WHERE department IN ('Engineering', 'HR')
""").show()

# BETWEEN operator (inclusive on both ends)
spark.sql("""
    SELECT * FROM employees
    WHERE salary BETWEEN 70000 AND 100000
""").show()

# LIKE operator — % is wildcard, _ is single char
spark.sql("""
    SELECT * FROM employees
    WHERE name LIKE 'A%'
""").show()

# IS NULL / IS NOT NULL
spark.sql("""
    SELECT * FROM employees
    WHERE department IS NOT NULL
""").show()

# NOT IN
spark.sql("""
    SELECT * FROM employees
    WHERE department NOT IN ('HR')
""").show()

spark.sql("""
    SELECT
        name,
        salary,
        CASE
            WHEN salary >= 100000 THEN 'Senior'
            WHEN salary >= 80000  THEN 'Mid'
            ELSE 'Junior'
        END AS seniority_level
    FROM employees
    ORDER BY salary DESC
""").show()

# Only return top 3 rows
spark.sql("SELECT * FROM employees ORDER BY salary DESC LIMIT 3").show()

# Count employees per department
spark.sql("""
    SELECT department, COUNT(*) AS headcount
    FROM employees
    GROUP BY department
""").show()

# Multiple aggregations per group
spark.sql("""
    SELECT
        department,
        COUNT(*)            AS headcount,
        SUM(salary)         AS total_salary,
        AVG(salary)         AS avg_salary,
        MIN(salary)         AS min_salary,
        MAX(salary)         AS max_salary
    FROM employees
    GROUP BY department
    ORDER BY avg_salary DESC
""").show()

# Extended dataset for multi-column grouping example
data2 = [
    ("Alice", 30, "Engineering", "NYC", 95000),
    ("Bob",   25, "Marketing",   "NYC", 72000),
    ("Carol", 35, "Engineering", "SF",  110000),
    ("Dave",  28, "Marketing",   "SF",  68000),
    ("Eve",   40, "Engineering", "NYC", 88000),
]
df2 = spark.createDataFrame(data2, ["name","age","department","city","salary"])
df2.createOrReplaceTempView("emp2")

# Group by department AND city
spark.sql("""
    SELECT department, city, COUNT(*) AS count, AVG(salary) AS avg_sal
    FROM emp2
    GROUP BY department, city
    ORDER BY department, city
""").show()

# Only show departments with avg salary above 80000
spark.sql("""
    SELECT department, COUNT(*) AS headcount, AVG(salary) AS avg_salary
    FROM employees
    GROUP BY department
    HAVING AVG(salary) > 80000
""").show()

# HAVING with COUNT — departments with more than 1 employee
spark.sql("""
    SELECT department, COUNT(*) AS headcount
    FROM employees
    GROUP BY department
    HAVING COUNT(*) > 1
""").show()

# WHERE + GROUP BY + HAVING together
spark.sql("""
    SELECT department, AVG(salary) AS avg_salary
    FROM employees
    WHERE age >= 28          -- filter rows first
    GROUP BY department      -- then group
    HAVING AVG(salary) > 80000  -- then filter groups
""").show()

# COUNT(*) vs COUNT(col) — difference matters for NULLs
# COUNT(*) counts ALL rows including NULL
# COUNT(col) counts only NON-NULL values in that column

spark.sql("""
    SELECT
        COUNT(*)                  AS total_rows,
        COUNT(department)         AS non_null_dept,
        COUNT(DISTINCT department) AS unique_depts,
        COLLECT_LIST(name)        AS all_names,
        COLLECT_SET(department)   AS unique_dept_set
    FROM employees
""").show(truncate=False)

# Sort by salary descending (highest first)
spark.sql("""
    SELECT name, department, salary
    FROM employees
    ORDER BY salary DESC
""").show()

# Sort by department ASC, then salary DESC within each department
spark.sql("""
    SELECT name, department, salary
    FROM employees
    ORDER BY department ASC, salary DESC
""").show()

# Sort by a computed column (alias works in ORDER BY in Spark)
spark.sql("""
    SELECT name, salary, salary * 1.1 AS new_salary
    FROM employees
    ORDER BY new_salary DESC
""").show()

# NULL appears LAST by default in ASC
spark.sql("SELECT name, salary FROM employees ORDER BY salary ASC").show()

# Force NULLs to appear FIRST
spark.sql("""
    SELECT name, salary FROM employees
    ORDER BY salary ASC NULLS FIRST
""").show()

# Force NULLs to appear LAST in DESC (default is NULLS FIRST for DESC)
spark.sql("""
    SELECT name, salary FROM employees
    ORDER BY salary DESC NULLS LAST
""").show()

# Basic CTE — compute dept averages, then filter
spark.sql("""
    WITH dept_stats AS (
        SELECT
            department,
            AVG(salary) AS avg_salary,
            COUNT(*)    AS headcount
        FROM employees
        GROUP BY department
    )
    SELECT *
    FROM dept_stats
    WHERE avg_salary > 80000
    ORDER BY avg_salary DESC
""").show()

spark.sql("""
    WITH

    -- Step 1: Filter senior employees (age >= 30)
    senior_emp AS (
        SELECT * FROM employees WHERE age >= 30
    ),

    -- Step 2: Compute stats only for senior employees
    senior_stats AS (
        SELECT
            department,
            COUNT(*)    AS senior_count,
            AVG(salary) AS senior_avg_salary
        FROM senior_emp
        GROUP BY department
    ),

    -- Step 3: Mark high-paying departments
    final AS (
        SELECT
            department,
            senior_count,
            senior_avg_salary,
            CASE WHEN senior_avg_salary > 90000 THEN 'High Pay'
                 ELSE 'Standard'
            END AS pay_tier
        FROM senior_stats
    )

    SELECT * FROM final ORDER BY senior_avg_salary DESC
""").show()

# Scalar subquery in SELECT — compare each salary to overall average
spark.sql("""
    SELECT
        name,
        salary,
        (SELECT AVG(salary) FROM employees) AS company_avg,
        salary - (SELECT AVG(salary) FROM employees) AS diff_from_avg
    FROM employees
    ORDER BY diff_from_avg DESC
""").show()

# Create a "high pay departments" reference
spark.sql("""
    SELECT name, department, salary
    FROM employees
    WHERE department IN (
        SELECT department
        FROM employees
        GROUP BY department
        HAVING AVG(salary) > 80000
    )
""").show()

# Set up a "flagged_departments" table
flagged = spark.createDataFrame([("Engineering",), ("HR",)], ["dept"])
flagged.createOrReplaceTempView("flagged_depts")

# EXISTS — find employees in any flagged department
spark.sql("""
    SELECT e.name, e.department
    FROM employees e
    WHERE EXISTS (
        SELECT 1 FROM flagged_depts f
        WHERE f.dept = e.department   -- correlated: inner refs outer
    )
""").show()

spark.sql("""
    SELECT dept_summary.department, dept_summary.avg_salary
    FROM (
        SELECT department, AVG(salary) AS avg_salary
        FROM employees
        GROUP BY department
    ) AS dept_summary
    WHERE dept_summary.avg_salary > 80000
""").show()

-- General syntax:
-- function_name() OVER (
--     PARTITION BY partition_columns  -- divide into groups (optional)
--     ORDER BY order_columns          -- define row order within group
--     ROWS/RANGE BETWEEN ...          -- frame size (optional)
-- )

# Rank employees by salary within each department
spark.sql("""
    SELECT
        name,
        department,
        salary,
        RANK()        OVER (PARTITION BY department ORDER BY salary DESC) AS rank,
        DENSE_RANK()  OVER (PARTITION BY department ORDER BY salary DESC) AS dense_rank,
        ROW_NUMBER()  OVER (PARTITION BY department ORDER BY salary DESC) AS row_num
    FROM employees
""").show()

# Sales data with dates
sales_data = [
    ("2024-01-01", 1000),
    ("2024-01-02", 1500),
    ("2024-01-03", 1200),
    ("2024-01-04", 1800),
    ("2024-01-05", 900),
]
df_sales = spark.createDataFrame(sales_data, ["date", "revenue"])
df_sales.createOrReplaceTempView("daily_sales")

spark.sql("""
    SELECT
        date,
        revenue,
        LAG(revenue, 1)  OVER (ORDER BY date) AS prev_day_revenue,
        LEAD(revenue, 1) OVER (ORDER BY date) AS next_day_revenue,
        revenue - LAG(revenue, 1) OVER (ORDER BY date) AS day_over_day_change
    FROM daily_sales
    ORDER BY date
""").show()

spark.sql("""
    SELECT
        date,
        revenue,

        -- Running total (all rows from start to current)
        SUM(revenue) OVER (
            ORDER BY date
            ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW
        ) AS running_total,

        -- 3-day moving average (current + 2 previous)
        AVG(revenue) OVER (
            ORDER BY date
            ROWS BETWEEN 2 PRECEDING AND CURRENT ROW
        ) AS moving_avg_3d,

        -- Cumulative max
        MAX(revenue) OVER (
            ORDER BY date
            ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW
        ) AS running_max

    FROM daily_sales
    ORDER BY date
""").show()

spark.sql("""
    SELECT
        name,
        department,
        salary,
        FIRST_VALUE(name) OVER (PARTITION BY department ORDER BY salary DESC)
            AS highest_earner,
        LAST_VALUE(name)  OVER (
            PARTITION BY department ORDER BY salary DESC
            ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING
        ) AS lowest_earner
    FROM employees
""").show()