# tools.py — tool definitions only
import math
import uuid
import re
import requests
from bs4 import BeautifulSoup

from smolagents import tool, FinalAnswerTool
from sentence_transformers import SentenceTransformer
import chromadb

# ── Memory singletons ───────────────────────────────────────────────────────
_embedder = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
_chroma   = chromadb.PersistentClient(path="./memory_db")
_col      = _chroma.get_or_create_collection("agent_memory")


# ── Tools ───────────────────────────────────────────────────────────────────

@tool
def search_tool(query: str) -> str:
    """Search the web for current facts, news, prices, or data.
    Args:
        query: Short precise factual search query.
    """
    try:
        from ddgs import DDGS
        results = list(DDGS().text(query, max_results=5))
        if not results:
            return "No results found."
        return "\n\n".join(
            f"{r['title']}\n{r['href']}\n{r['body']}" for r in results
        )
    except Exception as e:
        return f"search_tool error: {e}"


@tool
def fetch_webpage(url: str) -> str:
    """Fetch the full readable plain text of a webpage. Already stripped of HTML tags.
    Do NOT call BeautifulSoup on the result — it is already plain text, not HTML.
    Args:
        url: Full URL to fetch.
    """
    try:
        resp = requests.get(url, headers={"User-Agent": "Mozilla/5.0"}, timeout=15)
        resp.raise_for_status()
        soup = BeautifulSoup(resp.text, "html.parser")
        for tag in soup(["script", "style", "nav", "footer", "header"]):
            tag.decompose()
        return soup.get_text(separator="\n", strip=True)[:15000]
    except Exception as e:
        return f"fetch_webpage error: {e}"


@tool
def wikipedia_search(query: str, sentences: int = 5) -> str:
    """Look up a topic on Wikipedia and return a summary. Use for factual questions
    about people, places, history, science — faster than fetch_webpage for known entities.
    Args:
        query: Topic to look up (e.g. 'Australia', 'Albert Einstein').
        sentences: Number of summary sentences to return (default 5).
    """
    try:
        import wikipedia
        return wikipedia.summary(query, sentences=sentences, auto_suggest=True)
    except Exception as e:
        return f"wikipedia_search error: {e}"
@tool
def wikipedia_section(page_title: str, section: str) -> str:
    """Get a specific section from a Wikipedia page — use when you need a section
    like 'Discography', 'Career', 'Albums' rather than a short summary.
    Args:
        page_title: Exact Wikipedia page title (e.g. 'Mercedes Sosa').
        section: Section heading to find (e.g. 'Discography').
    """
    try:
        import wikipedia
        page = wikipedia.page(page_title, auto_suggest=False)
        text = page.content
        idx = text.find(section)
        return text[idx:idx + 3000] if idx != -1 else text[:3000]
    except Exception as e:
        return f"wikipedia_section error: {e}"


@tool
def read_pdf(file_path: str) -> str:
    """Extract all text from a PDF file.
    Args:
        file_path: Local path to the PDF.
    """
    try:
        from pypdf import PdfReader
        text = "\n".join(p.extract_text() or "" for p in PdfReader(file_path).pages).strip()
        return text[:12000] if text else "PDF has no extractable text (may be image-only)."
    except Exception as e:
        return f"read_pdf error: {e}"


@tool
def read_csv_file(file_path: str, max_rows: int = 200) -> str:
    """Load a CSV file and return its structure and data as text.
    Args:
        file_path: Local path to the CSV.
        max_rows: Maximum rows to return (default 200).
    """
    import pandas as pd
    for enc in ("utf-8", "latin-1", "cp1252", "utf-8-sig"):
        try:
            df = pd.read_csv(file_path, encoding=enc)
            return f"Shape: {df.shape[0]}×{df.shape[1]}\nColumns: {list(df.columns)}\n\n{df.head(max_rows).to_string(index=False)}"
        except UnicodeDecodeError:
            continue
        except Exception as e:
            return f"read_csv_file error: {e}"
    return "read_csv_file error: could not decode file with any known encoding"


