import os
import time
import psycopg2
from opentelemetry import trace, metrics
from opentelemetry.sdk.resources import Resource
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import BatchSpanProcessor
from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter
from opentelemetry.instrumentation.psycopg2 import Psycopg2Instrumentor
from opentelemetry.sdk.metrics import MeterProvider
from opentelemetry.sdk.metrics.export import PeriodicExportingMetricReader
from opentelemetry.exporter.otlp.proto.grpc.metric_exporter import OTLPMetricExporter

# Resource definition for observability metadata
resource = Resource.create({"service.name": "python-postgresql-app"})

# Tracing configuration
tracer_provider = TracerProvider(resource=resource)
trace.set_tracer_provider(tracer_provider)
span_processor = BatchSpanProcessor(OTLPSpanExporter(
    endpoint=os.getenv("OTEL_EXPORTER_OTLP_ENDPOINT", "http://jaeger:4317"),
    insecure=True
))
tracer_provider.add_span_processor(span_processor)

# Metrics configuration
metric_reader = PeriodicExportingMetricReader(OTLPMetricExporter(
    endpoint=os.getenv("OTEL_EXPORTER_OTLP_ENDPOINT", "http://otel-collector:4317"),
    insecure=True
))
meter_provider = MeterProvider(resource=resource, metric_readers=[metric_reader])
metrics.set_meter_provider(meter_provider)

# Get tracer and meter
tracer = trace.get_tracer(__name__)
meter = metrics.get_meter(__name__)

# Define a counter metric
db_operations_total = meter.create_counter(
    name="db.operations.total",
    description="Total number of database operations.",
    unit="{operations}"
)

# Instrument psycopg2 for auto-tracing DB queries
Psycopg2Instrumentor().instrument()

def get_db_connection():
    max_attempts = 5
    current_attempt = 0
    while current_attempt < max_attempts:
        try:
            conn = psycopg2.connect(
                host=os.getenv("DB_HOST", "postgres"),
                database=os.getenv("DB_NAME", "testdb"),
                user=os.getenv("DB_USER", "user"),
                password=os.getenv("DB_PASSWORD", "password")
            )
            print(f"Connected to PostgreSQL on attempt {current_attempt + 1}")
            return conn
        except psycopg2.OperationalError as e:
            print(f"Attempt {current_attempt + 1} failed: {e}")
            current_attempt += 1
            time.sleep(2)
    raise Exception("Failed to connect to PostgreSQL")


def create_table():
    with tracer.start_as_current_span("create_table"):
        conn = get_db_connection()
        cur = conn.cursor()
        cur.execute("DROP TABLE IF EXISTS users;")
        cur.execute("CREATE TABLE users (id SERIAL PRIMARY KEY, name VARCHAR(100), email VARCHAR(100));")
        conn.commit()
        cur.close()
        conn.close()
        db_operations_total.add(2, attributes={"operation": "create_table"})


def insert_user(name, email):
    with tracer.start_as_current_span("insert_user") as span:
        span.set_attribute("user.name", name)
        span.set_attribute("user.email", email)
        conn = get_db_connection()
        cur = conn.cursor()
        cur.execute("INSERT INTO users (name, email) VALUES (%s, %s);", (name, email))
        conn.commit()
        cur.close()
        conn.close()
        db_operations_total.add(1, attributes={"operation": "insert_user"})

if __name__ == "__main__":
    print("Starting OpenTelemetry PostgreSQL Demo...")
    create_table()
    insert_user("Alice", "[email protected]")
    insert_user("Bob", "[email protected]")
    get_users()
    update_user_email("Alice", "[email protected]")
    get_users()
    delete_user("Bob")
    get_users()
    print("Check Jaeger (http://localhost:16686) for traces and Prometheus (http://localhost:9090) for metrics.")