from flask import Flask, render_template, request, redirect, url_for, session, jsonify, send_file
import os
import pandas as pd
import json
from anthropic import Anthropic
from dotenv import load_dotenv
from functools import wraps

# === Load Environment ===
load_dotenv()
anthropic = Anthropic(api_key=os.getenv("ANTHROPIC_API_KEY"))

# === Flask Setup ===
app = Flask(__name__, template_folder='templates')
app.secret_key = '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'  # 🔐 Change to something secure

UPLOAD_FOLDER = 'uploads'
PROCESSED_FOLDER = 'processed'
os.makedirs(UPLOAD_FOLDER, exist_ok=True)
os.makedirs(PROCESSED_FOLDER, exist_ok=True)

# === In-Memory Storage ===
all_code_frames = []
all_coded_responses = []

# === Auth Configuration ===
VALID_USERNAME = 'admin'
VALID_PASSWORD = 'Ke@123'

def login_required(f):
    @wraps(f)
    def wrapper(*args, **kwargs):
        if 'user' not in session:
            return redirect(url_for('login'))
        return f(*args, **kwargs)
    return wrapper

# === Auth Routes ===
@app.route('/login', methods=['GET', 'POST'])
def login():
    if request.method == 'POST':
        if (request.form['username'] == VALID_USERNAME and
            request.form['password'] == VALID_PASSWORD):
            session['user'] = VALID_USERNAME
            return redirect(url_for('index'))
        else:
            return render_template('login.html', error="Invalid credentials")
    return render_template('login.html')

@app.route('/logout')
def logout():
    session.pop('user', None)
    return redirect(url_for('login'))

# === App Routes ===
@app.route('/')
@login_required
def index():
    return render_template('index.html')

@app.route('/upload', methods=['POST'])
@login_required
def upload_excel():
    file = request.files['file']
    if not file:
        return jsonify({'error': 'No file uploaded'}), 400

    filepath = os.path.join(UPLOAD_FOLDER, file.filename)
    file.save(filepath)

    df = pd.read_excel(filepath)
    sets = df['Set'].unique()
    output = {}

    for s in sets:
        subset = df[df['Set'] == s]
        question = subset['Question'].iloc[0]
        responses = subset['Open end responses'].dropna().tolist()

        output[f'Set {s}'] = {
            'question': question,
            'responses': responses,
            'claude_output': ''
        }

    return jsonify(output)

@app.route('/generate_codeframe', methods=['POST'])
@login_required
def generate_codeframe():
    data = request.json
    question = data.get('question')
    responses = data.get('responses')
    prompt_override = data.get('prompt_override', '').strip()

    if not question or not responses:
        return jsonify({'error': 'Missing inputs'}), 400

    try:
        claude_output = call_claude_for_codeframe_only(question, responses, prompt_override)
        return jsonify({'code_frame': claude_output})
    except Exception as e:
        print("Claude Error:", e)
        return jsonify({'error': str(e)}), 500

@app.route('/finalize_and_process', methods=['POST'])
@login_required
def finalize_and_process():
    data = request.json
    set_name = data.get('set_name')
    question = data.get('question')
    responses = data.get('responses')
    raw_code_frame = data.get('code_frame')

    if not all([set_name, question, responses, raw_code_frame]):
        return jsonify({'success': False, 'error': 'Missing input'}), 400

    try:
        if isinstance(raw_code_frame, list):
            raw_code_frame = "\n".join(raw_code_frame)
        elif not isinstance(raw_code_frame, str):
            raw_code_frame = str(raw_code_frame)

        code_lines = raw_code_frame.strip().splitlines()
        themes_for_set = []
        for line in code_lines:
            line = line.strip()
            if not line.startswith("- ") or ":" not in line:
                continue
            try:
                theme, description = line[2:].split(":", 1)
                themes_for_set.append({
                    "Theme": theme.strip(),
                    "Description": description.strip()
                })
            except ValueError:
                continue

        for i, theme_data in enumerate(themes_for_set, start=1):
            all_code_frames.append({
                "Set": set_name,
                "Code": i,
                "Theme": theme_data["Theme"],
                "Description": theme_data["Description"]
            })

        coded_data = call_claude_for_response_coding(question, responses, raw_code_frame)

        for item in coded_data:
            all_coded_responses.append({
                "Set": set_name,
                "Question": question,
                "Response": item["response"],
                "Themes": " | ".join(item["codes"])  # Use pipe symbol
            })

        return jsonify({'success': True})
    except Exception as e:
        print("Claude Error:", e)
        return jsonify({'success': False, 'error': str(e)})

