import requests
import numpy as np
import tensorflow as tf
from tensorflow.keras import layers  
import asyncio
from fastapi import FastAPI, Request
from fastapi.responses import StreamingResponse, PlainTextResponse
import sentencepiece as spm
import re
import math
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity

app = FastAPI()    


from fastapi.middleware.cors import CORSMiddleware

origins = [
    "https://insect5386.github.io",
    "https://insect5386.github.io/insect5386"
]

app.add_middleware(
    CORSMiddleware,
    allow_origins=origins,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)
      
sp = spm.SentencePieceProcessor()      
sp.load("kolig_unigram.model")      
      
pad_id = sp.piece_to_id("<pad>")      
if pad_id == -1: pad_id = 0      
start_id = sp.piece_to_id("<start>")      
if start_id == -1: start_id = 1      
end_id = sp.piece_to_id("<end>")      
if end_id == -1: end_id = 2      
unk_id = sp.piece_to_id("<unk>")      
if unk_id == -1: unk_id = 3      
      
vocab_size = sp.get_piece_size()      
max_len = 100      
      
def text_to_ids(text):      
    return sp.encode(text, out_type=int)      
      
def ids_to_text(ids):      
    return sp.decode(ids)      
    
class RotaryPositionalEmbedding(layers.Layer):      
    def __init__(self, dim):      
        super().__init__()      
        inv_freq = 1.0 / (10000 ** (np.arange(0, dim, 2) / dim))      
        self.inv_freq = tf.constant(inv_freq, dtype=tf.float32)      
      
    def call(self, x):      
        batch, heads, seq_len, depth = tf.unstack(tf.shape(x))      
        t = tf.range(seq_len, dtype=tf.float32)      
        freqs = tf.einsum('i,j->ij', t, self.inv_freq)      
        emb_sin = tf.sin(freqs)      
        emb_cos = tf.cos(freqs)      
        emb_cos = tf.reshape(emb_cos, [1, 1, seq_len, -1])      
        emb_sin = tf.reshape(emb_sin, [1, 1, seq_len, -1])      
        x1 = x[..., ::2]      
        x2 = x[..., 1::2]      
        x_rotated = tf.stack([      
            x1 * emb_cos - x2 * emb_sin,      
            x1 * emb_sin + x2 * emb_cos      
        ], axis=-1)      
        x_rotated = tf.reshape(x_rotated, tf.shape(x))      
        return x_rotated    
    
class SwiGLU(tf.keras.layers.Layer):    
    def __init__(self, d_model, d_ff):    
        super().__init__()    
        self.proj = tf.keras.layers.Dense(d_ff * 2)    
        self.out = tf.keras.layers.Dense(d_model)    
    
    def call(self, x):    
        x_proj = self.proj(x)    
        x_val, x_gate = tf.split(x_proj, 2, axis=-1)    
        return self.out(x_val * tf.nn.silu(x_gate))    
            