@tool
def read_excel_file(file_path: str, sheet_name: str = None, max_rows: int = 200) -> str:
    """Load an Excel file and return its structure and data as text.
    Lists all sheet names first, then reads the requested sheet.
    Args:
        file_path: Local path to the Excel file (.xlsx or .xls).
        sheet_name: Sheet name or index to read (default: first sheet).
        max_rows: Maximum rows to return (default 200).
    """
    try:
        import pandas as pd
        xl = pd.ExcelFile(file_path)
        sheets = xl.sheet_names
        target = sheet_name if sheet_name else sheets[0]
        df = pd.read_excel(file_path, sheet_name=target)
        header = f"Sheets: {sheets}\nReading: '{target}'\nShape: {df.shape[0]}×{df.shape[1]}\nColumns: {list(df.columns)}\n\n"
        return header + df.head(max_rows).to_string(index=False)
    except Exception as e:
        return f"read_excel_file error: {e}"


@tool
def calculator(expression: str) -> str:
    """Evaluate a mathematical expression and return the exact result.
    Supports: + - * / ** % sqrt sin cos tan log log10 pi e abs round.
    Args:
        expression: e.g. 'sqrt(144) + 2**10'
    """
    try:
        safe = {k: getattr(math, k) for k in dir(math) if not k.startswith("_")}
        safe.update({"abs": abs, "round": round, "int": int, "float": float})
        return str(eval(expression, {"__builtins__": {}}, safe))
    except Exception as e:
        return f"calculator error: {e}"


@tool
def count_and_find(text: str, count_type: str = "words", pattern: str = None) -> str:
    """Count elements in text or find regex pattern occurrences.
    Args:
        text: Text to analyse.
        count_type: 'words' | 'characters' | 'lines' | 'sentences' | 'pattern'.
        pattern: Regex string — required when count_type is 'pattern'.
    """
    try:
        if count_type == "words":      return str(len(text.split()))
        if count_type == "characters": return str(len(text))
        if count_type == "lines":      return str(len(text.splitlines()))
        if count_type == "sentences":  return str(len(re.split(r"[.!?]+", text.strip())))
        if count_type == "pattern" and pattern:
            m = re.findall(pattern, text)
            return f"{len(m)} matches: {m[:20]}"
        return "Invalid count_type. Use: words | characters | lines | sentences | pattern."
    except Exception as e:
        return f"count_and_find error: {e}"


@tool
def arxiv_search(query: str, max_results: int = 3) -> str:
    """Search arXiv for academic papers.
    Args:
        query: Search query.
        max_results: Number of results (default 3).
    """
    try:
        import arxiv
        search = arxiv.Search(query=query, max_results=max_results, sort_by=arxiv.SortCriterion.Relevance)
        parts = []
        for p in search.results():
            authors = ", ".join(a.name for a in p.authors[:3])
            parts.append(f"Title: {p.title}\nAuthors: {authors}\nDate: {p.published.strftime('%Y-%m')}\nAbstract: {p.summary[:400]}\nURL: {p.entry_id}")
        return "\n\n---\n\n".join(parts) if parts else "No results."
    except Exception as e:
        return f"arxiv_search error: {e}"


@tool
def flight_time_from_delhi(latitude: float, longitude: float, speed_kmph: float = 850.0) -> str:
    """Calculate flight distance and time from Delhi (Haversine).
    Args:
        latitude: Destination latitude.
        longitude: Destination longitude.
        speed_kmph: Flight speed km/h (default 850).
    """
    r = 6371.0
    la1, lo1 = math.radians(28.6139), math.radians(77.2090)
    la2, lo2 = math.radians(latitude), math.radians(longitude)
    a = math.sin((la2-la1)/2)**2 + math.cos(la1)*math.cos(la2)*math.sin((lo2-lo1)/2)**2
    d = r * 2 * math.asin(math.sqrt(a))
    return f"{d:.0f} km, {d/speed_kmph:.2f} hours"


@tool
def transcribe_audio(file_path: str) -> str:
    """Transcribe an audio file to text using speech recognition.
    Args:
        file_path: Local path to audio file (.mp3, .wav, .m4a, etc.)
    """
    import os
    import requests as _req
    try:
        ext = os.path.splitext(file_path)[1].lower().lstrip(".")
        mime_map = {
            "mp3": "audio/mpeg", "wav": "audio/wav",
            "m4a": "audio/mp4", "ogg": "audio/ogg", "flac": "audio/flac",
        }
        content_type = mime_map.get(ext, "audio/mpeg")
        token = os.environ.get("HF_TOKEN", "")
        with open(file_path, "rb") as f:
            audio_bytes = f.read()
        resp = _req.post(
            "https://api-inference.huggingface.co/models/openai/whisper-large-v3",
            headers={"Authorization": f"Bearer {token}", "Content-Type": content_type},
            data=audio_bytes,
            timeout=60,
        )
        resp.raise_for_status()
        result = resp.json()
        return result.get("text", str(result))
    except Exception as e:
        return f"transcribe_audio error: {e}"


