import requests
from thrift.protocol import TBinaryProtocol
from thrift.transport import TTransport
import sqlite3
import time
import os
import json
import argparse
import sys
import traceback
import lz4.block
import hashlib
import datetime
#import keyboard
from urllib.parse import urlparse, unquote


# Standardized parsing logic for the specific Thrift structure of this app
class QpleClient:
    def __init__(self, config_path="config.json", mode="Job"):
        if not os.path.exists(config_path):
            print(f"[Fatal] Error: {config_path} not found.")
            sys.exit(1)
            
        with open(config_path, 'r', encoding='utf-8') as f:
            self.cfg = json.load(f)

        self.session = requests.Session()
        self.base_url = "https://qple-api-prod.qpleapp.com"
        self.auth_token = self.cfg.get('auth_token', '')
        self.mode = mode 
        
        # Proper encoding for Korean headers
        self.enc_nick = self.cfg['nickname'].encode('utf-8').decode('latin-1')
        self.enc_branch = self.cfg['branch_name'].encode('utf-8').decode('latin-1')
        
        self.timeout = 60
        self.max_retries = 3
        
        self.common_headers = {
            'Accept': '*/*',
            'AppVersion': self.cfg['app_version'],
            'Content-Type': 'application/x-thrift',
            'User-Agent': f'UnityPlayer/{self.cfg["unity_version"]} (UnityWebRequest/1.0, libcurl/8.10.1-DEV)',
            'X-Unity-Version': self.cfg["unity_version"]
        }
        
        # Force login if token is missing
        if not self.auth_token:
            if not self.login():
                print("[Fatal] Initial login failed. Check your config.json.")
                sys.exit(1)

    def get_current_headers(self):
        headers = self.common_headers.copy()
        if self.auth_token:
            headers['Authorization'] = self.auth_token
        headers['Cookie'] = f'CHARCTER_INFO=nickname={self.enc_nick}&branchName={self.enc_branch}&feedType={self.mode}'
        return headers

    def login(self, attempt=1):
        print(f"\n[Auth] >>> Attempting login (Try {attempt}/{self.max_retries})...")
        login_data = self.cfg.get('login_credentials')
        if not login_data:
            print("[Error] No login credentials found in config.json")
            return False

        trans = TTransport.TMemoryBuffer()
        proto = TBinaryProtocol.TBinaryProtocol(trans)
        proto.writeFieldBegin(b'id', 11, 1); b = login_data['device_id'].encode('utf-8'); proto.writeI32(len(b)); proto.trans.write(b); proto.writeFieldEnd()
        proto.writeFieldBegin(b'id', 11, 2); b = login_data['account_id'].encode('utf-8'); proto.writeI32(len(b)); proto.trans.write(b); proto.writeFieldEnd()
        proto.writeFieldBegin(b'id', 11, 3); b = login_data['auth_secret'].encode('utf-8'); proto.writeI32(len(b)); proto.trans.write(b); proto.writeFieldEnd()
        proto.writeFieldStop()
        
        url = f"{self.base_url}/Login/UserLoginReq"
        try:
            r = self.session.post(url, headers=self.common_headers, data=trans.getvalue(), timeout=self.timeout)
            if r.status_code == 200:
                parsed = self.parse_thrift(r.content)
                new_token = parsed.get(6, parsed.get(7))
                if new_token and isinstance(new_token, str):
                    self.auth_token = f"Bearer {new_token}"
                    print(f"[Auth] Login Successful! New token acquired.")
                    return True
            print(f"[Auth] Login failed with status code: {r.status_code}")
        except Exception as e:
            print(f"[Auth] Login error: {e}")
            if attempt < self.max_retries:
                time.sleep(5)
                return self.login(attempt + 1)
        return False

    def parse_thrift(self, data):
        if not data: return {}
        trans = TTransport.TMemoryBuffer(data)
        proto = TBinaryProtocol.TBinaryProtocol(trans)
        def _read_val(ftype):
            if ftype == 2: return proto.readBool()
            elif ftype == 8: return proto.readI32()
            elif ftype == 4: return proto.readDouble()
            elif ftype == 11:
                size = proto.readI32()
                return trans.readAll(size).decode('utf-8', 'ignore')
            elif ftype == 12: return _read_struct()
            elif ftype == 15:
                etype, size = proto.readListBegin()
                return [_read_val(etype) for _ in range(size)]
            return None
        def _read_struct():
            s = {}
            while True:
                _, ftype, fid = proto.readFieldBegin()
                if ftype == 0: break
                s[fid] = _read_val(ftype)
                proto.readFieldEnd()
            return s
        try: return _read_struct()
        except: return {}

    def encode_str_struct(self, fid, val):
        trans = TTransport.TMemoryBuffer()
        proto = TBinaryProtocol.TBinaryProtocol(trans)
        proto.writeFieldBegin(b'id', 11, fid); b = str(val).encode('utf-8'); proto.writeI32(len(b)); proto.trans.write(b); proto.writeFieldEnd(); proto.writeFieldStop()
        return trans.getvalue()

    def encode_submit(self, guid, q_seq, ans):
        # Syntax fix for the duplicate parenthesis typo
        trans = TTransport.TMemoryBuffer()
        proto = TBinaryProtocol.TBinaryProtocol(trans)
        # Field 1: Subject GUID (string)
        proto.writeFieldBegin(b'id', 11, 1); b = guid.encode('utf-8'); proto.writeI32(len(b)); proto.trans.write(b); proto.writeFieldEnd()
        # Field 2: Question Sequence/ID (int32) - CRITICAL: Use Field 5 from response
        proto.writeFieldBegin(b'id', 8, 2); proto.writeI32(int(q_seq)); proto.writeFieldEnd()
        # Field 3: Answer String (string)
        proto.writeFieldBegin(b'id', 11, 3); b = str(ans).encode('utf-8'); proto.writeI32(len(b)); proto.trans.write(b); proto.writeFieldEnd()
        proto.writeFieldStop()
        return trans.getvalue()

    def post(self, path, data, attempt=1):
        if not self.auth_token: return None
        headers = self.get_current_headers()
        try:
            r = self.session.post(f"{self.base_url}/{path}", headers=headers, data=data, timeout=self.timeout)
            if r.status_code == 200:
                return self.parse_thrift(r.content)
            elif r.status_code == 409:
                print(f"\n[!] 409 Conflict on {path}. Cooling down...")
                for i in range(60, 0, -10):
                    print(f"     Wait... {i}s left")
                    time.sleep(10)
                return self.post(path, data, attempt)
            elif r.status_code == 401:
                if self.login(): return self.post(path, data, attempt)
                else: sys.exit(1)
            else:
                print(f"\n[!] HTTP {r.status_code} error on {path}")
        except requests.exceptions.Timeout:
            if attempt < self.max_retries:
                time.sleep(5)
                return self.post(path, data, attempt + 1)
        except Exception as e:
            print(f"\n[!] Request Error: {e}")
        return None

