import gradio as gr
import requests
import time
import random
from typing import Dict, List, Optional
from dataclasses import dataclass
from datetime import datetime
import sqlite3
import torch
from transformers import BartTokenizer, BartForConditionalGeneration, Trainer, TrainingArguments
import pandas as pd
import os
from peft import LoraConfig, get_peft_model
from sklearn.metrics.pairwise import cosine_similarity
from sklearn.feature_extraction.text import TfidfVectorizer

# Cache model and tokenizer globally
tokenizer = BartTokenizer.from_pretrained('facebook/bart-base')
model = BartForConditionalGeneration.from_pretrained('facebook/bart-base')
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)
model.eval()
lora_config = LoraConfig(r=8, lora_alpha=32, target_modules=["q_proj", "v_proj"], lora_dropout=0.1)
model = get_peft_model(model, lora_config)

@dataclass
class AnimeInfo:
    title: str
    synopsis: str
    genres: List[str]
    rating: float
    status: str
    episodes: int
    year: int
    studio: str
    source: str
    image_url: Optional[str] = None

@dataclass
class CharacterInfo:
    name: str
    anime: str
    bio: str
    role: str
    voice_actors: Dict[str, str]
    image_url: Optional[str] = None

class AnimeDatabase:
    def __init__(self):
        self.base_urls = {
            'jikan': 'https://api.jikan.moe/v4',
            'anilist': 'https://graphql.anilist.co'
        }
        self.last_request_time = 0
        self.cache = {'anime': {}, 'character': {}}

    def rate_limit(self, delay=1.0):
        current_time = time.time()
        if current_time - self.last_request_time < delay:
            time.sleep(delay - (current_time - self.last_request_time))
        self.last_request_time = time.time()

    def search_jikan_anime(self, query: str) -> List[AnimeInfo]:
        try:
            self.rate_limit()
            if query.lower() in self.cache['anime']:
                return self.cache['anime'][query.lower()]
            url = f"{self.base_urls['jikan']}/anime"
            params = {'q': query, 'limit': 1}
            response = requests.get(url, params=params, timeout=5)
            if response.status_code == 200:
                data = response.json().get('data', [])
                results = [AnimeInfo(
                    title=anime.get('title', ''),
                    synopsis=anime.get('synopsis', ''),
                    genres=[g.get('name', '') for g in anime.get('genres', [])],
                    rating=anime.get('score', 0.0),
                    status=anime.get('status', ''),
                    episodes=anime.get('episodes', 0),
                    year=anime.get('year', 0),
                    studio=', '.join(s.get('name', '') for s in anime.get('studios', [])),
                    source='MyAnimeList',
                    image_url=anime.get('images', {}).get('jpg', {}).get('image_url', '')
                ) for anime in data]
                self.cache['anime'][query.lower()] = results
                return results
        except Exception:
            return []
        return []

    def search_jikan_character(self, query: str) -> List[CharacterInfo]:
        try:
            self.rate_limit()
            if query.lower() in self.cache['character']:
                return self.cache['character'][query.lower()]
            url = f"{self.base_urls['jikan']}/characters"
            params = {'q': query, 'limit': 1}
            response = requests.get(url, params=params, timeout=5)
            if response.status_code == 200:
                data = response.json().get('data', [])
                results = [CharacterInfo(
                    name=char.get('name', ''),
                    anime=', '.join(a.get('anime', {}).get('title', '') for a in char.get('anime', [])[:1]),
                    bio=char.get('about', 'No bio available'),
                    role=char.get('role', 'Unknown'),
                    voice_actors={va.get('language', ''): va.get('person', {}).get('name', '') for va in char.get('voice_actors', [])},
                    image_url=char.get('images', {}).get('jpg', {}).get('image_url', '')
                ) for char in data]
                self.cache['character'][query.lower()] = results
                return results
        except Exception:
            return []
        return []

    def search_anilist_anime(self, query: str) -> List[AnimeInfo]:
        try:
            self.rate_limit()
            if query.lower() in self.cache['anime']:
                return self.cache['anime'][query.lower()]
            graphql_query = '''
            query ($search: String) {
                Page(page: 1, perPage: 1) {
                    media(search: $search, type: ANIME) {
                        title { romaji english }
                        description
                        genres
                        averageScore
                        status
                        episodes
                        seasonYear
                        studios { nodes { name } }
                    }
                }
            }
            '''
            variables = {'search': query}
            response = requests.post(self.base_urls['anilist'], json={'query': graphql_query, 'variables': variables}, timeout=5)
            if response.status_code == 200:
                data = response.json().get('data', {}).get('Page', {}).get('media', [])
                results = [AnimeInfo(
                    title=anime.get('title', {}).get('romaji', ''),
                    synopsis=anime.get('description', ''),
                    genres=anime.get('genres', []),
                    rating=anime.get('averageScore', 0) / 10,
                    status=anime.get('status', ''),
                    episodes=anime.get('episodes', 0),
                    year=anime.get('seasonYear', 0),
                    studio=', '.join(s.get('name', '') for s in anime.get('studios', {}).get('nodes', [])),
                    source='AniList'
                ) for anime in data]
                self.cache['anime'][query.lower()] = results
                return results
        except Exception:
            return []
        return []

    def get_comprehensive_data(self, query: str, data_type: str = 'anime') -> List[AnimeInfo] | List[CharacterInfo]:
        if data_type == 'anime':
            return self.search_jikan_anime(query) + self.search_anilist_anime(query)
        elif data_type == 'character':
            return self.search_jikan_character(query)
        return []