@tool
def extract_text_from_image(file_path: str) -> str:
    """Extract all text and describe visual content from an image using a vision model.
    Use for screenshots, diagrams, charts, tables, or any image containing text.
    Args:
        file_path: Local path to image file (.png, .jpg, .jpeg, .gif, .webp, etc.)
    """
    import os
    import base64
    from huggingface_hub import InferenceClient

    if not os.path.exists(file_path):
        return f"extract_text_from_image error: file not found: {file_path}"

    try:
        ext = file_path.rsplit(".", 1)[-1].lower()
        mime = {
            "jpg": "image/jpeg", "jpeg": "image/jpeg",
            "png": "image/png", "gif": "image/gif", "webp": "image/webp",
        }.get(ext, "image/jpeg")

        with open(file_path, "rb") as f:
            b64 = base64.b64encode(f.read()).decode("utf-8")

        client = InferenceClient(token=os.environ.get("HF_TOKEN", ""))
        prompt = (
            "1. Extract EVERY piece of text visible in this image exactly as written.\n"
            "2. If there is a table, reproduce it row by row with all values.\n"
            "3. If there is a chart or graph, state all axis labels, data points, and exact values.\n"
            "4. If there is a chess board, describe the position in FEN notation.\n"
            "5. State any numbers, dates, names, or codes precisely."
        )
        resp = client.chat_completion(
            model="Qwen/Qwen2.5-VL-72B-Instruct",
            messages=[{
                "role": "user",
                "content": [
                    {"type": "image_url", "image_url": {"url": f"data:{mime};base64,{b64}"}},
                    {"type": "text", "text": prompt},
                ],
            }],
            max_tokens=2000,
        )
        return resp.choices[0].message.content
    except Exception as e:
        return f"extract_text_from_image error: {e}"


@tool
def get_youtube_transcript(video_url: str) -> str:
    """Get the full transcript of a YouTube video. Use for any question referencing a YouTube URL.
    Args:
        video_url: Full YouTube URL (e.g. https://www.youtube.com/watch?v=...)
    """
    try:
        import re as _re
        from youtube_transcript_api import YouTubeTranscriptApi
        vid = _re.search(r"(?:v=|youtu\.be/)([^&\n?#]+)", video_url)
        if not vid:
            return "get_youtube_transcript error: could not extract video ID from URL"
        video_id = vid.group(1)
        ytt = YouTubeTranscriptApi()
        transcript = ytt.fetch(video_id)
        return " ".join(t.text for t in transcript)[:10000]
    except Exception as e:
        return f"get_youtube_transcript error: {e}"


@tool
def extract_table_from_url(url: str) -> str:
    """Extract tables from a webpage as readable text.
    Use when the answer is likely inside an HTML table (sports stats, rankings, schedules, etc.).
    Args:
        url: Full URL of the page containing the table.
    """
    try:
        import pandas as pd
        tables = pd.read_html(url)
        if not tables:
            return "No tables found on this page."
        parts = []
        for i, t in enumerate(tables[:5]):
            parts.append(f"[Table {i+1}]\n{t.to_string(index=False)}")
        return "\n\n".join(parts)[:12000]
    except Exception as e:
        return f"extract_table_from_url error: {e}"


