"""
ai_loop.py — Autonomous music production loop using an Anthropic-compatible AI.

The loop:
  1. Calls an Anthropic-compatible endpoint to generate a SongScore JSON
  2. Validates and saves the score to scores/
  3. Optionally renders it into Ableton Live

Configuration (environment variables OR command-line args):
  AI_BASE_URL   → API base URL     (default: https://api.anthropic.com)
  AI_API_KEY    → API key          (required)
  AI_MODEL      → model name       (default: GLM-5-Turbo)
  AI_MAX_TOKENS → max output tokens (default: 4096)
  RENDER_AFTER  → "1" to auto-render each score in Ableton (default: 0)
  LOOP_COUNT    → how many songs to produce (default: 10, 0 = infinite)
  LOOP_DELAY    → seconds between generations (default: 5)
  LIB_ROOT      → path to libreria/reggaeton (auto-detected)

Usage examples:
  # OpenRouter with Claude Haiku
  AI_BASE_URL=https://openrouter.ai/api/v1 AI_API_KEY=sk-xxx python ai_loop.py

  # Local LM Studio (Anthropic-compatible)
  AI_BASE_URL=http://localhost:1234/v1 AI_API_KEY=sk-any python ai_loop.py --count 5

  # Real Anthropic + auto-render
  AI_API_KEY=sk-ant-xxx RENDER_AFTER=1 python ai_loop.py
"""

import argparse
import json
import logging
import os
import sys
import time
from datetime import datetime
from pathlib import Path

_THIS_DIR  = Path(__file__).resolve().parent
_PROJ_DIR  = _THIS_DIR.parent
_BASE_DIR  = _PROJ_DIR.parent

for _p in (str(_THIS_DIR), str(_PROJ_DIR)):
    if _p not in sys.path:
        sys.path.insert(0, _p)

from score_engine   import SongScore, SCORES_DIR
from score_renderer import ScoreRenderer

logging.basicConfig(
    level  = logging.INFO,
    format = "%(asctime)s [ai_loop] %(levelname)s: %(message)s",
)
log = logging.getLogger("ai_loop")

_DEFAULT_LIB_ROOT = str(_BASE_DIR / "libreria" / "reggaeton")

SYSTEM_PROMPT = """\
You are a professional reggaeton and Latin urban music producer AI.
Your ONLY job is to output a valid SongScore JSON object for each request.
Do NOT include any explanation, markdown code fences, or commentary.
Output ONLY raw JSON that starts with { and ends with }.

SongScore schema:
{
  "meta": {
    "title":          "<unique Spanish/English song title>",
    "tempo":          <85-105>,
    "key":            "<Am|Dm|Em|Fm|Gm|C#m|C|F|G|Bb>",
    "genre":          "reggaeton",
    "time_signature": "4/4",
    "gap_bars":       <1.0-4.0>
  },
  "structure": [
    { "name": "<section name>", "duration_bars": <integer> },
    ...
  ],
  "tracks": [
    {
      "id":   "<unique_id>",
      "name": "<Track Name>",
      "type": "<audio|midi>",

      "clips": [
        { "section": "<section name>", "sample": "kick/auto", "loop": true }
      ],

      "instrument": "<Wavetable|Operator>",

      "mixer": { "volume": <0-1>, "pan": <-1 to 1>, "eq_preset": "<optional>" }
    }
  ]
}

Available sample categories — use EXACTLY "category/auto" in the "sample" field:
  "kick/auto"          -> Kick drums (23 samples: main + reggaeton 3 + SentimientoLatino)
  "snare/auto"         -> Snares (29 samples)
  "hihat/auto"         -> Hi-hats (6 samples)
  "drumloops/auto"     -> Drum loops with BPM (70 samples, 83-160 BPM range)
  "perc/auto"          -> Percussion loops (21 samples)
  "bass/auto"          -> Bass samples (41 samples)
  "fx/auto"            -> FX and transitions (45 samples)
  "synth/auto"         -> Synth leads, plucks, arps (54 samples)
  "pad/auto"           -> Pads and textures (23 samples)
  "keys/auto"          -> Piano, rhodes, keys (13 samples)
  "vocals/auto"        -> Vocal chops, phrases, ad-libs (42 samples)
  "oneshots/auto"      -> One-shot melodic hits (63 samples)
  "impact/auto"        -> Impact hits (7 samples)
  "fill/auto"          -> Drum fills (5 samples)
  "bells/auto"         -> Bells and mallets (16 samples)
  "chords/auto"        -> Chord samples and MIDI (56 samples)
  "guitar/auto"        -> Guitar loops (3 samples)
  "brass/auto"         -> Brass hits (included in oneshots)
  "music_loop/auto"    -> Full music loops (7 samples)

The system automatically picks the BEST sample matching the project BPM and key.

Available MIDI patterns (use in "pattern" field for type:"midi" tracks):
  dembow_minimal  dembow_standard  dembow_double
  bass_sub  bass_pluck  bass_octaves  bass_sustained
  chords_verse  chords_chorus  melody_simple

Available EQ presets: kick, kick_sub, kick_punch, snare, snare_body, snare_crack,
  bass, bass_clean, synth, synth_air, pad_warm, master

Rules:
- Every track MUST have at least one clip.
- Every clip MUST reference a valid section name from the structure array.
- Always include at minimum: kick, snare or drum_loop, dembow, bass tracks.
- Use 6-12 tracks for a full production. Be creative with synths, pads, vocals, bells.
- Vary everything: title, tempo, key, gap_bars, structure length (40-90 total bars).
- Use realistic reggaeton/latin structures (Intro, Verse, Pre-Chorus, Chorus, Bridge, Outro).
- Mix audio and MIDI tracks creatively. Use diverse sample categories.
- Section names MUST be unique. Use numbered suffixes: "Intro", "Verse A", "Pre-Chorus",
  "Chorus A", "Verse B", "Chorus B", "Bridge", "Outro". NEVER repeat a section name.
- Do NOT include "start_bar" in sections. The engine calculates it automatically.
- Audio tracks use "sample" field. MIDI tracks use "pattern" field. Do NOT mix them.
- Output ONLY the JSON object. Nothing else.
"""