def get_answer_from_db(q_text, field9):
    db_file = 'qple_quiz.db'
    if not os.path.exists(db_file):
        return None
    passage = ""
    if isinstance(field9, str):
        if not field9.startswith("http"):
            passage = field9
    try:
        conn = sqlite3.connect(db_file)
        cursor = conn.cursor()
        cursor.execute("SELECT correct_answer_str FROM QuizQuestions WHERE source = 'qple' AND question_text = ? AND passage = ?", (q_text, passage))
        row = cursor.fetchone()
        conn.close()
        if row:
            return row[0]
    except Exception as e:
        print(f"[DB Error] {e}")
    return None


def get_lz4_image_filename(url):
    parsed = urlparse(url)
    pure_path = parsed.scheme + "://" + parsed.netloc + parsed.path
    pure_path = unquote(pure_path)
    base_path, ext = os.path.splitext(pure_path)
    if not ext:
        ext = ".png"
    compressed = lz4.block.compress(base_path.encode('utf-8'))
    hex_name = compressed.hex()
    if len(hex_name) > 150:
        h = hashlib.md5(hex_name.encode('utf-8')).hexdigest()
        hex_name = hex_name[:150] + "_" + h
    return hex_name + ext

def log_image_mapping(original_url, local_filename):
    base_dir = os.path.dirname(os.path.abspath(__file__))
    # 1. JSON lookup map update
    map_file = os.path.join(base_dir, "image_map.json")
    mapping = {}
    if os.path.exists(map_file):
        try:
            with open(map_file, "r", encoding="utf-8") as f:
                mapping = json.load(f)
        except Exception:
            mapping = {}
            
    mapping[local_filename] = original_url
    
    try:
        with open(map_file, "w", encoding="utf-8") as f:
            json.dump(mapping, f, ensure_ascii=False, indent=2)
    except Exception as e:
        print(f"[Log Error] Failed to write JSON map: {e}")

    # 2. Sequential text download log
    log_file = os.path.join(base_dir, "image_download.log")
    timestamp = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    log_line = f"[{timestamp}] {local_filename} -> {original_url}\n"
    try:
        with open(log_file, "a", encoding="utf-8") as f:
            f.write(log_line)
    except Exception as e:
        print(f"[Log Error] Failed to write text log: {e}")

