app.py

import os
import sqlite3
from flask import Flask, render_template, request, redirect, url_for, session
from werkzeug.utils import secure_filename
from werkzeug.security import generate_password_hash, check_password_hash

from agents import AgentOrchestrator


def get_db_connection(db_name):
    conn = sqlite3.connect(db_name, timeout=10)
    conn.execute("PRAGMA busy_timeout = 10000")
    return conn

# --- Initialize DB and create users table with role ---
def init_db():
    with get_db_connection("users.db") as conn:
        cursor = conn.cursor()
        cursor.execute("""
            CREATE TABLE IF NOT EXISTS users (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                username TEXT UNIQUE NOT NULL,
                password TEXT NOT NULL,
                role TEXT NOT NULL DEFAULT 'user'
            )
        """)
        conn.commit()


def init_reports_db():
    with get_db_connection("users.db") as conn:
        cursor = conn.cursor()
        cursor.execute(
            """
            CREATE TABLE IF NOT EXISTS reports (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                username TEXT NOT NULL,
                description TEXT NOT NULL,
                attachment TEXT,
                status TEXT NOT NULL DEFAULT 'Open',
                confidence INTEGER,
                final_solution TEXT,
                created_at TEXT DEFAULT CURRENT_TIMESTAMP,
                updated_at TEXT DEFAULT CURRENT_TIMESTAMP
            )
            """
        )
        conn.commit()


def create_report_record(username, description, attachment):
    with get_db_connection("users.db") as conn:
        cursor = conn.cursor()
        cursor.execute(
            """
            INSERT INTO reports (username, description, attachment, status)
            VALUES (?, ?, ?, 'Open')
            """,
            (username, description, attachment),
        )
        conn.commit()
        return cursor.lastrowid


def update_report_analysis(report_id, final_solution, confidence):
    if not report_id:
        return
    with get_db_connection("users.db") as conn:
        cursor = conn.cursor()
        cursor.execute(
            """
            UPDATE reports
            SET final_solution = ?,
                confidence = ?,
                status = 'Analyzed',
                updated_at = CURRENT_TIMESTAMP
            WHERE id = ?
            """,
            (final_solution, confidence, report_id),
        )
        conn.commit()


def get_reports_for_user(username):
    with get_db_connection("users.db") as conn:
        cursor = conn.cursor()
        cursor.execute(
            """
            SELECT id, description, status, confidence, created_at, updated_at
            FROM reports
            WHERE username = ?
            ORDER BY id DESC
            """,
            (username,),
        )
        rows = cursor.fetchall()
    return rows

init_db()
init_reports_db()

# --- Flask Setup ---
app = Flask(__name__)
app.config["UPLOAD_FOLDER"] = os.path.join(os.path.dirname(__file__), "uploads")
os.makedirs(app.config["UPLOAD_FOLDER"], exist_ok=True)
app.secret_key = "dev-secret-change-me"
ALLOWED_EXTENSIONS = {"png", "jpg", "jpeg", "txt", "log"}
agentic = AgentOrchestrator()

# --- Helpers ---
def allowed_file(filename):
    return "." in filename and filename.rsplit(".", 1)[1].lower() in ALLOWED_EXTENSIONS

# --- Routes ---

# Home
@app.route("/")
def home():
    return render_template("home.html")


@app.route("/my_reports")
def my_reports():
    if "user" not in session:
        return redirect(url_for("home"))
    username = session.get("user")
    reports = get_reports_for_user(username)
    return render_template("my_reports.html", reports=reports, username=username)

# Register
@app.route("/register", methods=["GET", "POST"])
def register():
    if request.method == "POST":
        username = request.form.get("username", "").strip()
        password = request.form.get("password", "").strip()
        role = request.form.get("role", "user")
        if not username or not password:
            return "Enter both username and password."

        hashed_pw = generate_password_hash(password)
        try:
            with get_db_connection("users.db") as conn:
                cursor = conn.cursor()
                cursor.execute("INSERT INTO users (username, password, role) VALUES (?, ?, ?)", 
                               (username, hashed_pw, role))
                conn.commit()
            return redirect(url_for("login"))
        except sqlite3.IntegrityError:
            return "Username already exists."
    return render_template("register.html")

# User Login
@app.route("/login", methods=["GET", "POST"])
def login():
    if request.method == "POST":
        username = request.form.get("username", "").strip()
        password = request.form.get("password", "").strip()
        if not username or not password:
            return "Enter both username and password."

        with get_db_connection("users.db") as conn:
            cursor = conn.cursor()
            cursor.execute("SELECT password, role FROM users WHERE username=?", (username,))
            row = cursor.fetchone()

        if row and check_password_hash(row[0], password):
            session["user"] = username
            session["role"] = row[1]
            if row[1] == "admin":
                return redirect(url_for("admin_dashboard"))
            else:
                return redirect(url_for("report"))
        else:
            return "Invalid credentials."
    return render_template("login.html")

# Admin Login (separate page optional)
@app.route("/admin_login", methods=["GET", "POST"])
def admin_login():
    if request.method == "POST":
        username = request.form.get("username", "").strip()
        password = request.form.get("password", "").strip()
        if not username or not password:
            return "Enter both username and password."

        with get_db_connection("users.db") as conn:
            cursor = conn.cursor()
            cursor.execute("SELECT password, role FROM users WHERE username=?", (username,))
            row = cursor.fetchone()

        if row and check_password_hash(row[0], password) and row[1] == "admin":
            session["user"] = username
            session["role"] = "admin"
            return redirect(url_for("admin_dashboard"))
        else:
            return "Invalid admin credentials."
    return render_template("admin_login.html")


