import duckdb
import polars as pl
from contextlib import contextmanager


@contextmanager
def get_connection(db_path: str = "output/pipeline.duckdb"):
    con = duckdb.connect(db_path)
    try:
        yield con
    finally:
        con.close()


def init_tables(con: duckdb.DuckDBPyConnection):
    con.execute("""
    CREATE TABLE IF NOT EXISTS prices_raw (
        timestamp BIGINT NOT NULL UNIQUE,
        price DOUBLE NOT NULL
    )
    """)

    con.execute("""
    CREATE TABLE IF NOT EXISTS prices_clean (
        timestamp TIMESTAMP NOT NULL UNIQUE,
        price DOUBLE NOT NULL
    )
    """)

    con.execute("""
    CREATE TABLE IF NOT EXISTS weather_raw (
        timestamp TIMESTAMP NOT NULL UNIQUE,
        temperature DOUBLE,
        wind_speed DOUBLE,
        solar DOUBLE,
        sunshine DOUBLE,
        cloud_cover DOUBLE,
        precipitation DOUBLE
    )
    """)

    con.execute("""
    CREATE TABLE IF NOT EXISTS combined (
        timestamp TIMESTAMP NOT NULL UNIQUE,
        price DOUBLE NOT NULL,
        temperature DOUBLE,
        wind_speed DOUBLE,
        solar DOUBLE,
        sunshine DOUBLE,
        cloud_cover DOUBLE,
        precipitation DOUBLE
    )
    """)


def upsert_prices(con: duckdb.DuckDBPyConnection, df: pl.DataFrame):
    con.execute(
        """INSERT INTO prices_raw SELECT * FROM df ON CONFLICT (timestamp) DO NOTHING;"""
    )


def upsert_prices_clean(con: duckdb.DuckDBPyConnection, df: pl.DataFrame):
    con.execute(
        """INSERT INTO prices_clean SELECT * FROM df ON CONFLICT (timestamp) DO NOTHING;"""
    )


def upsert_weather(con: duckdb.DuckDBPyConnection, df: pl.DataFrame):
    con.execute(
        """INSERT INTO weather_raw SELECT * FROM df ON CONFLICT (timestamp) DO NOTHING;"""
    )


def upsert_combined(con: duckdb.DuckDBPyConnection, df: pl.DataFrame):
    con.execute(
        """INSERT INTO combined SELECT * FROM df ON CONFLICT (timestamp) DO NOTHING;"""
    )