class GPTBlock(tf.keras.layers.Layer):    
    def __init__(self, d_model, d_ff, num_heads=8, dropout_rate=0.1, adapter_dim=64):      
        super().__init__()      
        self.ln1 = tf.keras.layers.LayerNormalization(epsilon=1e-5)      
        self.mha = tf.keras.layers.MultiHeadAttention(num_heads=num_heads, key_dim=d_model // num_heads)      
        self.dropout1 = tf.keras.layers.Dropout(dropout_rate)     
        self.adapter_down = tf.keras.layers.Dense(adapter_dim, activation='gelu')     
        self.adapter_up = tf.keras.layers.Dense(d_model)      
      
        self.ln2 = tf.keras.layers.LayerNormalization(epsilon=1e-5)      
        self.ffn = SwiGLU(d_model, d_ff)      
        self.dropout2 = tf.keras.layers.Dropout(dropout_rate)     
        self.rope = RotaryPositionalEmbedding(d_model // num_heads)      
      
    def call(self, x, training=False):      
        x_norm = self.ln1(x)      
        b, s, _ = tf.shape(x_norm)[0], tf.shape(x_norm)[1], tf.shape(x_norm)[2]      
        h = self.mha.num_heads      
        d = x_norm.shape[-1] // h      
      
        qkv = tf.reshape(x_norm, [b, s, h, d])      
        qkv = tf.transpose(qkv, [0, 2, 1, 3])      
        q = self.rope(qkv)      
        k = self.rope(qkv)      
        q = tf.reshape(tf.transpose(q, [0, 2, 1, 3]), [b, s, h * d])      
        k = tf.reshape(tf.transpose(k, [0, 2, 1, 3]), [b, s, h * d])      
      
        attn_out = self.mha(query=q, value=x_norm, key=k, use_causal_mask=True, training=training)      
        attn_out = self.dropout1(attn_out, training=training)      
    
        adapter_out = self.adapter_up(self.adapter_down(attn_out))    
        attn_out = attn_out + adapter_out      
      
        x = x + attn_out      
        ffn_out = self.ffn(self.ln2(x))      
        x = x + self.dropout2(ffn_out, training=training)      
        return x    
    
class InteractGPT(tf.keras.Model):      
    def __init__(self, vocab_size, seq_len, d_model, d_ff, n_layers, num_heads=8, dropout_rate=0.1):      
        super().__init__()      
        self.token_embedding = tf.keras.layers.Embedding(vocab_size, d_model)      
        self.blocks = [GPTBlock(d_model, d_ff, num_heads, dropout_rate) for _ in range(n_layers)]      
        self.ln_f = tf.keras.layers.LayerNormalization(epsilon=1e-5)      
      
    def call(self, x, training=False):      
        x = self.token_embedding(x)      
        for block in self.blocks:      
            x = block(x, training=training)      
        x = self.ln_f(x)      
        logits = tf.matmul(x, self.token_embedding.embeddings, transpose_b=True)      
        return logits      
    
model = InteractGPT(vocab_size=vocab_size, seq_len=max_len, d_model=256, d_ff=1024, n_layers=6)        
      
dummy_input = tf.zeros((1, max_len), dtype=tf.int32)  # 배치1, 시퀀스길이 max_len      
_ = model(dummy_input)  # 모델이 빌드됨      
model.load_weights("Flexi.weights.h5")      
print("모델 가중치 로드 완료!")      

  
def is_greedy_response_acceptable(text):  
    text = text.strip()  
  
    # 너무 짧은 문장 거르기  
    if len(text) < 5:  
        return False  
  
    # 단어 수 너무 적은 것도 거름  
    if len(text.split()) < 3:  
        return False  
  
    # ㅋㅋㅋ 같은 자모 연속만 있으면 거름 (단, 'ㅋㅋ' 포함되면 허용)  
    if re.search(r'[ㄱ-ㅎㅏ-ㅣ]{3,}', text) and 'ㅋㅋ' not in text:  
        return False  
  
    # 문장 끝이 어색한 경우 (다/요/죠 등 일반적 형태로 끝나지 않으면 거름)  
    if not re.search(r'(다|요|죠|다\.|요\.|죠\.|다!|요!|죠!|\!|\?|\.)$', text):  
        return False  
  
    return True

def generate_text_sample(model, prompt, max_len=100, max_gen=98,
                         temperature=0.7, top_k=40, min_len=12):
    model_input = text_to_ids(f"<start> {prompt} <sep>")
    model_input = model_input[:max_len]
    generated = list(model_input)

    for _ in range(max_gen):
        pad_len = max(0, max_len - len(generated))
        input_padded = np.pad(generated, (0, pad_len), constant_values=pad_id)
        input_tensor = tf.convert_to_tensor([input_padded])
        logits = model(input_tensor, training=False)
        next_logits = logits[0, len(generated) - 1].numpy()

        # Temperature 적용
        next_logits = next_logits / temperature
        probs = np.exp(next_logits - np.max(next_logits))
        probs = probs / probs.sum()

        # Top-K 필터링
        if top_k is not None and top_k > 0:
            indices_to_remove = probs < np.sort(probs)[-top_k]
            probs[indices_to_remove] = 0
            probs /= probs.sum()

        # 샘플링
        next_token = np.random.choice(len(probs), p=probs)
        generated.append(int(next_token))

        # 디코딩 및 후처리
        decoded = sp.decode(generated)
        for t in ["<start>", "<sep>", "<end>"]:
            decoded = decoded.replace(t, "")
        decoded = decoded.strip()

        if len(generated) >= min_len and (next_token == end_id or decoded.endswith(('요', '다', '.', '!', '?'))):
            if is_greedy_response_acceptable(decoded):
                return decoded
            else:
                continue

    decoded = sp.decode(generated)
    for t in ["<start>", "<sep>", "<end>"]:
        decoded = decoded.replace(t, "")
    return decoded.strip()

    
def mismatch_tone(input_text, output_text):  
    if "ㅋㅋ" in input_text and not re.search(r'ㅋㅋ|ㅎ|재밌|놀|만나|맛집|여행', output_text):  
        return True  
    return False

# 유효한 응답인지 검사
def is_valid_response(response):
    if len(response.strip()) < 2:
        return False
    if re.search(r'[ㄱ-ㅎㅏ-ㅣ]{3,}', response):
        return False
    if len(response.split()) < 2:
        return False
    if response.count(' ') < 2:
        return False
    if any(tok in response.lower() for tok in ['hello', 'this', 'ㅋㅋ']):
        return False
    return True

# 위키 요약 관련
def extract_main_query(text):
    sentences = re.split(r'[.?!]\s*', text)
    sentences = [s.strip() for s in sentences if s.strip()]
    if not sentences:
        return text
    last = sentences[-1]
    last = re.sub(r'[^가-힣a-zA-Z0-9 ]', '', last)
    particles = ['이', '가', '은', '는', '을', '를', '의', '에서', '에게', '한테', '보다']
    for p in particles:
        last = re.sub(rf'\b(\w+){p}\b', r'\1', last)
    return last.strip()

def get_wikipedia_summary(query):
    cleaned_query = extract_main_query(query)
    url = f"https://ko.wikipedia.org/api/rest_v1/page/summary/{cleaned_query}"
    res = requests.get(url)
    if res.status_code == 200:
        return res.json().get("extract", "요약 정보를 찾을 수 없습니다.")
    else:
        return "위키백과에서 정보를 가져올 수 없습니다."

def textrank_summarize(text, top_n=3):
    sentences = re.split(r'(?<=[.!?])\s+', text.strip())
    sentences = [s.strip() for s in sentences if len(s.strip()) > 10]
    if len(sentences) <= top_n:
        return text
    vectorizer = TfidfVectorizer()
    tfidf_matrix = vectorizer.fit_transform(sentences)
    sim_matrix = cosine_similarity(tfidf_matrix)
    np.fill_diagonal(sim_matrix, 0)
    def pagerank(matrix, damping=0.85, max_iter=100, tol=1e-4):
        N = matrix.shape[0]
        ranks = np.ones(N) / N
        row_sums = np.sum(matrix, axis=1)
        row_sums[row_sums == 0] = 1
        for _ in range(max_iter):
            prev_ranks = ranks.copy()
            for i in range(N):
                incoming = matrix[:, i]
                ranks[i] = (1 - damping) / N + damping * np.sum(incoming * prev_ranks / row_sums)
            if np.linalg.norm(ranks - prev_ranks) < tol:
                break
        return ranks
    scores = pagerank(sim_matrix)
    ranked_idx = np.argsort(scores)[::-1]
    selected_idx = sorted(ranked_idx[:top_n])
    summary = ' '.join([sentences[i] for i in selected_idx])
    return summary

def summarize_from_wikipedia(query, top_n=3):  
    raw_summary = get_wikipedia_summary(query)  
    first_summary = textrank_summarize(raw_summary, top_n=top_n)  
    second_summary = textrank_summarize(first_summary, top_n=top_n)  
    return second_summary


def simple_intent_classifier(text):
    text = text.lower()
    greet_keywords = ["안녕", "반가워", "이름", "누구", "소개", "어디서 왔", "정체", "몇 살", "너 뭐야"]
    info_keywords = ["설명", "정보", "무엇", "뭐야", "어디", "누구", "왜", "어떻게", "종류", "개념"]
    math_keywords = ["더하기", "빼기", "곱하기", "나누기", "루트", "제곱", "+", "-", "*", "/", "=", "^", "√", "계산", "몇이야", "얼마야"]
    if any(kw in text for kw in greet_keywords):
        return "인사"
    elif any(kw in text for kw in info_keywords):
        return "정보질문"
    elif any(kw in text for kw in math_keywords):
        return "수학질문"
    else:
        return "일상대화"

def parse_math_question(text):
    text = text.replace("곱하기", "*").replace("더하기", "+").replace("빼기", "-").replace("나누기", "/").replace("제곱", "*2")
    text = re.sub(r'루트\s(\d+)', r'math.sqrt(\1)', text)
    try:
        result = eval(text)
        return f"정답은 {result}입니다."
    except:
        return "계산할 수 없는 수식이에요. 다시 한번 확인해 주세요!"

# 최종 응답 함수
def respond(input_text):
    intent = simple_intent_classifier(input_text)

    if "이름" in input_text:
        return "제 이름은 Flexi입니다."

    if "누구" in input_text:
        return "저는 Flexi라고 해요."

    if intent == "수학질문":
        return parse_math_question(input_text)

    if intent == "인사":
        return "반가워요! 무엇을 도와드릴까요?"

    if intent == "정보질문":
        keyword = re.sub(r"(에 대해|에 대한|에 대해서)?\s*(설명해줘|알려줘|뭐야|개념|정의|정보)?", "", input_text).strip()
        if not keyword:
            return "어떤 주제에 대해 궁금한가요?"
        summary = summarize_from_wikipedia(keyword)
        return f"{summary}\n다른 궁금한 점 있으신가요?"

    # 일상 대화: 샘플링 + fallback
    response = generate_text_sample(model, input_text)
    if not is_valid_response(response) or mismatch_tone(input_text, response):
        response = generate_text_sample(model, input_text)
    return response

@app.get("/generate", response_class=PlainTextResponse)
async def generate(request: Request):
    prompt = request.query_params.get("prompt", "안녕하세요")
    response_text = respond(prompt)
    return response_text