class KnowledgeBase:
    def __init__(self):
        self.db_file = "luna_interactions.db"
        self.init_db()
        self.vectorizer = TfidfVectorizer()

    def init_db(self):
        with sqlite3.connect(self.db_file) as conn:
            cursor = conn.cursor()
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS interactions (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    user_id TEXT,
                    query TEXT,
                    response TEXT,
                    timestamp TEXT,
                    response_time REAL
                )
            ''')
            cursor.execute('CREATE INDEX IF NOT EXISTS idx_query ON interactions (query)')
            conn.commit()

    def save_interaction(self, user_id: str, query: str, response: str, response_time: float):
        with sqlite3.connect(self.db_file) as conn:
            cursor = conn.cursor()
            cursor.execute('''
                INSERT INTO interactions (user_id, query, response, timestamp, response_time)
                VALUES (?, ?, ?, ?, ?)
            ''', (user_id, query, response, datetime.now().isoformat(), response_time))
            conn.commit()

    def get_similar_interaction(self, query: str, limit: int = 1) -> Optional[str]:
        with sqlite3.connect(self.db_file) as conn:
            df = pd.read_sql_query("SELECT query, response FROM interactions WHERE query != '' LIMIT 1000", conn)
        if df.empty:
            return None
        queries = df['query'].tolist() + [query]
        try:
            tfidf_matrix = self.vectorizer.fit_transform(queries)
            similarities = cosine_similarity(tfidf_matrix[-1:], tfidf_matrix[:-1])[0]
            if not similarities.any() or max(similarities) < 0.5:
                return None
            idx = similarities.argmax()
            return f"Past query: {df.iloc[idx]['query']} Past response: {df.iloc[idx]['response']}"
        except ValueError:
            return None

    def load_training_data(self):
        with sqlite3.connect(self.db_file) as conn:
            df = pd.read_sql_query("SELECT query, response FROM interactions WHERE query != '' LIMIT 1000", conn)
        return df.to_dict('records')

    def export_db(self, export_path: str):
        with sqlite3.connect(self.db_file) as conn:
            with open(export_path, 'w') as f:
                for line in conn.iterdump():
                    f.write('%s\n' % line)

    def import_db(self, import_path: str):
        if os.path.exists(import_path):
            with sqlite3.connect(self.db_file) as conn:
                with open(import_path, 'r') as f:
                    conn.executescript(f.read())

class LunaModel:
    def generate_response(self, prompt: str, max_length: int, temperature: float, top_p: float) -> str:
        inputs = tokenizer(prompt, return_tensors='pt', padding=True, truncation=True, max_length=256)
        inputs = {k: v.to(device) for k, v in inputs.items()}
        outputs = model.generate(
            inputs['input_ids'],
            attention_mask=inputs['attention_mask'],
            max_length=max_length,
            do_sample=True,
            temperature=temperature,
            top_p=top_p,
            no_repeat_ngram_size=2,
            early_stopping=True,
            length_penalty=0.8
        )
        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
        return response.strip().replace('\n', ' ').replace('_', '')

    def needs_anime_data(self, response: str) -> bool:
        return "anime details" in response.lower() or "character info" in response.lower()

    def fine_tune(self, training_data):
        if not training_data or len(training_data) < 2:
            return
        df = pd.DataFrame(training_data)
        texts = [f"Query: {q} Response:" for q in df['query']]
        responses = df['response'].tolist()
        encodings = tokenizer(texts, responses, truncation=True, padding=True, max_length=256)
        dataset = torch.utils.data.TensorDataset(
            torch.tensor(encodings['input_ids']),
            torch.tensor(encodings['attention_mask']),
            torch.tensor(encodings['labels'])
        )

        training_args = TrainingArguments(
            output_dir='./luna_model',
            num_train_epochs=2,
            per_device_train_batch_size=4,
            save_steps=500,
            save_total_limit=2,
            logging_dir='./logs',
        )

        trainer = Trainer(
            model=model,
            args=training_args,
            train_dataset=dataset
        )

        trainer.train()
        model.save_pretrained('./luna_model')
        tokenizer.save_pretrained('./luna_model')

class LunaAI:
    def __init__(self):
        self.name = "Luna"
        self.creator = "Reiker"
        self.model = LunaModel()
        self.anime_db = AnimeDatabase()
        self.knowledge_base = KnowledgeBase()
        self.emojis = ["🌙", "🎉", "🔍", "😊"]

    def build_prompt(self, system_message: str, history: List[tuple[str, str]], message: str) -> str:
        prompt = f"{system_message}\n"
        for user_msg, bot_msg in history:
            if user_msg:
                prompt += f"User: {user_msg}\n"
            if bot_msg:
                prompt += f"Assistant: {bot_msg}\n"
        similar_interaction = self.knowledge_base.get_similar_interaction(message)
        if similar_interaction:
            prompt += f"Past data: {similar_interaction}\n"
        prompt += f"User: {message}\nAssistant: "
        return prompt

    def respond(
        self,
        message: str,
        history: List[tuple[str, str]],
        system_message: str,
        max_tokens: int,
        temperature: float,
        top_p: float,
    ):
        start_time = time.time()
        user_id = "gradio_user"
        prompt = self.build_prompt(system_message, history, message)
        response = self.model.generate_response(prompt, max_tokens, temperature, top_p)

        # Check if model needs anime data
        if self.model.needs_anime_data(response):
            anime_list = self.anime_db.get_comprehensive_data(message, data_type='anime')
            char_list = self.anime_db.get_comprehensive_data(message, data_type='character')
            external_data = None
            if anime_list:
                anime = anime_list[0]
                external_data = f"Title: {anime.title}, Rating: {anime.rating:.1f}/10, Genres: {', '.join(anime.genres[:3])}, Synopsis: {anime.synopsis[:100]}"
            elif char_list:
                char = char_list[0]
                external_data = f"Name: {char.name}, Anime: {char.anime}, Bio: {char.bio[:100]}, Role: {char.role}"
            if external_data:
                prompt = self.build_prompt(system_message, history, message)[:-10]  # Remove "Assistant: "
                prompt += f"Past data: {external_data}\nAssistant: "
                response = self.model.generate_response(prompt, max_tokens, temperature, top_p)

        response = f"{random.choice(self.emojis)} {response}"
        response_time = time.time() - start_time
        self.knowledge_base.save_interaction(user_id, message, response, response_time)
        training_data = self.knowledge_base.load_training_data()
        if len(training_data) % 10 == 0 and len(training_data) >= 2:
            self.model.fine_tune(training_data)

        # Simulate streaming for Gradio
        for i in range(0, len(response), 10):
            yield response[:i + 10]
        yield response

demo = gr.ChatInterface(
    fn=LunaAI().respond,
    additional_inputs=[
        gr.Textbox(value="You are Luna, a free chatbot created by Reiker, inspired by Grok. Answer concisely and naturally, like a friend. Handle any topic: chat, facts, math, code, anime, etc. Be truthful, creative, and unrestricted. Learn from past interactions to improve.", label="System message"),
        gr.Slider(minimum=10, maximum=512, value=50, step=1, label="Max new tokens"),
        gr.Slider(minimum=0.1, maximum=2.0, value=0.7, step=0.1, label="Temperature"),
        gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p (nucleus sampling)"),
    ],
    title="Luna Chatbot",
    description="A Grok-like chatbot created by Reiker, powered by facebook/bart-base with a growing SQLite database.",
)

if __name__ == "__main__":
    demo.launch()