app.py

import streamlit as st
import io
import os
import requests
import time
from PyPDF2 import PdfReader
import pdfplumber
from ui import header, left_panel_upload, left_panel_actions, right_panel_extracted, show_results

# ---------- Helpers ----------
def get_api_key():
    return st.secrets.get("GROQ_API_KEY") or os.environ.get("GROQ_API_KEY")

def get_ocr_key():
    return st.secrets.get("OCR_SPACE_API_KEY") or os.environ.get("OCR_SPACE_API_KEY")

def get_assemblyai_key():
    return st.secrets.get("ASSEMBLYAI_API_KEY") or os.environ.get("ASSEMBLYAI_API_KEY")

# ---------- Text Extraction (unchanged) ----------
def extract_text_from_pdf_bytes(pdf_bytes: bytes) -> str:
    text = ""
    try:
        with pdfplumber.open(io.BytesIO(pdf_bytes)) as pdf:
            for page in pdf.pages:
                page_text = page.extract_text()
                if page_text:
                    text += page_text + "\n"
    except:
        try:
            reader = PdfReader(io.BytesIO(pdf_bytes))
            for page in reader.pages:
                text += (page.extract_text() or "") + "\n"
        except:
            pass
    return text.strip()

def extract_text_from_image_bytes(img_bytes: bytes) -> str:
    ocr_api_key = get_ocr_key()
    if not ocr_api_key:
        st.warning("OCR API key missing. Image OCR will not work.")
        return ""
    try:
        response = requests.post(
            "https://api.ocr.space/parse/image",
            files={"file": ("image.png", img_bytes)},
            data={"apikey": ocr_api_key, "language": "eng"}
        )
        result = response.json()
        parsed_text = ""
        if result.get("ParsedResults"):
            parsed_text = "\n".join([r.get("ParsedText", "") for r in result["ParsedResults"]])
        return parsed_text.strip()
    except Exception as e:
        st.warning(f"Image OCR failed: {e}")
        return ""

def transcribe_audio_assemblyai(file_bytes: bytes, api_key: str) -> str:
    try:
        st.info("📤 Uploading audio to AssemblyAI...")
        upload_url = 'https://api.assemblyai.com/v2/upload'
        headers = {'authorization': api_key}

        # Upload audio in chunks
        def upload_stream(data):
            response = requests.post(upload_url, headers=headers, data=data)
            response.raise_for_status()
            return response.json()['upload_url']

        audio_url = upload_stream(file_bytes)
        st.success("✅ Audio uploaded successfully!")

        # Request transcription
        st.info("🧠 Transcription started... please wait.")
        transcript_request = {'audio_url': audio_url, 'language_code': 'en_us'}
        transcript_response = requests.post(
            'https://api.assemblyai.com/v2/transcript',
            headers=headers,
            json=transcript_request
        )
        transcript_response.raise_for_status()
        transcript_id = transcript_response.json()['id']

        # Poll for completion
        while True:
            check_response = requests.get(
                f'https://api.assemblyai.com/v2/transcript/{transcript_id}',
                headers=headers
            )
            status = check_response.json()['status']
            if status == 'completed':
                st.success("✅ Transcription complete!")
                return check_response.json()['text']
            elif status == 'failed':
                st.error("❌ Transcription failed.")
                return ""
            time.sleep(5)
    except Exception as e:
        st.warning(f"Audio transcription failed: {e}")
        return ""

def extract_text(file):
    if file.type == "application/pdf":
        return extract_text_from_pdf_bytes(file.read())
    elif file.type in ["image/png", "image/jpeg", "image/jpg"]:
        return extract_text_from_image_bytes(file.read())
    elif file.type in ["audio/mpeg", "audio/wav"]:
        api_key = get_assemblyai_key()
        if not api_key:
            st.warning("AssemblyAI key missing. Audio transcription will not work.")
            return ""
        return transcribe_audio_assemblyai(file.read(), api_key)
    else:
        return ""

# ---------- GROQ AI Processing (unchanged) ----------
def call_groq_api(prompt: str, api_key: str, max_words=300, max_points=12) -> list:
    url = "https://api.groq.com/openai/v1/chat/completions"
    headers = {"Authorization": f"Bearer {api_key}"}
    payload = {
        "model": "groq/compound",
        "messages": [{"role": "user", "content": prompt}],
        "temperature": 0.2,
        "max_tokens": 500
    }
    try:
        response = requests.post(url, headers=headers, json=payload)
        response.raise_for_status()
        data = response.json()
        content = data.get("choices", [{}])[0].get("message", {}).get("content", "")

        lines = [line.strip("-•* \n") for line in content.split("\n") if line.strip()]
        output = []
        word_count = 0
        for i, line in enumerate(lines):
            if i >= max_points:
                break
            words = len(line.split())
            if word_count + words <= max_words:
                output.append(line)
                word_count += words
            else:
                break
        return output
    except Exception as e:
        st.error(f"Error calling AI API: {e}")
        return ["Error calling AI API"]

# ---------- Main App ----------
def consentclear_app():
    st.set_page_config(layout="wide", page_title="ConsentClear")
    header()  # new styled header from ui.py

    # small layout tweak: left narrower, right bigger
    col1, col2 = st.columns([1, 2])
    with col1:
        files = left_panel_upload()
        analyze_btn, threat_btn = left_panel_actions()

    all_text = ""
    if files:
        for file in files:
            all_text += extract_text(file) + "\n"

    with col2:
        right_panel_extracted(all_text)

    api_key = get_api_key()
    if files and api_key:
        if analyze_btn:
            prompt_analysis = f"""
Analyze the following document thoroughly. Extract the most important points in clear, concise sentences.
Do NOT include titles, headings, numbering, or explanation of the summary. Focus only on actionable and relevant information.
Provide each point as a separate line, directly usable.

Document content:
{all_text}
"""
            analysis_results = call_groq_api(prompt_analysis, api_key, max_words=300, max_points=10)
            show_results(analysis_results, "AI Analysis", color="#111")

        if threat_btn:
            prompt_threats = f"""
Analyze the following document and extract all potential threats clearly and concisely.
Do NOT include headings, numbering, or explanations.
Document content:
{all_text}
"""
            threats_results = call_groq_api(prompt_threats, api_key, max_words=80, max_points=5)
            show_results(threats_results, "Threats", color="#ff4b4b")

            prompt_warnings = f"""
Analyze the following document and extract all potential warnings clearly and concisely.
Do NOT include headings, numbering, or explanations.
Document content:
{all_text}
"""
            warnings_results = call_groq_api(prompt_warnings, api_key, max_words=80, max_points=5)
            show_results(warnings_results, "Warnings", color="#ff914d")