@tool
def download_and_read(url: str) -> str:
    """Download a file from a URL and return its content.
    Supports PDF, CSV, Excel, plain text. Use when a question links directly to a file.
    Args:
        url: Direct URL to the file.
    """
    import os
    import tempfile
    try:
        resp = requests.get(url, headers={"User-Agent": "Mozilla/5.0"}, timeout=30)
        resp.raise_for_status()
        content_type = resp.headers.get("Content-Type", "")
        suffix = ".bin"
        if "pdf" in content_type:
            suffix = ".pdf"
        elif "csv" in content_type or url.endswith(".csv"):
            suffix = ".csv"
        elif "excel" in content_type or url.endswith((".xlsx", ".xls")):
            suffix = ".xlsx"
        elif "text" in content_type:
            return resp.text[:12000]

        with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
            tmp.write(resp.content)
            tmp_path = tmp.name

        try:
            if suffix == ".pdf":
                from pypdf import PdfReader
                text = "\n".join(p.extract_text() or "" for p in PdfReader(tmp_path).pages)
                return text[:12000]
            elif suffix == ".csv":
                import pandas as pd
                df = pd.read_csv(tmp_path)
                return f"Shape: {df.shape}\n{df.head(100).to_string(index=False)}"
            elif suffix == ".xlsx":
                import pandas as pd
                df = pd.read_excel(tmp_path)
                return f"Shape: {df.shape}\n{df.head(100).to_string(index=False)}"
            else:
                return resp.text[:12000]
        finally:
            os.unlink(tmp_path)
    except Exception as e:
        return f"download_and_read error: {e}"


@tool
def analyze_chess_position(fen: str, depth: int = 15) -> str:
    """Analyze a chess position and return the best move and evaluation.
    Use for any question involving a chess board position or asking for the best move.
    Args:
        fen: FEN string of the chess position (e.g. 'rnbqkbnr/pppppppp/8/8/4P3/8/PPPP1PPP/RNBQKBNR b KQkq e3 0 1').
        depth: Search depth (default 15).
    """
    try:
        import chess
        import chess.engine
        import shutil

        board = chess.Board(fen)

        # Try stockfish if available
        stockfish_path = shutil.which("stockfish")
        if stockfish_path:
            with chess.engine.SimpleEngine.popen_uci(stockfish_path) as eng:
                result = eng.analyse(board, chess.engine.Limit(depth=depth))
                best_move = result["pv"][0] if result.get("pv") else None
                score = result["score"].white()
                move_san = board.san(best_move) if best_move else "unknown"
                return f"Best move: {move_san}\nFEN: {fen}\nScore: {score}\nTurn: {'White' if board.turn else 'Black'}"

        # Fallback: list legal moves and basic board info
        legal = [board.san(m) for m in board.legal_moves]
        checks = [m for m in legal if "+" in m or "#" in m]
        return (
            f"Position: {fen}\nTurn: {'White' if board.turn else 'Black'}\n"
            f"Legal moves ({len(legal)}): {', '.join(legal[:30])}\n"
            f"Checks/mates: {checks if checks else 'none'}"
        )
    except Exception as e:
        return f"analyze_chess_position error: {e}"


@tool
def save_memory(text: str, key: str = "", tag: str = "") -> str:
    """Save a fact to long-term memory.
    Call ONLY when user explicitly says: remember / save / store / note that.
    Args:
        text: Fact to save.
        key: Memory key (e.g. 'office', 'son').
        tag: Optional category.
    """
    emb = _embedder.encode(text).tolist()
    mid = key if key else str(uuid.uuid4())
    _col.upsert(ids=[mid], embeddings=[emb], documents=[text], metadatas=[{"key": key, "tag": tag}])
    return f"Saved: {text}"


@tool
def retrieve_memory(query: str) -> str:
    """Retrieve saved personal facts from long-term memory.
    Call whenever user refers to 'my office / son / broker / trip / project / city'.
    Args:
        query: What to search for.
    """
    q_emb = _embedder.encode(query).tolist()
    res = _col.query(query_embeddings=[q_emb], n_results=3)
    if not res["documents"] or not res["documents"][0]:
        return "No saved memory found."
    return res["documents"][0][0]


# ── Registry ────────────────────────────────────────────────────────────────

TOOL_LIST = [
    search_tool, fetch_webpage,
    wikipedia_search, wikipedia_section, FinalAnswerTool(),
    read_pdf, read_csv_file, read_excel_file,
    calculator, count_and_find, arxiv_search,
    flight_time_from_delhi,
    get_youtube_transcript, extract_table_from_url, download_and_read,
    analyze_chess_position,
    transcribe_audio, extract_text_from_image,
    save_memory, retrieve_memory,
]

AUTHORIZED_IMPORTS = [
    "pandas", "numpy", "math", "re", "json", "io",
    "openpyxl", "requests", "bs4", "pypdf",
    "datetime", "collections", "statistics", "arxiv",
]