USER_PROMPT_TEMPLATE = """\
Generate song number {index} of {total}.
Make it unique. Use creative Spanish/English titles.
Output only the SongScore JSON.
"""


def _build_client(base_url: str, api_key: str):
    try:
        import anthropic
    except ImportError:
        log.error("anthropic package not installed. Run: pip install anthropic")
        sys.exit(1)

    kwargs = {"api_key": api_key}
    if base_url and "anthropic.com" not in base_url:
        kwargs["base_url"] = base_url

    return anthropic.Anthropic(**kwargs)


def _generate_score(client, model: str, max_tokens: int,
                    index: int, total: int) -> str:
    user_prompt = USER_PROMPT_TEMPLATE.format(index=index, total=total)

    message = client.messages.create(
        model      = model,
        max_tokens = max_tokens,
        system     = SYSTEM_PROMPT,
        messages   = [{"role": "user", "content": user_prompt}],
    )

    content = message.content
    if isinstance(content, list):
        text_blocks = [b.text for b in content if hasattr(b, "text")]
        return "\n".join(text_blocks).strip()
    return str(content).strip()


def _fix_brackets(text: str) -> str:
    """Fix common LLM bracket mistakes: } where ] is needed, missing }, etc."""
    import re
    # GLM-5-Turbo sometimes closes "structure": [...] with } instead of ]
    # Pattern: },\n  "tracks" -> ],\n  "tracks"
    text = re.sub(r'\},(\s*\n\s*)"tracks"', r'],\1"tracks"', text, count=1)
    # Also: }\n] (array of objects closed with } then ]) -> }\n]
    text = re.sub(r'\}\s*\]', '}\n]', text)
    # Trailing comma before closing bracket
    text = re.sub(r',(\s*\})', r'\1', text)
    text = re.sub(r',(\s*\])', r'\1', text)
    return text


def _parse_score(raw: str, index: int) -> SongScore:
    import re

    raw = raw.strip()
    if raw.startswith("```"):
        lines = raw.split("\n")
        raw = "\n".join(lines[1:-1] if lines[-1].strip() == "```" else lines[1:])

    start = raw.find("{")
    end   = raw.rfind("}") + 1
    if start < 0 or end <= start:
        raise ValueError("No JSON object found in AI response")
    raw = raw[start:end]

    # Attempt 1: direct parse
    try:
        data = json.loads(raw)
        return SongScore.from_dict(data)
    except json.JSONDecodeError:
        pass

    # Attempt 2: fix common bracket errors from LLMs
    fixed = _fix_brackets(raw)
    try:
        data = json.loads(fixed)
        log.info("JSON bracket fix succeeded on attempt 2")
        return SongScore.from_dict(data)
    except json.JSONDecodeError:
        pass

    # Attempt 3: remove // comments + trailing commas + bracket fix
    cleaned = re.sub(r'//.*$', '', fixed, flags=re.MULTILINE)
    cleaned = re.sub(r',(\s*\})', r'\1', cleaned)
    cleaned = re.sub(r',(\s*\])', r'\1', cleaned)

    try:
        data = json.loads(cleaned)
        log.info("JSON cleaned successfully on attempt 3")
        return SongScore.from_dict(data)
    except json.JSONDecodeError as exc:
        # Attempt 4: brute-force close unclosed brackets
        open_b = cleaned.count('{') - cleaned.count('}')
        open_br = cleaned.count('[') - cleaned.count(']')
        if open_b > 0 or open_br > 0:
            repaired = cleaned.rstrip().rstrip(',')
            repaired += ']' * max(0, open_br)
            repaired += '}' * max(0, open_b)
            try:
                data = json.loads(repaired)
                log.info("JSON repaired (bracket closure) on attempt 4")
                return SongScore.from_dict(data)
            except json.JSONDecodeError as exc4:
                pass

        raise ValueError(
            "JSON parse failed after all attempts: %s\nLast output:\n%s"
            % (exc, cleaned[:800])
        )