@app.route("/admin_register", methods=["GET", "POST"])
def admin_register():
    if request.method == "POST":
        username = request.form.get("username", "").strip()
        password = request.form.get("password", "").strip()

        if not username or not password:
            return "Enter both username and password."

        hashed_pw = generate_password_hash(password)
        try:
            with get_db_connection("users.db") as conn:
                cursor = conn.cursor()
                cursor.execute("INSERT INTO users (username, password, role) VALUES (?, ?, ?)",
                               (username, hashed_pw, "admin"))
                conn.commit()
            return redirect(url_for("admin_login"))
        except sqlite3.IntegrityError:
            return "Admin username already exists."
    return render_template("admin_register.html")

# Admin Dashboard
@app.route("/admin")
def admin_dashboard():
    if session.get("role") != "admin":
        return "Access denied!"
    return render_template("admin.html")

# Guest Access
@app.route("/guest")
def guest():
    session["user"] = "Guest"
    session["role"] = "guest"
    return redirect(url_for("report"))

# Bug Report Page
@app.route("/report", methods=["GET", "POST"])
def report():
    if "user" not in session:
        return redirect(url_for("home"))

    if request.method == "POST":
        description = request.form.get("description", "").strip()
        file = request.files.get("attachment")
        file_path = None

        if file and allowed_file(file.filename):
            filename = secure_filename(file.filename)
            file_path = os.path.join(app.config["UPLOAD_FOLDER"], filename)
            file.save(file_path)

        if not description and not file_path:
            return "Please provide a description or attach a file."

        session["description"] = description
        session["attachment"] = file_path
        report_id = create_report_record(session.get("user", "Guest"), description, file_path)
        session["report_id"] = report_id
        session["known_info"] = {}
        session["conversation_history"] = []
        session["current_question_key"] = None
        session["process_cache"] = {
            "turn_summaries": [],
            "asked_keys": [],
            "resolved_keys": [],
            "stage": "diagnosing",
            "latest_issue_summary": "",
        }

        return redirect(url_for("chat"))

    return render_template("report.html")

# Chat Page
@app.route("/chat", methods=["GET", "POST"])
def chat():
    description = session.get("description")
    attachment = session.get("attachment")
    known_info = session.get("known_info", {})
    conversation_history = session.get("conversation_history", [])
    current_question_key = session.get("current_question_key")
    process_cache = session.get("process_cache", {})
    report_id = session.get("report_id")
    action = "send"
    handled_reply = False

    if request.method == "POST":
        action = request.form.get("action", "send").strip().lower()
        user_message = request.form.get("message", "").strip()
        file = request.files.get("attachment")
        if file and allowed_file(file.filename):
            filename = secure_filename(file.filename)
            file_path = os.path.join(app.config["UPLOAD_FOLDER"], filename)
            file.save(file_path)
            attachment = file_path
            session["attachment"] = attachment
            user_message = f"Attached file: {filename}" if not user_message else user_message

        if user_message:
            conversation_history.append({"role": "user", "content": user_message})
            known_info, process_cache = agentic.update_known_info(
                known_info,
                current_question_key,
                user_message,
                process_cache,
            )
            session["known_info"] = known_info
            session["process_cache"] = process_cache
            session["current_question_key"] = None

        if action == "continue":
            if not user_message:
                user_message = "Still not fixed after previous solution."
                conversation_history.append({"role": "user", "content": user_message})
                known_info, process_cache = agentic.update_known_info(
                    known_info,
                    current_question_key,
                    user_message,
                    process_cache,
                )
                session["known_info"] = known_info
                session["process_cache"] = process_cache

            ai_reply, next_key, should_generate_solution, process_cache = agentic.continue_reply(
                description,
                known_info,
                conversation_history,
                process_cache,
            )
            if not ai_reply:
                ai_reply = "Share what happened after the last fix and I will refine the solution."
            conversation_history.append({"role": "assistant", "content": ai_reply})
            session["conversation_history"] = conversation_history
            session["current_question_key"] = next_key
            session["process_cache"] = process_cache
            handled_reply = True

    ai_reply = None
    should_generate_solution = False
    next_key = None

    if (request.method == "POST" and not handled_reply) or not conversation_history:
        ai_reply, next_key, should_generate_solution, process_cache = agentic.next_reply(
            description,
            known_info,
            conversation_history,
            process_cache,
        )
        if not ai_reply:
            ai_reply = "Can you provide more details or any error messages?"
        conversation_history.append({"role": "assistant", "content": ai_reply})
        session["conversation_history"] = conversation_history
        session["current_question_key"] = next_key
        session["process_cache"] = process_cache

    final_solution = None
    final_solution_parsed = None
    if should_generate_solution:
        final_solution = agentic.generate_solution(
            description,
            known_info,
            conversation_history,
            process_cache,
        )
        confidence = agentic.extract_confidence(final_solution)
        update_report_analysis(report_id, final_solution, confidence)
        final_solution_parsed = agentic.parse_solution(final_solution)

    return render_template(
        "chat.html",
        history=conversation_history,
        final_solution=final_solution,
        final_solution_parsed=final_solution_parsed
    )

if __name__ == "__main__":
    app.run(debug=True)