from flask import Flask, request, jsonify, render_template, send_file
import pandas as pd
import os
import json
from anthropic import Anthropic
from dotenv import load_dotenv

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

# App setup
app = Flask(__name__, template_folder='templates')
UPLOAD_FOLDER = 'uploads'
PROCESSED_FOLDER = 'processed'
os.makedirs(UPLOAD_FOLDER, exist_ok=True)
os.makedirs(PROCESSED_FOLDER, exist_ok=True)

# Store results across all sets
all_code_frames = []
all_coded_responses = []

# Claude: Generate Code Frame
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])])

    # Use custom prompt if provided, otherwise default
    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}]
    )

    # Handle new-style Claude output
    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()




# Claude: Code responses
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": ["Easy to Use", "Helpful"] }},
  {{ "response": "asdfasdf", "codes": ["Junk Words"] }}
]
Start with [ and end with ]. Do not include any explanation.
    """

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

    # ✅ Extract text safely
    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)

    print("RAW CLAUDE RESPONSE:\n", output_text)

    # ✅ Extract and parse JSON only
    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}")



# Home page
@app.route('/')
def index():
    return render_template('index.html')

# Upload Excel
@app.route('/upload', methods=['POST'])
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)

# Generate Code Frame
@app.route('/generate_codeframe', methods=['POST'])
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


# Finalize & Code Responses
@app.route('/finalize_and_process', methods=['POST'])
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:
        # Normalize code frame input
        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)

        # Step 1: Parse clean Theme + Description lines from 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  # Skip malformed lines

        # Step 2: Append to global code frame (with code numbers)
        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"]
            })

        # Step 3: Send to Claude for response coding
        coded_data = call_claude_for_response_coding(question, responses, raw_code_frame)

        # Step 4: Append coded responses to global output
        for item in coded_data:
            all_coded_responses.append({
                "Set": set_name,
                "Question": question,
                "Response": item["response"],
                "Themes": " | ".join(item["codes"])  # ✅ use pipe delimiter
            })

        return jsonify({'success': True})

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


# Download Final Combined Excel
@app.route('/download/final_output')
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)

@app.route('/reset_session', methods=['POST'])
def reset_session():
    all_code_frames.clear()
    all_coded_responses.clear()
    return jsonify({'success': True, 'message': 'Session cleared.'})


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