@app.route('/download/final_output')
@login_required
def download_final_output():
    if not all_code_frames or not all_coded_responses:
        return "No data processed yet", 400

    output_path = os.path.join(PROCESSED_FOLDER, 'final_output.xlsx')
    df1 = pd.DataFrame(all_code_frames)
    df2 = pd.DataFrame(all_coded_responses)

    with pd.ExcelWriter(output_path) as writer:
        df1.to_excel(writer, index=False, sheet_name='Code Frame')
        df2.to_excel(writer, index=False, sheet_name='Coded Responses')

    return send_file(output_path, as_attachment=True)

# === Claude Functions ===

def call_claude_for_codeframe_only(question, responses, custom_prompt=None):
    clean_responses = [str(r).strip() for r in responses if isinstance(r, str) or isinstance(r, int)]
    sample_responses = "\n".join([f"{i+1}. {r}" for i, r in enumerate(clean_responses[:30])])

    if custom_prompt:
        prompt = f"""
Survey Question:
"{question}"

Responses:
{sample_responses}

{custom_prompt}
"""
    else:
        prompt = f"""
You are a professional qualitative researcher.

Survey Question:
"{question}"

Responses:
{sample_responses}

Task:
Generate a code frame with 10–20 meaningful, reusable themes.
Each theme should have a short, professional description.
Return only the code frame as bullet points in this format:

- Theme Name: Description
- Theme Name: Description
        """

    response = anthropic.messages.create(
        model="claude-3-5-haiku-20241022",
        max_tokens=1000,
        temperature=0.4,
        messages=[{"role": "user", "content": prompt}]
    )

    if hasattr(response, 'content') and isinstance(response.content, list):
        output_text = "".join(block.text for block in response.content if hasattr(block, 'text'))
    else:
        output_text = str(response.content)

    return output_text.strip()

def call_claude_for_response_coding(question, responses, code_frame):
    clean_responses = [str(r).strip() for r in responses if isinstance(r, str) or isinstance(r, int)]
    formatted_responses = "\n".join([f"{i+1}. {r}" for i, r in enumerate(clean_responses)])

    prompt = f"""
You are a qualitative coding assistant.

Question:
{question}

Code Frame (Themes):
{code_frame}

Responses:
{formatted_responses}

Your task:
1. Assign 1–4 relevant themes to each response from the code frame basis the corresponding question.
2. If the response is gibberish, tag it as "Junk Words".
3. Return ONLY a valid JSON array like this:
[
  {{ "response": "This is good", "codes": ["Theme 1", "Theme 2"] }},
  {{ "response": "gibberish", "codes": ["Junk Words"] }}
]
Start with [ and end with ]. Do not add any explanation.
"""

    response = anthropic.messages.create(
        model="claude-3-5-haiku-20241022",
        max_tokens=4096,
        temperature=0.4,
        messages=[{"role": "user", "content": prompt}]
    )

    if hasattr(response, 'content') and isinstance(response.content, list):
        output_text = "".join(block.text for block in response.content if hasattr(block, 'text'))
    else:
        output_text = str(response.content)

    # Extract clean JSON
    json_start = output_text.find("[")
    json_end = output_text.rfind("]") + 1
    json_string = output_text[json_start:json_end]

    try:
        return json.loads(json_string)
    except json.JSONDecodeError as e:
        raise ValueError(f"Invalid JSON from Claude: {e}")

# === Run App ===
if __name__ == '__main__':
    app.run(debug=True)