def download_and_save_image(image_url, target_dir="./assets"):
    if not image_url or not image_url.startswith("http"):
        return image_url
    try:
        os.makedirs(target_dir, exist_ok=True)
        filename = get_lz4_image_filename(image_url)
        dest_path = os.path.join(target_dir, filename)
        
        # Log the mapping of compressed filename to original url
        log_image_mapping(image_url, filename)
        
        if os.path.exists(dest_path):
            return f"/assets/{filename}"
        
        headers = {
            "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36"
        }
        r = requests.get(image_url, headers=headers, timeout=15)
        if r.status_code == 200:
            with open(dest_path, "wb") as f:
                f.write(r.content)
            print(f"     [Image] Downloaded {image_url} -> {dest_path}")
            return f"/assets/{filename}"
        else:
            print(f"     [Image Warning] Failed to download {image_url}, status code: {r.status_code}")
    except Exception as e:
        print(f"     [Image Error] Error downloading {image_url}: {e}")
    return image_url

def sync_db():
    db_file = 'qple_quiz.db'
    app_db_file = r'./qple_quiz.db'
    if os.path.exists(db_file):
        try:
            import shutil
            shutil.copy2(db_file, app_db_file)
            print(f"     [Sync] Copied {db_file} -> {app_db_file}")
        except Exception as e:
            print(f"     [Sync Error] Database synchronization failed: {e}")

def save_to_db(q, cat, diff, correct_ans_str):
    db_file = 'qple_quiz.db'
    conn = sqlite3.connect(db_file)
    cursor = conn.cursor()
    
    q_seq = int(q.get(3, 0))
    txt = str(q.get(7, ''))
    choices = q.get(10, [])
    if not isinstance(choices, list): choices = []
    c1, c2, c3, c4 = (choices + ['', '', '', ''])[:4]
    explanation = str(q.get(12, ''))
    field9 = q.get(9, '')
    passage, image_url = "", ""
    if isinstance(field9, str):
        if field9.startswith("http"): 
            image_url = field9
        else: 
            passage = field9
            
    if image_url:
        image_url = download_and_save_image(image_url)
        
    raw_json = json.dumps(q, ensure_ascii=False)
    
    correct_idx = 0
    for i, c in enumerate(choices):
        if c == correct_ans_str: correct_idx = i; break

    # Try to find existing question with same category, difficulty, question_text, and passage for qple
    cursor.execute("SELECT id FROM QuizQuestions WHERE source = 'qple' AND category = ? AND difficulty = ? AND question_text = ? AND passage = ?", (cat, diff, txt, passage))
    row = cursor.fetchone()
    if row:
        q_id = row[0]
    else:
        # Assign the first vacant numeric key from the front
        cursor.execute("SELECT id FROM QuizQuestions WHERE source = 'qple' ORDER BY id ASC")
        existing_ids = [r[0] for r in cursor.fetchall()]
        if not existing_ids:
            q_id = 1
        elif existing_ids[0] > 1:
            q_id = 1
        else:
            q_id = existing_ids[-1] + 1
            for i in range(len(existing_ids) - 1):
                if existing_ids[i+1] - existing_ids[i] > 1:
                    q_id = existing_ids[i] + 1
                    break

    cursor.execute('''INSERT OR REPLACE INTO QuizQuestions (id, source, category, question_text, choice1, choice2, choice3, choice4, correct_answer, correct_answer_str, explanation, difficulty, passage, image_url, raw_json, oid) 
                      VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)''', (q_id, 'qple', cat, txt, c1, c2, c3, c4, correct_idx, correct_ans_str, explanation, diff, passage, image_url, raw_json, q_seq))
    conn.commit()
    conn.close()

