# code.py
import re
import torch
import pandas as pd
import anthropic
import os
from dotenv import load_dotenv
from config import model, tokenizer, label_mapping, big5_dimensions, emotion_big5_priors

# .env 파일 로드
load_dotenv()

def parse_speaker_text(text):
    """대화 텍스트를 파싱하여 화자별 발화를 추출하는 함수"""
    speaker_dict = {}
    lines = text.strip().split('\n')
    
    print(f"📝 입력된 총 줄 수: {len(lines)}")  # 디버그용
    
    for i, line in enumerate(lines):
        line = line.strip()  # 앞뒤 공백 제거
        if not line:  # 빈 줄 건너뛰기
            continue
            
        print(f"🔍 처리 중인 줄 {i+1}: '{line}'")  # 디버그용
        
        # 다양한 패턴 지원
        patterns = [
            r'^(\d+)\s*:\s*(.+)',  # "1: 안녕하세요"
            r'^(\d+)\s*\.\s*(.+)',  # "1. 안녕하세요"
            r'^(\d+)\s+(.+)',      # "1 안녕하세요"
            r'^화자\s*(\d+)\s*:\s*(.+)',  # "화자1: 안녕하세요"
        ]
        
        matched = False
        for pattern in patterns:
            match = re.match(pattern, line)
            if match:
                speaker_id = int(match.group(1))
                utterance = match.group(2).strip()
                
                if speaker_id not in speaker_dict:
                    speaker_dict[speaker_id] = []
                
                speaker_dict[speaker_id].append(utterance)
                print(f"✅ 매칭 성공: 화자{speaker_id} -> '{utterance}'")  # 디버그용
                matched = True
                break
        
        if not matched:
            print(f"❌ 매칭 실패: '{line}'")  # 디버그용
    
    print(f"🎯 최종 결과: {len(speaker_dict)}명의 화자 발견")  # 디버그용
    for speaker_id, utterances in speaker_dict.items():
        print(f"   화자{speaker_id}: {len(utterances)}개 발화")
    
    return speaker_dict

def analyze_emotions(utterances, model, tokenizer, label_mapping):
    """발화 리스트에 대해 감정 분석을 수행하는 함수"""
    results = {}
    
    for idx, text in enumerate(utterances):
        inputs = tokenizer(text, return_tensors="pt")
        with torch.no_grad():
            outputs = model(**inputs)
        emotions = torch.softmax(outputs.logits, dim=-1)
        
        values = emotions.cpu().detach().numpy()
        df = pd.DataFrame(values, columns=[label_mapping[i] for i in range(60)])
        df = df.T.reset_index()
        df.columns = ['Emotion', 'Probability']
        df = df.sort_values(by='Probability', ascending=False).head(5)
        
        results[f"utterance_{idx+1}"] = df
    
    # 열 방향으로 병합
    merged_df = None
    for key, df in results.items():
        df = df.set_index("Emotion")
        df.columns = [key]
        if merged_df is None:
            merged_df = df
        else:
            merged_df = merged_df.join(df, how='outer')
    
    return merged_df

def calculate_probabilistic_mapping(merged_df, emotion_big5_priors):
    """확률적 매핑을 통해 각 발화별 Big5 성향 점수를 계산하는 함수"""
    big5_results = {dim: [] for dim in big5_dimensions}
    utterance_names = []

    for utterance in merged_df.columns:
        utterance_names.append(utterance)
        observed_emotions = merged_df[utterance].dropna()

        if len(observed_emotions) == 0:
            for dim in big5_dimensions:
                big5_results[dim].append(0.0)
            continue

        big5_scores_utterance = {}
        total_weight = sum(observed_emotions.values)

        for dim in big5_dimensions:
            weighted_sum = 0.0
            for emotion, intensity in observed_emotions.items():
                if emotion in emotion_big5_priors:
                    weighted_sum += emotion_big5_priors[emotion][dim] * intensity

            big5_scores_utterance[dim] = weighted_sum / total_weight if total_weight > 0 else 0.0
            big5_results[dim].append(big5_scores_utterance[dim])

    big5_df = pd.DataFrame(big5_results, index=utterance_names)
    return big5_df

def analyze_emotion_patterns(big5_df):
    """감정 패턴 분석"""
    display_df = big5_df.round(3)
    return display_df

def run_probabilistic_mapping(merged_df):
    """확률적 매핑 전체 프로세스 실행"""
    big5_df = calculate_probabilistic_mapping(merged_df, emotion_big5_priors)
    result_summary = analyze_emotion_patterns(big5_df)
    return big5_df, result_summary

def calculate_big5_averages(df):
    """Big5 성격특성 데이터프레임을 입력받아 각 특성의 평균을 계산하는 함수"""
    averages = {}
    for column in df.columns:
        averages[column] = df[column].mean()
    return averages

def analyze_all_speakers(speaker_dict, model, tokenizer, label_mapping):
    """모든 화자에 대해 Big5 분석을 수행하는 함수"""
    all_results = {}

    for speaker_id, utterances in speaker_dict.items():
        emotion_results = analyze_emotions(utterances, model, tokenizer, label_mapping)
        big5_scores, summary = run_probabilistic_mapping(emotion_results)
        big5_avg = calculate_big5_averages(big5_scores)
        all_results[speaker_id] = big5_avg

    return all_results

def stream_response(user_content: str, api_key: str):
    """Anthropic Claude API를 사용하여 시나리오를 생성하는 함수"""
    if not api_key or not api_key.strip():
        return "❌ API 키를 입력해주세요."
    
    if not api_key.startswith('sk-ant-'):
        return "❌ 올바른 Anthropic API 키 형식이 아닙니다. 'sk-ant-'로 시작해야 합니다."
    
    try:
        client = anthropic.Anthropic(api_key=api_key.strip())
        
        stream = client.messages.create(
            model="claude-sonnet-4-20250514",
            max_tokens=3000,
            system="당신은 몰입감 넘치는 드라마틱한 시나리오를 만드는 전문 작가입니다. 심리학적 성격 분석을 바탕으로 인물 간의 갈등과 화학작용이 생생하게 느껴지는 장면을 창조하세요.",
            messages=[
                {"role": "user", "content": user_content}
            ],
            stream=True
        )

        result = ""
        for event in stream:
            if event.type == "content_block_delta":
                result += event.delta.text
        
        return result
        
    except Exception as e:
        return f"❌ API 호출 중 오류가 발생했습니다: {str(e)}"