from flask import Flask, request, jsonify, render_template_string
import psycopg2
import os

app = Flask(__name__)

def get_db_connection():
    return psycopg2.connect(
        host=os.environ.get("POSTGRES_HOST"),
        database=os.environ.get("POSTGRES_DB"),
        user=os.environ.get("POSTGRES_USER"),
        password=os.environ.get("POSTGRES_PASSWORD"),
        port=os.environ.get("POSTGRES_PORT", 5432),
    )

def init_db():
    conn = get_db_connection()
    cur = conn.cursor()
    cur.execute("""
        CREATE TABLE IF NOT EXISTS students (
            id SERIAL PRIMARY KEY,
            name TEXT NOT NULL,
            room_number TEXT NOT NULL,
            faculty TEXT NOT NULL
        );
    """)
    conn.commit()
    cur.close()
    conn.close()

@app.route("/")
def index():
    return "<h1>Welcome!</h1><p>Go to /add to add a student, /students to view all.</p>"

# HTML Form + POST Submission
@app.route("/add", methods=["GET", "POST"])
def add_student():
    if request.method == "POST":
        name = request.form.get("name")
        room_number = request.form.get("room_number")
        faculty = request.form.get("faculty")

        if not name or not room_number or not faculty:
            return "All fields are required!", 400

        conn = get_db_connection()
        cur = conn.cursor()
        cur.execute("INSERT INTO students (name, room_number, faculty) VALUES (%s, %s, %s)", (name, room_number, faculty))
        conn.commit()
        cur.close()
        conn.close()
        return "<p>Student added successfully!</p><a href='/add'>Add another</a>"

    # HTML form
    return render_template_string("""
        <h2>Add Student</h2>
        <form method="post">
            Name: <input name="name"><br>
            Room Number: <input name="room_number"><br>
            Faculty: <input name="faculty"><br>
            <input type="submit" value="Add Student">
        </form>
    """)

# View all students
@app.route("/students", methods=["GET"])
def get_students():
    conn = get_db_connection()
    cur = conn.cursor()
    cur.execute("SELECT id, name, room_number, faculty FROM students")
    rows = cur.fetchall()
    cur.close()
    conn.close()

    return jsonify([
        {"id": row[0], "name": row[1], "room_number": row[2], "faculty": row[3]} for row in rows
    ])

if __name__ == "__main__":
    init_db()
    app.run(host="0.0.0.0", port=8000)