def main():
    parser = argparse.ArgumentParser(description="Qple Scraper")
    parser.add_argument("--mode", type=str, choices=["Job", "Certificate","Financial"], default="Job")
    parser.add_argument("--category", type=str)
    parser.add_argument("--delay", type=float, default=1.5)
    args = parser.parse_args()

    client = QpleClient(mode=args.mode)
    select_list_path = f"Quiz/{args.mode}SelectListReq"
    subject_list_path = f"Quiz/{args.mode}SubjectListReq"
    
    print(f"\n>>> Starting {args.mode} Extraction Loop...")
    job_list = client.post(select_list_path, b'\x00')
    if not job_list or 1 not in job_list: sys.exit(1)

    targets = [j for j in job_list[1] if not args.category or j.get(2) == args.category]
    if not targets: return

    for job in targets:
        job_guid, job_name = job.get(1), job.get(2, 'Unknown')
        print(f"\n>>> [Category] {job_name}")
        subject_list = client.post(subject_list_path, client.encode_str_struct(1, job_guid))
        if not subject_list: continue
            
        for subj in subject_list[1]:
            guid, diff, q_count = subj.get(2), subj.get(3, 'Normal'), subj.get(5, 0)
            if q_count == 0: continue
            
            print(f"  -> [Subject] {diff} (Pool: {q_count})")
            unique_questions = set()
            consecutive_dupes = 0
            
            # Load existing questions from DB to skip them and prevent duplicate limit trigger immediately
            try:
                conn = sqlite3.connect('qple_quiz.db')
                c = conn.cursor()
                c.execute("SELECT question_text, passage FROM QuizQuestions WHERE source = 'qple' AND category = ? AND difficulty = ?", (job_name, diff))
                for row in c.fetchall():
                    txt_db, pas_db = row[0], row[1] or ""
                    unique_questions.add(txt_db + pas_db)
                conn.close()
                print(f"     [Database] Pre-loaded {len(unique_questions)} existing questions to skip.")
            except Exception as db_err:
                print(f"     [Database Warning] Failed to load existing questions: {db_err}")

            while len(unique_questions) < q_count:
                print(f"     [Status] {len(unique_questions)+1}/{q_count} - Fetching...")
                q = client.post("Quiz/QuizQuestionInfoReq", client.encode_str_struct(1, guid))
                if not q or 7 not in q: break
                
                txt = q.get(7)
                q_seq = q.get(5, 0) # Use Field 5 as Sequence ID for submission
                choices = q.get(10, [])
                field9 = q.get(9, '')
                ukey = txt + str(field9)
                
                # Check if this question is already in DB and we know the answer
                known_ans = get_answer_from_db(txt, field9)
                
                res = None
                if known_ans:
                    print(f"     [Database] Found existing question. Submitting correct answer directly: '{known_ans}'")
                    res = client.post("Quiz/QuizSubmitSelectedAnswerReq", client.encode_submit(guid, q_seq, known_ans))
                    true_ans = known_ans
                else:
                    # Advance & Learn (Submit dummy 'O' or first choice to learn answer)
                    dummy_ans = choices[0] if choices else "O"
                    print(f"     [Status] {len(unique_questions)+1}/{q_count} - Submitting dummy '{dummy_ans}' to learn answer with Seq ID: {q_seq}...")
                    res = client.post("Quiz/QuizSubmitSelectedAnswerReq", client.encode_submit(guid, q_seq, dummy_ans))
                    
                    true_ans = ""
                    if res:
                        if res.get(1) == True:
                            true_ans = dummy_ans
                        else:
                            true_ans = res.get(2, '')
                    if not true_ans:
                        true_ans = dummy_ans
                
                # Save to DB if new
                if ukey not in unique_questions:
                    save_to_db(q, job_name, diff, true_ans)
                    unique_questions.add(ukey)
                    consecutive_dupes = 0
                    print(f"     [Saved] {len(unique_questions)}/{q_count} - {txt[:35]}...")
                    # Sync DB periodically for visual feedback
                    #if len(unique_questions) % 50 == 0:
                        #sync_db()
                    if len(unique_questions) >= q_count: break
                else:
                    consecutive_dupes += 1
                    print(f"     [...] Duplicate #{consecutive_dupes} skipped.")
                    if consecutive_dupes >= 35:
                        print(f"     [Status] Hitted {consecutive_dupes} consecutive duplicates. Breaking loop.")
                        break
                
                if res and res.get(3) == True:
                    print("     => Server marked session as finished. (GotoLearningResultPage)")
                    print("     => Visiting TodayLearningResultReq...")
                    client.post("Quiz/TodayLearningResultReq", b'\x00')
                    print("     => Waiting 5 seconds before moving to next subject...")
                    time.sleep(5.0)
                    # break
                
                #if keyboard.is_pressed('esc'):
                 #   break
                
                time.sleep(args.delay)
                
            print("     => Visiting TodayLearningResultReq...")
            client.post("Quiz/TodayLearningResultReq", b'\x00')
            # Sync DB at the end of each subject
            # sync_db()
            print("     => Waiting 5 seconds before moving to next target...")
            time.sleep(5.0)
    print("\nMission Complete.")
    #sync_db()

if __name__ == "__main__":
    main()