import streamlit as st
import os
import pathlib
import pandas as pd
from collections import defaultdict
import json
import copy
import plotly.express as px



@st.cache_data
def load_local_corpus(corpus_file, columns_to_combine=["title", "text"]):
    if corpus_file is None:
        return None
    did2text = {}
    id_key = "_id"
    with corpus_file as f:
        for idx, line in enumerate(f):
            uses_bytes = not (type(line) == str)
            if uses_bytes:
                if idx == 0 and "doc_id" in line.decode("utf-8"):
                    continue
                inst = json.loads(line.decode("utf-8"))
            else:
                if idx == 0 and "doc_id" in line:
                    continue
                inst = json.loads(line)
            all_text = " ".join([inst[col] for col in columns_to_combine if col in inst])
            if id_key not in inst:
                id_key = "doc_id"
            did2text[inst[id_key]] = {
                "text": all_text,
                "title": inst["title"] if "title" in inst else "",
            }
    return did2text


@st.cache_data
def load_local_queries(queries_file):
    if queries_file is None:
        return None
    qid2text = {}
    id_key = "_id"
    with queries_file as f:
        for idx, line in enumerate(f):
            uses_bytes = not (type(line) == str)
            if uses_bytes:
                if idx == 0 and "query_id" in line.decode("utf-8"):
                    continue
                inst = json.loads(line.decode("utf-8"))
            else:
                if idx == 0 and "query_id" in line:
                    continue
                inst = json.loads(line)
            if id_key not in inst:
                id_key = "query_id"
            qid2text[inst[id_key]] = inst["text"]
    return qid2text


@st.cache_data
def load_local_qrels(qrels_file):
    if qrels_file is None:
        return None
    qid2did2label = defaultdict(dict)
    with qrels_file as f:
        for idx, line in enumerate(f):
            uses_bytes = not (type(line) == str)
            if uses_bytes:
                if idx == 0 and "qid" in line.decode("utf-8") or "query-id" in line.decode("utf-8"):
                    continue
                cur_line = line.decode("utf-8")
            else:
                if idx == 0 and "qid" in line or "query-id" in line:
                    continue
                cur_line = line
            try:
                qid, _, doc_id, label = cur_line.split()
            except:
                qid, doc_id, label = cur_line.split()
            qid2did2label[str(qid)][str(doc_id)] = int(label)

    return qid2did2label



@st.cache_data
def load_jsonl(f):
    did2text = defaultdict(list)
    sub_did2text = {}

    for idx, line in enumerate(f):
        inst = json.loads(line)
        if "question" in inst:
            docid = inst["metadata"][0]["passage_id"] if "doc_id" not in inst else inst["doc_id"]
            did2text[docid].append(inst["question"])
        elif "text" in inst:
            docid = inst["doc_id"] if "doc_id" in inst else inst["did"]
            did2text[docid].append(inst["text"])
            sub_did2text[inst["did"]] = inst["text"]
        elif "query" in inst:
            docid = inst["doc_id"] if "doc_id" in inst else inst["did"]
            did2text[docid].append(inst["query"])
        else:
            breakpoint()
            raise NotImplementedError("Need to handle this case")
                
    return did2text, sub_did2text



@st.cache_data(persist="disk")
def get_dataset(dataset_name: str, input_fields_doc, input_fields_query):
    if type(input_fields_doc) == str:
        input_fields_doc = input_fields_doc.strip().split(",")
    if type(input_fields_query) == str:
        input_fields_query = input_fields_query.strip().split(",")

    if dataset_name == "":
        return {}, {}, {}
    
    else:
        raise NotImplementedError("Dataset not implemented")