def run_loop(
    base_url:       str,
    api_key:        str,
    model:          str,
    max_tokens:     int,
    count:          int,
    delay:          float,
    render:         bool,
    lib_root:       str,
    output_prefix:  str = "ai_song",
    dry_run:        bool = False,
):
    client   = _build_client(base_url, api_key)
    renderer = ScoreRenderer(lib_root) if (render and not dry_run) else None
    total    = count if count > 0 else "inf"

    log.info("Starting AI production loop — model=%s  count=%s  render=%s",
             model, total, render)
    log.info("Scores will be saved to: %s", SCORES_DIR)
    if render:
        log.info("Library root: %s", lib_root)
    if dry_run:
        log.info("DRY RUN — Ableton will NOT be touched")

    produced = 0
    iteration = 0

    while True:
        iteration += 1
        if count > 0 and produced >= count:
            break

        log.info("Generating song %d / %s", iteration, total)

        try:
            raw_json = _generate_score(client, model, max_tokens, iteration, count or 999)
            log.debug("Raw AI output:\n%s", raw_json[:500])

            score = _parse_score(raw_json, iteration)

            warnings = score.validate()
            if warnings:
                log.warning("Validation warnings: %s", warnings)

            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            filename  = "%s_%03d_%s.json" % (output_prefix, iteration, timestamp)
            saved_path = SCORES_DIR / filename
            score.save(saved_path)
            log.info("Saved: %s  (%d tracks, %.0f bars)",
                     filename, len(score.tracks), score.total_bars())

            if renderer:
                log.info("Rendering into Ableton...")
                result = renderer.render(score, clear_first=True)
                if result.get("success"):
                    log.info("Rendered OK  tracks=%d  clips=%d  bars=%.0f",
                             len(result["tracks_created"]),
                             result["clips_created"],
                             score.total_bars())
                else:
                    log.warning("Render completed with errors:")
                    for err in result.get("errors", []):
                        log.warning("  - %s", err)

            produced += 1

        except KeyboardInterrupt:
            log.info("Loop interrupted by user. %d songs produced.", produced)
            break
        except json.JSONDecodeError as exc:
            log.error("JSON parse error on iteration %d: %s", iteration, exc)
        except Exception as exc:
            log.exception("Unexpected error on iteration %d: %s", iteration, exc)

        if count == 0 or produced < count:
            if delay > 0:
                log.info("Waiting %.0fs before next generation...", delay)
                time.sleep(delay)

    log.info("Loop complete. %d songs produced and saved to %s", produced, SCORES_DIR)


def main():
    parser = argparse.ArgumentParser(
        description="Autonomous AI music production loop (Anthropic-compatible)",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=__doc__,
    )
    parser.add_argument("--base-url",  default=os.environ.get("AI_BASE_URL",   "https://api.anthropic.com"))
    parser.add_argument("--api-key",   default=os.environ.get("AI_API_KEY",    ""))
    parser.add_argument("--model",    default=os.environ.get("AI_MODEL",     "GLM-5-Turbo"))
    parser.add_argument("--max-tokens",default=int(os.environ.get("AI_MAX_TOKENS", "4096")), type=int)
    parser.add_argument("--count",    default=int(os.environ.get("LOOP_COUNT",    "10")),   type=int,
                        help="Songs to produce (0 = infinite)")
    parser.add_argument("--delay",    default=float(os.environ.get("LOOP_DELAY",  "5")),   type=float,
                        help="Seconds between generations")
    parser.add_argument("--render",    action="store_true",
                        default=os.environ.get("RENDER_AFTER", "0") == "1",
                        help="Render each score into Ableton immediately")
    parser.add_argument("--lib-root", default=os.environ.get("LIB_ROOT", _DEFAULT_LIB_ROOT))
    parser.add_argument("--prefix",   default="ai_song",
                        help="Filename prefix for saved scores")
    parser.add_argument("--dry-run",  action="store_true",
                        help="Generate + validate + save but do NOT call Ableton")
    parser.add_argument("--list",     action="store_true",
                        help="List saved scores and exit")

    args = parser.parse_args()

    if args.list:
        scores = sorted(SCORES_DIR.glob("*.json"))
        if not scores:
            print("No scores saved yet.")
        else:
            for f in scores:
                size = f.stat().st_size
                print("  %s  (%d bytes)" % (f.name, size))
        return

    if not args.api_key:
        parser.error("API key required. Set --api-key or AI_API_KEY env variable.")

    run_loop(
        base_url      = args.base_url,
        api_key       = args.api_key,
        model         = args.model,
        max_tokens    = args.max_tokens,
        count         = args.count,
        delay         = args.delay,
        render        = args.render,
        lib_root      = args.lib_root,
        output_prefix = args.prefix,
        dry_run       = args.dry_run,
    )


if __name__ == "__main__":
    main()
